more fixes for adaptive/learner/balancing_learner.py

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -13,9 +13,7 @@
 from adaptive.utils import cache_latest, named_product, restore
 
 
-def dispatch(
-    child_functions: List[Callable], arg: Any,
-) -> Union[int, np.float64, float]:
+def dispatch(child_functions: List[Callable], arg: Any,) -> Union[Any]:
     index, x = arg
     return child_functions[index](x)
 
@@ -94,14 +92,14 @@ def __init__(
         self.strategy = strategy
 
     @property
-    def data(self) -> Dict[Tuple[int, int], int]:
+    def data(self) -> Dict[Tuple[int, Any], Any]:
         data = {}
         for i, l in enumerate(self.learners):
             data.update({(i, p): v for p, v in l.data.items()})
         return data
 
     @property
-    def pending_points(self) -> Set[Tuple[int, int]]:
+    def pending_points(self) -> Set[Tuple[int, Any]]:
         pending_points = set()
         for i, l in enumerate(self.learners):
             pending_points.update({(i, p) for p in l.pending_points})
@@ -140,7 +138,9 @@ def strategy(self, strategy):
                 ' strategy="npoints", or strategy="cycle" is implemented.'
             )
 
-    def _ask_and_tell_based_on_loss_improvements(self, n: int) -> Any:
+    def _ask_and_tell_based_on_loss_improvements(
+        self, n: int
+    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -165,11 +165,7 @@ def _ask_and_tell_based_on_loss_improvements(self, n: int) -> Any:
 
     def _ask_and_tell_based_on_loss(
         self, n: int
-    ) -> Union[
-        Tuple[List[Tuple[int, float]], List[np.float64]],
-        Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
-        Tuple[List[Tuple[int, int]], List[float]],
-    ]:
+    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -192,11 +188,7 @@ def _ask_and_tell_based_on_loss(
 
     def _ask_and_tell_based_on_npoints(
         self, n: int
-    ) -> Union[
-        Tuple[List[Union[Tuple[np.int64, int], Tuple[np.int64, float]]], List[float]],
-        Tuple[List[Tuple[np.int64, float]], List[np.float64]],
-        Tuple[List[Tuple[np.int64, int]], List[float]],
-    ]:
+    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -214,11 +206,7 @@ def _ask_and_tell_based_on_npoints(
 
     def _ask_and_tell_based_on_cycle(
         self, n: int
-    ) -> Union[
-        Tuple[List[Tuple[int, float]], List[np.float64]],
-        Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
-        Tuple[List[Tuple[int, int]], List[float]],
-    ]:
+    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
         points, loss_improvements = [], []
         for _ in range(n):
             index = next(self._cycle)
@@ -229,7 +217,9 @@ def _ask_and_tell_based_on_cycle(
 
         return points, loss_improvements
 
-    def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[Any], List[float]]:
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
         """Chose points for learners."""
         if n == 0:
             return [], []
@@ -240,26 +230,20 @@ def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[Any], List[float]
         else:
             return self._ask_and_tell(n)
 
-    def tell(
-        self,
-        x: Any,
-        y: Union[int, np.float64, float, Tuple[int, int], Tuple[np.int64, int]],
-    ) -> None:
+    def tell(self, x: Tuple[int, Any], y: Any,) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self._pending_loss.pop(index, None)
         self.learners[index].tell(x, y)
 
-    def tell_pending(self, x: Any) -> None:
+    def tell_pending(self, x: Tuple[int, Any]) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self.learners[index].tell_pending(x)
 
-    def _losses(
-        self, real: bool = True
-    ) -> Union[List[float], List[np.float64], List[Union[float, np.float64]]]:
+    def _losses(self, real: bool = True) -> List[float]:
         losses = []
         loss_dict = self._loss if real else self._pending_loss
 
@@ -271,7 +255,7 @@ def _losses(
         return losses
 
     @cache_latest
-    def loss(self, real: bool = True) -> Union[np.float64, float]:
+    def loss(self, real: bool = True) -> Union[float]:
         losses = self._losses(real)
         return max(losses)