clean up code

clocktower/__init__.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 # import numpy.typing as npt
@@ -9,15 +9,17 @@ from scipy.stats.stats import F_onewayResult
 
 class ClocktowerManager:
     __slots__ = ['bounds', 'data', 'rng', 'significance', 'best_guess', 'min_correct_tries',
-                 'anova', 'use_smart_strategy']
+                 'anova', 'use_adaptive', 'min_required_count', 'favor_candidate_chance']
     bounds: Tuple[int, int]
     data: Dict[int, np.ndarray]
     rng: npr.Generator
     significance: float
-    min_correct_tries: int
     best_guess: Optional[int]
     anova: Optional[F_onewayResult]
-    use_smart_strategy: bool
+    use_adaptive: bool
+    min_correct_tries: int
+    min_required_count: int
+    favor_candidate_chance: float
 
     def __init__(self, bounds: Tuple[int, int] = (0, 256),
                  significance: float = 0.01,
@@ -32,29 +34,45 @@ class ClocktowerManager:
         self.min_correct_tries = min_correct_tries
         self.best_guess = None
         self.anova = None
-        self.use_smart_strategy = True
+        self.use_adaptive = True
+        self.min_required_count = 3
+        self.favor_candidate_chance = 0.5
 
     def next_guess(self) -> Optional[int]:
+        if self.use_adaptive:
+            return self._next_guess_adaptive()
+        else:
+            return self._next_guess_naive()
+
+    def _next_guess_adaptive(self) -> Optional[int]:
+        counts = [len(x) for x in self.data.values()]
+        min_count = min(counts)
+        low_count = np.quantile(counts, 0.5)
+        low_keys = [k for k in self.data.keys() if len(self.data[k]) <= low_count]
+        min_keys = [k for k in self.data.keys() if len(self.data[k]) == min_count]
+
+        if (self.best_guess is not None and
+                len(self.data[self.best_guess]) >= self.min_correct_tries):
+            return None
+
+        if min_count < self.min_required_count:
+            return self.rng.choice(min_keys)
+        elif self.rng.uniform(0, 1) > self.favor_candidate_chance:
+            return self.rng.choice(low_keys)
+        elif self.best_guess is not None:
+            return self.best_guess
+        else:
+            means = {k: v.mean() for k, v in self.data.items()}
+            return max(means.items(), key=lambda x: x[1])[0]
+
+    def _next_guess_naive(self) -> Optional[int]:
         min_count = min([len(x) for x in self.data.values()])
         min_keys = [k for k in self.data.keys() if len(self.data[k]) == min_count]
 
-        if self.use_smart_strategy:
-            if (self.best_guess is not None and
-                    len(self.data[self.best_guess]) >= self.min_correct_tries):
-                return None
-
-            if min_count < 3 or self.rng.uniform(0, 1) > 0.8:
-                return self.rng.choice(min_keys)
-            elif self.best_guess is not None:
-                return self.best_guess
-            else:
-                means = {k: v.mean() for k, v in self.data.items()}
-                return max(means.items(), key=lambda x: x[1])[0]
+        if self.best_guess is not None:
+            return None
         else:
-            if self.best_guess is not None:
-                return None
-            else:
-                return self.rng.choice(min_keys)
+            return self.rng.choice(min_keys)
 
     def update(self, guess: int, value: float) -> None:
         if guess not in self.data.keys():
@@ -68,7 +86,8 @@ class ClocktowerManager:
             return
         self.anova = f_oneway(*inputs)
         if self.anova.pvalue <= self.significance:
-            self.best_guess = max(self.data.items(), key=lambda v: v[1].mean())[0]
+            self.best_guess = max(self.data.items(),
+                                  key=lambda v: v[1].mean() if len(v[1]) > 0 else 0)[0]
         else:
             self.best_guess = None
 
@@ -83,12 +102,14 @@ def main() -> None:
     # stdev = 7059
     # u0 = 500000
     # u1 = 506046
-    # stdev = 1000
-    # u0 = 500000
-    # u1 = 506046
-    stdev = 7000 * 2
+
+    stdev = 1000
     u0 = 500000
     u1 = 506046
+
+    # stdev = 7000 * 2
+    # u0 = 500000
+    # u1 = 506046
     correct_guess = 0x42
 
     def sample(guess):
@@ -111,3 +132,4 @@ def main() -> None:
     print("state", mgr.anova)
     print(f"answer 0x{mgr.get_best_guess():02x}")
     print("took", num_guesses, "guesses")
+    print("stats:", ", ".join([f"{i:02x} {len(mgr.data[i]):04d}" for i in range(256)]))