feat(ai): add get_wells

src/ai/fitness/fitness.py

@@ -13,10 +13,6 @@ def calculate_fitness(game: Game) -> float:
     return fitness
 
 
-def get_wells(field: np.ndarray) -> int:
-    pass
-
-
 def _calc_holes(field: np.ndarray) -> float:
     height, width = field.shape
     penalty = 0

src/ai/fitness/wells.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import numpy as np
+
+from ai.log import log
+
+from .peaks import get_peaks
+
+
+def get_wells(
+    *, peaks: Optional[np.ndarray] = None, field: Optional[np.ndarray] = None
+) -> np.ndarray:
+    if peaks is None and field is None:
+        raise ValueError("peaks and field cannot both be None")
+    elif peaks is None:
+        peaks = get_peaks(field)
+
+    wells = np.zeros_like(peaks)
+    log.debug(f"{peaks=}")
+
+    first_well = peaks[1] - peaks[0]
+    wells[0] = first_well if first_well > 0 else 0
+
+    last_well = peaks[-2] - peaks[-1]
+    wells[-1] = last_well if last_well > 0 else 0
+
+    for idx in range(1, len(peaks) - 1):
+        well_l = peaks[idx - 1] - peaks[idx]
+        well_l = well_l if well_l > 0 else 0
+
+        well_r = peaks[idx + 1] - peaks[idx]
+        well_r = well_r if well_r > 0 else 0
+
+        wells[idx] = well_l if well_l >= well_r else well_r
+
+    log.debug(f"{wells=}")
+    return wells

tests/ai/test_fitness.py

@@ -8,6 +8,7 @@ from ai.fitness.transitions import (
     get_col_transition,
     get_row_transition,
 )
+from ai.fitness.wells import get_wells
 
 
 class TestFitness(unittest.TestCase):
@@ -58,3 +59,12 @@ class TestFitness(unittest.TestCase):
         )
         for field, answer in zip(self.fields, answers):
             self.assertTrue(np.array_equal(get_holes(field), answer))
+
+    def test_get_wells(self):
+        answers = (
+            np.array([1, 0, 2, 1, 0]),
+            np.array([0, 0, 0, 0, 0]),
+            np.array([2, 0, 2, 0, 0]),
+        )
+        for field, answer in zip(self.fields, answers):
+            self.assertTrue(np.array_equal(get_wells(field=field), answer))