add type annotations for adaptive/learner/learner2D.py

Author: basnijholt

adaptive/learner/learner2D.py

@@ -2,10 +2,14 @@
 import warnings
 from collections import OrderedDict
 from copy import copy
+from functools import partial
 from math import sqrt
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import numpy as np
+from numpy import bool_, float64, ndarray
 from scipy import interpolate
+from scipy.interpolate.interpnd import LinearNDInterpolator
 
 from adaptive.learner.base_learner import BaseLearner
 from adaptive.learner.triangulation import simplex_volume_in_embedding
@@ -15,7 +19,7 @@
 # Learner2D and helper functions.
 
 
-def deviations(ip):
+def deviations(ip: LinearNDInterpolator) -> List[ndarray]:
     """Returns the deviation of the linear estimate.
 
     Is useful when defining custom loss functions.
@@ -52,7 +56,7 @@ def deviation(p, v, g):
     return devs
 
 
-def areas(ip):
+def areas(ip: LinearNDInterpolator) -> ndarray:
     """Returns the area per triangle of the triangulation inside
     a `LinearNDInterpolator` instance.
 
@@ -73,7 +77,7 @@ def areas(ip):
     return areas
 
 
-def uniform_loss(ip):
+def uniform_loss(ip: LinearNDInterpolator) -> ndarray:
     """Loss function that samples the domain uniformly.
 
     Works with `~adaptive.Learner2D` only.
@@ -104,7 +108,7 @@ def uniform_loss(ip):
     return np.sqrt(areas(ip))
 
 
-def resolution_loss_function(min_distance=0, max_distance=1):
+def resolution_loss_function(min_distance: int = 0, max_distance: int = 1) -> Callable:
     """Loss function that is similar to the `default_loss` function, but you
     can set the maximimum and minimum size of a triangle.
 
@@ -143,7 +147,7 @@ def resolution_loss(ip):
     return resolution_loss
 
 
-def minimize_triangle_surface_loss(ip):
+def minimize_triangle_surface_loss(ip: LinearNDInterpolator) -> ndarray:
     """Loss function that is similar to the distance loss function in the
     `~adaptive.Learner1D`. The loss is the area spanned by the 3D
     vectors of the vertices.
@@ -189,7 +193,7 @@ def _get_vectors(points):
     return np.linalg.norm(np.cross(a, b) / 2, axis=1)
 
 
-def default_loss(ip):
+def default_loss(ip: LinearNDInterpolator) -> ndarray:
     """Loss function that combines `deviations` and `areas` of the triangles.
 
     Works with `~adaptive.Learner2D` only.
@@ -209,7 +213,7 @@ def default_loss(ip):
     return losses
 
 
-def choose_point_in_triangle(triangle, max_badness):
+def choose_point_in_triangle(triangle: ndarray, max_badness: int) -> ndarray:
     """Choose a new point in inside a triangle.
 
     If the ratio of the longest edge of the triangle squared
@@ -348,7 +352,14 @@ class Learner2D(BaseLearner):
     over each triangle.
     """
 
-    def __init__(self, function, bounds, loss_per_triangle=None):
+    def __init__(
+        self,
+        function: Union[partial, Callable],
+        bounds: Union[
+            List[Tuple[int, int]], Tuple[Tuple[int, int], Tuple[int, int]], ndarray
+        ],
+        loss_per_triangle: Optional[Callable] = None,
+    ) -> None:
         self.ndim = len(bounds)
         self._vdim = None
         self.loss_per_triangle = loss_per_triangle or default_loss
@@ -369,28 +380,28 @@ def __init__(self, function, bounds, loss_per_triangle=None):
         self.stack_size = 10
 
     @property
-    def xy_scale(self):
+    def xy_scale(self) -> ndarray:
         xy_scale = self._xy_scale
         if self.aspect_ratio == 1:
             return xy_scale
         else:
             return np.array([xy_scale[0], xy_scale[1] / self.aspect_ratio])
 
-    def _scale(self, points):
+    def _scale(self, points: Any) -> ndarray:
         points = np.asarray(points, dtype=float)
         return (points - self.xy_mean) / self.xy_scale
 
-    def _unscale(self, points):
+    def _unscale(self, points: ndarray) -> ndarray:
         points = np.asarray(points, dtype=float)
         return points * self.xy_scale + self.xy_mean
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         """Number of evaluated points."""
         return len(self.data)
 
     @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`.
@@ -406,7 +417,7 @@ def vdim(self):
         return self._vdim or 1
 
     @property
-    def bounds_are_done(self):
+    def bounds_are_done(self) -> bool:
         return not any(
             (p in self.pending_points or p in self._stack) for p in self._bounds_points
         )
@@ -443,7 +454,7 @@ def interpolated_on_grid(self, n=None):
         xs, ys = self._unscale(np.vstack([xs, ys]).T).T
         return xs, ys, zs
 
-    def _data_in_bounds(self):
+    def _data_in_bounds(self) -> Tuple[ndarray, ndarray]:
         if self.data:
             points = np.array(list(self.data.keys()))
             values = np.array(list(self.data.values()), dtype=float)
@@ -452,7 +463,7 @@ def _data_in_bounds(self):
             return points[inds], values[inds].reshape(-1, self.vdim)
         return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
 
-    def _data_interp(self):
+    def _data_interp(self) -> Any:
         if self.pending_points:
             points = list(self.pending_points)
             if self.bounds_are_done:
@@ -465,7 +476,7 @@ def _data_interp(self):
             return points, values
         return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
 
-    def _data_combined(self):
+    def _data_combined(self) -> Tuple[ndarray, ndarray]:
         points, values = self._data_in_bounds()
         if not self.pending_points:
             return points, values
@@ -483,7 +494,7 @@ def ip(self):
         )
         return self.interpolator(scaled=True)
 
-    def interpolator(self, *, scaled=False):
+    def interpolator(self, *, scaled=False) -> LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data.
 
@@ -514,7 +525,7 @@ def interpolator(self, *, scaled=False):
             points, values = self._data_in_bounds()
             return interpolate.LinearNDInterpolator(points, values)
 
-    def _interpolator_combined(self):
+    def _interpolator_combined(self) -> LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data *and* interpolated data of
         the `pending_points`."""
@@ -524,12 +535,29 @@ def _interpolator_combined(self):
             self._ip_combined = interpolate.LinearNDInterpolator(points, values)
         return self._ip_combined
 
-    def inside_bounds(self, xy):
+    def inside_bounds(
+        self,
+        xy: Union[
+            Tuple[int, int],
+            Tuple[float64, float],
+            Tuple[float64, float64],
+            Tuple[float, float64],
+        ],
+    ) -> Union[bool, bool_]:
         x, y = xy
         (xmin, xmax), (ymin, ymax) = self.bounds
         return xmin <= x <= xmax and ymin <= y <= ymax
 
-    def tell(self, point, value):
+    def tell(
+        self,
+        point: Union[
+            Tuple[int, int],
+            Tuple[float64, float],
+            Tuple[float64, float64],
+            Tuple[float, float64],
+        ],
+        value: Union[List[int], float64, float],
+    ) -> None:
         point = tuple(point)
         self.data[point] = value
         if not self.inside_bounds(point):
@@ -538,15 +566,43 @@ def tell(self, point, value):
         self._ip = None
         self._stack.pop(point, None)
 
-    def tell_pending(self, point):
+    def tell_pending(
+        self,
+        point: Union[
+            Tuple[int, int],
+            Tuple[float64, float],
+            Tuple[float64, float64],
+            Tuple[float, float64],
+        ],
+    ) -> None:
         point = tuple(point)
         if not self.inside_bounds(point):
             return
         self.pending_points.add(point)
         self._ip_combined = None
         self._stack.pop(point, None)
 
-    def _fill_stack(self, stack_till=1):
+    def _fill_stack(
+        self, stack_till: int = 1
+    ) -> Union[
+        Tuple[List[Tuple[float64, float64]], List[float64]],
+        Tuple[
+            List[
+                Union[
+                    Tuple[float64, float64],
+                    Tuple[float, float64],
+                    Tuple[float64, float],
+                ]
+            ],
+            List[float64],
+        ],
+        Tuple[
+            List[Union[Tuple[float, float64], Tuple[float64, float64]]], List[float64]
+        ],
+        Tuple[
+            List[Union[Tuple[float64, float64], Tuple[float64, float]]], List[float64]
+        ],
+    ]:
         if len(self.data) + len(self.pending_points) < self.ndim + 1:
             raise ValueError("too few points...")
 
@@ -585,7 +641,7 @@ def _fill_stack(self, stack_till=1):
 
         return points_new, losses_new
 
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True) -> Any:
         # Even if tell_pending is False we add the point such that _fill_stack
         # will return new points, later we remove these points if needed.
         points = list(self._stack.keys())
@@ -616,14 +672,14 @@ def ask(self, n, tell_pending=True):
         return points[:n], loss_improvements[:n]
 
     @cache_latest
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float64:
         if not self.bounds_are_done:
             return np.inf
         ip = self.interpolator(scaled=True) if real else self._interpolator_combined()
         losses = self.loss_per_triangle(ip)
         return losses.max()
 
-    def remove_unfinished(self):
+    def remove_unfinished(self) -> None:
         self.pending_points = set()
         for p in self._bounds_points:
             if p not in self.data:
@@ -697,10 +753,10 @@ def plot(self, n=None, tri_alpha=0):
 
         return im.opts(style=im_opts) * tris.opts(style=tri_opts, **no_hover)
 
-    def _get_data(self):
+    def _get_data(self) -> OrderedDict:
         return self.data
 
-    def _set_data(self, data):
+    def _set_data(self, data: OrderedDict) -> None:
         self.data = data
         # Remove points from stack if they already exist
         for point in copy(self._stack):