add type annotations for adaptive/learner/learner1D.py

Author: basnijholt

adaptive/learner/learner1D.py

@@ -2,10 +2,15 @@
 import math
 from collections.abc import Iterable
 from copy import deepcopy
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import sortedcollections
 import sortedcontainers
+from numpy import float64, ndarray
+from sortedcollections.recipes import ItemSortedDict
+from sortedcontainers.sorteddict import SortedDict
 
 from adaptive.learner.base_learner import BaseLearner, uses_nth_neighbors
 from adaptive.learner.learnerND import volume
@@ -15,7 +20,10 @@
 
 
 @uses_nth_neighbors(0)
-def uniform_loss(xs, ys):
+def uniform_loss(
+    xs: Union[Tuple[float, float], Tuple[float64, float64]],
+    ys: Union[Tuple[float, float], Tuple[float64, float64]],
+) -> Union[float64, float]:
     """Loss function that samples the domain uniformly.
 
     Works with `~adaptive.Learner1D` only.
@@ -35,7 +43,15 @@ def uniform_loss(xs, ys):
 
 
 @uses_nth_neighbors(0)
-def default_loss(xs, ys):
+def default_loss(
+    xs: Union[
+        Tuple[float, float],
+        Tuple[float64, float],
+        Tuple[float64, float64],
+        Tuple[float, float64],
+    ],
+    ys: Union[Tuple[float, float], Tuple[ndarray, ndarray], Tuple[float64, float64]],
+) -> float64:
     """Calculate loss on a single interval.
 
     Currently returns the rescaled length of the interval. If one of the
@@ -52,7 +68,7 @@ def default_loss(xs, ys):
 
 
 @uses_nth_neighbors(1)
-def triangle_loss(xs, ys):
+def triangle_loss(xs: Any, ys: Any) -> Union[float64, float]:
     xs = [x for x in xs if x is not None]
     ys = [y for y in ys if y is not None]
 
@@ -69,7 +85,9 @@ def triangle_loss(xs, ys):
     return sum(vol(pts[i : i + 3]) for i in range(N)) / N
 
 
-def curvature_loss_function(area_factor=1, euclid_factor=0.02, horizontal_factor=0.02):
+def curvature_loss_function(
+    area_factor: int = 1, euclid_factor: float = 0.02, horizontal_factor: float = 0.02
+) -> Callable:
     # XXX: add a doc-string
     @uses_nth_neighbors(1)
     def curvature_loss(xs, ys):
@@ -88,7 +106,9 @@ def curvature_loss(xs, ys):
     return curvature_loss
 
 
-def linspace(x_left, x_right, n):
+def linspace(
+    x_left: Union[int, float64, float], x_right: Union[int, float64, float], n: int
+) -> Union[List[float], List[float64]]:
     """This is equivalent to
     'np.linspace(x_left, x_right, n, endpoint=False)[1:]',
     but it is 15-30 times faster for small 'n'."""
@@ -100,7 +120,7 @@ def linspace(x_left, x_right, n):
         return [x_left + step * i for i in range(1, n)]
 
 
-def _get_neighbors_from_list(xs):
+def _get_neighbors_from_list(xs: ndarray) -> SortedDict:
     xs = np.sort(xs)
     xs_left = np.roll(xs, 1).tolist()
     xs_right = np.roll(xs, -1).tolist()
@@ -110,7 +130,9 @@ def _get_neighbors_from_list(xs):
     return sortedcontainers.SortedDict(neighbors)
 
 
-def _get_intervals(x, neighbors, nth_neighbors):
+def _get_intervals(
+    x: Union[int, float64, float], neighbors: SortedDict, nth_neighbors: int
+) -> Any:
     nn = nth_neighbors
     i = neighbors.index(x)
     start = max(0, i - nn - 1)
@@ -163,7 +185,12 @@ class Learner1D(BaseLearner):
     decorator for more information.
     """
 
-    def __init__(self, function, bounds, loss_per_interval=None):
+    def __init__(
+        self,
+        function: Union[Callable, partial],
+        bounds: Union[Tuple[int, int], Tuple[float, float], ndarray],
+        loss_per_interval: Optional[Callable] = None,
+    ) -> None:
         self.function = function
 
         if hasattr(loss_per_interval, "nth_neighbors"):
@@ -205,7 +232,7 @@ def __init__(self, function, bounds, loss_per_interval=None):
         self._vdim = None
 
     @property
-    def vdim(self):
+    def vdim(self) -> int:
         """Length of the output of ``learner.function``.
         If the output is unsized (when it's a scalar)
         then `vdim = 1`.
@@ -225,35 +252,41 @@ def vdim(self):
         return self._vdim
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> Union[int, float64, float]:
         losses = self.losses if real else self.losses_combined
         if not losses:
             return np.inf
         max_interval, max_loss = losses.peekitem(0)
         return max_loss
 
-    def _scale_x(self, x):
+    def _scale_x(
+        self, x: Optional[Union[float, int, float64]]
+    ) -> Optional[Union[float, float64]]:
         if x is None:
             return None
         return x / self._scale[0]
 
-    def _scale_y(self, y):
+    def _scale_y(
+        self, y: Optional[Union[int, ndarray, float64, float]]
+    ) -> Optional[Union[float, float64, ndarray]]:
         if y is None:
             return None
         y_scale = self._scale[1] or 1
         return y / y_scale
 
-    def _get_point_by_index(self, ind):
+    def _get_point_by_index(self, ind: int) -> Optional[Union[int, float64, float]]:
         if ind < 0 or ind >= len(self.neighbors):
             return None
         return self.neighbors.keys()[ind]
 
-    def _get_loss_in_interval(self, x_left, x_right):
+    def _get_loss_in_interval(
+        self, x_left: Union[int, float64, float], x_right: Union[int, float64, float]
+    ) -> Union[int, float64, float]:
         assert x_left is not None and x_right is not None
 
         if x_right - x_left < self._dx_eps:
@@ -273,7 +306,9 @@ def _get_loss_in_interval(self, x_left, x_right):
         # we need to compute the loss for this interval
         return self.loss_per_interval(xs_scaled, ys_scaled)
 
-    def _update_interpolated_loss_in_interval(self, x_left, x_right):
+    def _update_interpolated_loss_in_interval(
+        self, x_left: Union[int, float64, float], x_right: Union[int, float64, float]
+    ) -> None:
         if x_left is None or x_right is None:
             return
 
@@ -289,7 +324,7 @@ def _update_interpolated_loss_in_interval(self, x_left, x_right):
             self.losses_combined[a, b] = (b - a) * loss / dx
             a = b
 
-    def _update_losses(self, x, real=True):
+    def _update_losses(self, x: Union[int, float64, float], real: bool = True) -> None:
         """Update all losses that depend on x"""
         # When we add a new point x, we should update the losses
         # (x_left, x_right) are the "real" neighbors of 'x'.
@@ -332,7 +367,7 @@ def _update_losses(self, x, real=True):
             self.losses_combined[x, b] = float("inf")
 
     @staticmethod
-    def _find_neighbors(x, neighbors):
+    def _find_neighbors(x: Union[int, float64, float], neighbors: SortedDict) -> Any:
         if x in neighbors:
             return neighbors[x]
         pos = neighbors.bisect_left(x)
@@ -341,14 +376,18 @@ def _find_neighbors(x, neighbors):
         x_right = keys[pos] if pos != len(neighbors) else None
         return x_left, x_right
 
-    def _update_neighbors(self, x, neighbors):
+    def _update_neighbors(
+        self, x: Union[int, float64, float], neighbors: SortedDict
+    ) -> None:
         if x not in neighbors:  # The point is new
             x_left, x_right = self._find_neighbors(x, neighbors)
             neighbors[x] = [x_left, x_right]
             neighbors.get(x_left, [None, None])[1] = x
             neighbors.get(x_right, [None, None])[0] = x
 
-    def _update_scale(self, x, y):
+    def _update_scale(
+        self, x: Union[int, float64, float], y: Union[float, int, float64, ndarray]
+    ) -> None:
         """Update the scale with which the x and y-values are scaled.
 
         For a learner where the function returns a single scalar the scale
@@ -375,7 +414,7 @@ def _update_scale(self, x, y):
                 self._bbox[1][1] = max(self._bbox[1][1], y)
                 self._scale[1] = self._bbox[1][1] - self._bbox[1][0]
 
-    def tell(self, x, y):
+    def tell(self, x: Union[int, float64, float], y: Any) -> None:
         if x in self.data:
             # The point is already evaluated before
             return
@@ -410,15 +449,15 @@ def tell(self, x, y):
 
             self._oldscale = deepcopy(self._scale)
 
-    def tell_pending(self, x):
+    def tell_pending(self, x: Union[int, float64, float]) -> None:
         if x in self.data:
             # The point is already evaluated before
             return
         self.pending_points.add(x)
         self._update_neighbors(x, self.neighbors_combined)
         self._update_losses(x, real=False)
 
-    def tell_many(self, xs, ys, *, force=False):
+    def tell_many(self, xs: Any, ys: Any, *, force=False) -> None:
         if not force and not (len(xs) > 0.5 * len(self.data) and len(xs) > 2):
             # Only run this more efficient method if there are
             # at least 2 points and the amount of points added are
@@ -486,7 +525,7 @@ def tell_many(self, xs, ys, *, force=False):
                 # have an inf loss.
                 self._update_interpolated_loss_in_interval(*ival)
 
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True) -> Any:
         """Return 'n' points that are expected to maximally reduce the loss."""
         points, loss_improvements = self._ask_points_without_adding(n)
 
@@ -496,7 +535,7 @@ def ask(self, n, tell_pending=True):
 
         return points, loss_improvements
 
-    def _ask_points_without_adding(self, n):
+    def _ask_points_without_adding(self, n: int) -> Any:
         """Return 'n' points that are expected to maximally reduce the loss.
         Without altering the state of the learner"""
         # Find out how to divide the n points over the intervals
@@ -574,7 +613,7 @@ def _ask_points_without_adding(self, n):
 
         return points, loss_improvements
 
-    def _loss(self, mapping, ival):
+    def _loss(self, mapping: ItemSortedDict, ival: Any) -> Any:
         loss = mapping[ival]
         return finite_loss(ival, loss, self._scale[0])
 
@@ -613,20 +652,20 @@ def plot(self, *, scatter_or_line="scatter"):
 
         return p.redim(x=dict(range=plot_bounds))
 
-    def remove_unfinished(self):
+    def remove_unfinished(self) -> None:
         self.pending_points = set()
         self.losses_combined = deepcopy(self.losses)
         self.neighbors_combined = deepcopy(self.neighbors)
 
-    def _get_data(self):
+    def _get_data(self) -> Dict[Union[int, float], float]:
         return self.data
 
-    def _set_data(self, data):
+    def _set_data(self, data: Dict[Union[int, float], float]) -> None:
         if data:
             self.tell_many(*zip(*data.items()))
 
 
-def loss_manager(x_scale):
+def loss_manager(x_scale: Union[int, float64, float]) -> ItemSortedDict:
     def sort_key(ival, loss):
         loss, ival = finite_loss(ival, loss, x_scale)
         return -loss, ival
@@ -635,7 +674,9 @@ def sort_key(ival, loss):
     return sorted_dict
 
 
-def finite_loss(ival, loss, x_scale):
+def finite_loss(
+    ival: Any, loss: Union[int, float64, float], x_scale: Union[int, float64, float]
+) -> Any:
     """Get the socalled finite_loss of an interval in order to be able to
     sort intervals that have infinite loss."""
     # If the loss is infinite we return the