add type annotations for adaptive/learner/balancing_learner.py

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -4,15 +4,25 @@
 from contextlib import suppress
 from functools import partial
 from operator import itemgetter
+from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import numpy as np
+from numpy import float64, int64
 
+from adaptive.learner.average_learner import AverageLearner
 from adaptive.learner.base_learner import BaseLearner
+from adaptive.learner.learner1D import Learner1D
+from adaptive.learner.learner2D import Learner2D
+from adaptive.learner.learnerND import LearnerND
+from adaptive.learner.sequence_learner import SequenceLearner, _IgnoreFirstArgument
 from adaptive.notebook_integration import ensure_holoviews
 from adaptive.utils import cache_latest, named_product, restore
 
 
-def dispatch(child_functions, arg):
+def dispatch(
+    child_functions: Union[List[Callable], List[partial], List[_IgnoreFirstArgument]],
+    arg: Any,
+) -> Union[int, float64, float]:
     index, x = arg
     return child_functions[index](x)
 
@@ -68,7 +78,19 @@ class BalancingLearner(BaseLearner):
     behave in an undefined way. Change the `strategy` in that case.
     """
 
-    def __init__(self, learners, *, cdims=None, strategy="loss_improvements"):
+    def __init__(
+        self,
+        learners: Union[
+            List[SequenceLearner],
+            List[AverageLearner],
+            List[Learner2D],
+            List[Learner1D],
+            List[LearnerND],
+        ],
+        *,
+        cdims=None,
+        strategy="loss_improvements"
+    ) -> None:
         self.learners = learners
 
         # Naively we would make 'function' a method, but this causes problems
@@ -89,21 +111,21 @@ def __init__(self, learners, *, cdims=None, strategy="loss_improvements"):
         self.strategy = strategy
 
     @property
-    def data(self):
+    def data(self) -> Dict[Tuple[int, int], int]:
         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):
+    def pending_points(self) -> Set[Tuple[int, int]]:
         pending_points = set()
         for i, l in enumerate(self.learners):
             pending_points.update({(i, p) for p in l.pending_points})
         return pending_points
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         return sum(l.npoints for l in self.learners)
 
     @property
@@ -135,7 +157,7 @@ def strategy(self, strategy):
                 ' strategy="npoints", or strategy="cycle" is implemented.'
             )
 
-    def _ask_and_tell_based_on_loss_improvements(self, n):
+    def _ask_and_tell_based_on_loss_improvements(self, n: int) -> Any:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
@@ -158,7 +180,13 @@ def _ask_and_tell_based_on_loss_improvements(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_loss(self, n):
+    def _ask_and_tell_based_on_loss(
+        self, n: int
+    ) -> Union[
+        Tuple[List[Tuple[int, float]], List[float64]],
+        Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
+        Tuple[List[Tuple[int, int]], 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):
@@ -179,7 +207,13 @@ def _ask_and_tell_based_on_loss(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_npoints(self, n):
+    def _ask_and_tell_based_on_npoints(
+        self, n: int
+    ) -> Union[
+        Tuple[List[Union[Tuple[int64, int], Tuple[int64, float]]], List[float]],
+        Tuple[List[Tuple[int64, float]], List[float64]],
+        Tuple[List[Tuple[int64, int]], 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):
@@ -195,7 +229,13 @@ def _ask_and_tell_based_on_npoints(self, n):
         points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
-    def _ask_and_tell_based_on_cycle(self, n):
+    def _ask_and_tell_based_on_cycle(
+        self, n: int
+    ) -> Union[
+        Tuple[List[Tuple[int, float]], List[float64]],
+        Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
+        Tuple[List[Tuple[int, int]], List[float]],
+    ]:
         points, loss_improvements = [], []
         for _ in range(n):
             index = next(self._cycle)
@@ -206,7 +246,7 @@ def _ask_and_tell_based_on_cycle(self, n):
 
         return points, loss_improvements
 
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True) -> Any:
         """Chose points for learners."""
         if n == 0:
             return [], []
@@ -217,20 +257,24 @@ def ask(self, n, tell_pending=True):
         else:
             return self._ask_and_tell(n)
 
-    def tell(self, x, y):
+    def tell(
+        self, x: Any, y: Union[int, float64, float, Tuple[int, int], Tuple[int64, int]]
+    ) -> 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):
+    def tell_pending(self, x: 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=True):
+    def _losses(
+        self, real: bool = True
+    ) -> Union[List[float], List[float64], List[Union[float, float64]]]:
         losses = []
         loss_dict = self._loss if real else self._pending_loss
 
@@ -242,7 +286,7 @@ def _losses(self, real=True):
         return losses
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> Union[float64, float]:
         losses = self._losses(real)
         return max(losses)
 
@@ -372,7 +416,7 @@ def from_product(cls, f, learner_type, learner_kwargs, combos):
             learners.append(learner)
         return cls(learners, cdims=arguments)
 
-    def save(self, fname, compress=True):
+    def save(self, fname: Callable, compress: bool = True) -> None:
         """Save the data of the child learners into pickle files
         in a directory.
 
@@ -410,7 +454,7 @@ def save(self, fname, compress=True):
             for l in self.learners:
                 l.save(fname(l), compress=compress)
 
-    def load(self, fname, compress=True):
+    def load(self, fname: Callable, compress: bool = True) -> None:
         """Load the data of the child learners from pickle files
         in a directory.