add type annotations for adaptive/learner/integrator_learner.py

Author: basnijholt

adaptive/learner/integrator_learner.py

@@ -2,10 +2,13 @@
 
 import sys
 from collections import defaultdict
+from functools import partial
 from math import sqrt
 from operator import attrgetter
+from typing import Any, Callable, List, Optional, Set, Tuple, Union
 
 import numpy as np
+from numpy import bool_, float64, ndarray, ufunc
 from scipy.linalg import norm
 from sortedcontainers import SortedSet
 
@@ -30,7 +33,7 @@
 )
 
 
-def _downdate(c, nans, depth):
+def _downdate(c: ndarray, nans: List[int], depth: int) -> ndarray:
     # This is algorithm 5 from the thesis of Pedro Gonnet.
     b = b_def[depth].copy()
     m = ns[depth] - 1
@@ -48,7 +51,7 @@ def _downdate(c, nans, depth):
     return c
 
 
-def _zero_nans(fx):
+def _zero_nans(fx: ndarray) -> List[int]:
     """Caution: this function modifies fx."""
     nans = []
     for i in range(len(fx)):
@@ -58,7 +61,7 @@ def _zero_nans(fx):
     return nans
 
 
-def _calc_coeffs(fx, depth):
+def _calc_coeffs(fx: ndarray, depth: int) -> ndarray:
     """Caution: this function modifies fx."""
     nans = _zero_nans(fx)
     c_new = V_inv[depth] @ fx
@@ -138,7 +141,9 @@ class _Interval:
         "removed",
     ]
 
-    def __init__(self, a, b, depth, rdepth):
+    def __init__(
+        self, a: Union[int, float64], b: Union[int, float64], depth: int, rdepth: int
+    ) -> None:
         self.children = []
         self.data = {}
         self.a = a
@@ -150,15 +155,15 @@ def __init__(self, a, b, depth, rdepth):
         self.removed = False
 
     @classmethod
-    def make_first(cls, a, b, depth=2):
+    def make_first(cls, a: int, b: int, depth: int = 2) -> "_Interval":
         ival = _Interval(a, b, depth, rdepth=1)
         ival.ndiv = 0
         ival.parent = None
         ival.err = sys.float_info.max  # needed because inf/2 == inf
         return ival
 
     @property
-    def T(self):
+    def T(self) -> ndarray:
         """Get the correct shift matrix.
 
         Should only be called on children of a split interval.
@@ -169,24 +174,24 @@ def T(self):
         assert left != right
         return T_left if left else T_right
 
-    def refinement_complete(self, depth):
+    def refinement_complete(self, depth: int) -> bool:
         """The interval has all the y-values to calculate the intergral."""
         if len(self.data) < ns[depth]:
             return False
         return all(p in self.data for p in self.points(depth))
 
-    def points(self, depth=None):
+    def points(self, depth: Optional[int] = None) -> ndarray:
         if depth is None:
             depth = self.depth
         a = self.a
         b = self.b
         return (a + b) / 2 + (b - a) * xi[depth] / 2
 
-    def refine(self):
+    def refine(self) -> "_Interval":
         self.depth += 1
         return self
 
-    def split(self):
+    def split(self) -> List["_Interval"]:
         points = self.points()
         m = points[len(points) // 2]
         ivals = [
@@ -201,10 +206,10 @@ def split(self):
 
         return ivals
 
-    def calc_igral(self):
+    def calc_igral(self) -> None:
         self.igral = (self.b - self.a) * self.c[0] / sqrt(2)
 
-    def update_heuristic_err(self, value):
+    def update_heuristic_err(self, value: Union[float64, float]) -> None:
         """Sets the error of an interval using a heuristic (half the error of
         the parent) when the actual error cannot be calculated due to its
         parents not being finished yet. This error is propagated down to its
@@ -217,7 +222,7 @@ def update_heuristic_err(self, value):
                 continue
             child.update_heuristic_err(value / 2)
 
-    def calc_err(self, c_old):
+    def calc_err(self, c_old: ndarray) -> float:
         c_new = self.c
         c_diff = np.zeros(max(len(c_old), len(c_new)))
         c_diff[: len(c_old)] = c_old
@@ -229,7 +234,7 @@ def calc_err(self, c_old):
                 child.update_heuristic_err(self.err / 2)
         return c_diff
 
-    def calc_ndiv(self):
+    def calc_ndiv(self) -> None:
         div = self.parent.c00 and self.c00 / self.parent.c00 > 2
         self.ndiv += div
 
@@ -240,15 +245,17 @@ def calc_ndiv(self):
             for child in self.children:
                 child.update_ndiv_recursively()
 
-    def update_ndiv_recursively(self):
+    def update_ndiv_recursively(self) -> None:
         self.ndiv += 1
         if self.ndiv > ndiv_max and 2 * self.ndiv > self.rdepth:
             raise DivergentIntegralError
 
         for child in self.children:
             child.update_ndiv_recursively()
 
-    def complete_process(self, depth):
+    def complete_process(
+        self, depth: int
+    ) -> Union[Tuple[bool, bool], Tuple[bool, bool_]]:
         """Calculate the integral contribution and error from this interval,
         and update the done leaves of all ancestor intervals."""
         assert self.depth_complete is None or self.depth_complete == depth - 1
@@ -323,7 +330,7 @@ def complete_process(self, depth):
 
         return force_split, remove
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         lst = [
             f"(a, b)=({self.a:.5f}, {self.b:.5f})",
             f"depth={self.depth}",
@@ -335,7 +342,12 @@ def __repr__(self):
 
 
 class IntegratorLearner(BaseLearner):
-    def __init__(self, function, bounds, tol):
+    def __init__(
+        self,
+        function: Union[partial, ufunc, Callable],
+        bounds: Tuple[int, int],
+        tol: float,
+    ) -> None:
         """
         Parameters
         ----------
@@ -384,10 +396,10 @@ def __init__(self, function, bounds, tol):
         self.first_ival = ival
 
     @property
-    def approximating_intervals(self):
+    def approximating_intervals(self) -> Set["_Interval"]:
         return self.first_ival.done_leaves
 
-    def tell(self, point, value):
+    def tell(self, point: float64, value: float64) -> None:
         if point not in self.x_mapping:
             raise ValueError(f"Point {point} doesn't belong to any interval")
         self.data[point] = value
@@ -423,7 +435,7 @@ def tell(self, point, value):
     def tell_pending(self):
         pass
 
-    def propagate_removed(self, ival):
+    def propagate_removed(self, ival: "_Interval") -> None:
         def _propagate_removed_down(ival):
             ival.removed = True
             self.ivals.discard(ival)
@@ -433,7 +445,7 @@ def _propagate_removed_down(ival):
 
         _propagate_removed_down(ival)
 
-    def add_ival(self, ival):
+    def add_ival(self, ival: "_Interval") -> None:
         for x in ival.points():
             # Update the mappings
             self.x_mapping[x].add(ival)
@@ -444,15 +456,19 @@ def add_ival(self, ival):
                 self._stack.append(x)
         self.ivals.add(ival)
 
-    def ask(self, n, tell_pending=True):
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> Union[Tuple[List[float64], List[float64]], Tuple[List[float64], List[float]]]:
         """Choose points for learners."""
         if not tell_pending:
             with restore(self):
                 return self._ask_and_tell_pending(n)
         else:
             return self._ask_and_tell_pending(n)
 
-    def _ask_and_tell_pending(self, n):
+    def _ask_and_tell_pending(
+        self, n: int
+    ) -> Union[Tuple[List[float64], List[float64]], Tuple[List[float64], List[float]]]:
         points, loss_improvements = self.pop_from_stack(n)
         n_left = n - len(points)
         while n_left > 0:
@@ -468,7 +484,13 @@ def _ask_and_tell_pending(self, n):
 
         return points, loss_improvements
 
-    def pop_from_stack(self, n):
+    def pop_from_stack(
+        self, n: int
+    ) -> Union[
+        Tuple[List[float64], List[float64]],
+        Tuple[List[Any], List[Any]],
+        Tuple[List[float64], List[float]],
+    ]:
         points = self._stack[:n]
         self._stack = self._stack[n:]
         loss_improvements = [
@@ -479,7 +501,7 @@ def pop_from_stack(self, n):
     def remove_unfinished(self):
         pass
 
-    def _fill_stack(self):
+    def _fill_stack(self) -> List[float64]:
         # XXX: to-do if all the ivals have err=inf, take the interval
         # with the lowest rdepth and no children.
         force_split = bool(self.priority_split)
@@ -515,16 +537,16 @@ def _fill_stack(self):
         return self._stack
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @property
-    def igral(self):
+    def igral(self) -> float64:
         return sum(i.igral for i in self.approximating_intervals)
 
     @property
-    def err(self):
+    def err(self) -> float64:
         if self.approximating_intervals:
             err = sum(i.err for i in self.approximating_intervals)
             if err > sys.float_info.max: