type hint fixes for adaptive/learner/learner2D.py

Author: basnijholt

adaptive/learner/learner2D.py

@@ -2,9 +2,8 @@
 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
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 from scipy import interpolate
@@ -107,7 +106,9 @@ def uniform_loss(ip: LinearNDInterpolator) -> np.ndarray:
     return np.sqrt(areas(ip))
 
 
-def resolution_loss_function(min_distance: int = 0, max_distance: int = 1) -> Callable:
+def resolution_loss_function(
+    min_distance: float = 0, max_distance: float = 1
+) -> Callable:
     """Loss function that is similar to the `default_loss` function, but you
     can set the maximimum and minimum size of a triangle.
 
@@ -353,10 +354,8 @@ class Learner2D(BaseLearner):
 
     def __init__(
         self,
-        function: Union[partial, Callable],
-        bounds: Union[
-            List[Tuple[int, int]], Tuple[Tuple[int, int], Tuple[int, int]], np.ndarray
-        ],
+        function: Callable,
+        bounds: Tuple[Tuple[int, int], Tuple[int, int]],
         loss_per_triangle: Optional[Callable] = None,
     ) -> None:
         self.ndim = len(bounds)
@@ -462,7 +461,7 @@ def _data_in_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
             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) -> Any:
+    def _data_interp(self) -> Tuple[np.ndarray, np.ndarray]:
         if self.pending_points:
             points = list(self.pending_points)
             if self.bounds_are_done:
@@ -493,7 +492,7 @@ def ip(self):
         )
         return self.interpolator(scaled=True)
 
-    def interpolator(self, *, scaled=False) -> LinearNDInterpolator:
+    def interpolator(self, *, scaled: bool = False) -> LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data.
 
@@ -534,28 +533,13 @@ def _interpolator_combined(self) -> LinearNDInterpolator:
             self._ip_combined = interpolate.LinearNDInterpolator(points, values)
         return self._ip_combined
 
-    def inside_bounds(
-        self,
-        xy: Union[
-            Tuple[int, int],
-            Tuple[float, float],
-            Tuple[float, float],
-            Tuple[float, float],
-        ],
-    ) -> Union[bool, np.bool_]:
+    def inside_bounds(self, xy: Tuple[float, float],) -> Union[bool, np.bool_]:
         x, y = xy
         (xmin, xmax), (ymin, ymax) = self.bounds
         return xmin <= x <= xmax and ymin <= y <= ymax
 
     def tell(
-        self,
-        point: Union[
-            Tuple[int, int],
-            Tuple[float, float],
-            Tuple[float, float],
-            Tuple[float, float],
-        ],
-        value: Union[List[int], float, float],
+        self, point: Tuple[float, float], value: Union[float, Iterable[float]],
     ) -> None:
         point = tuple(point)
         self.data[point] = value
@@ -565,15 +549,7 @@ def tell(
         self._ip = None
         self._stack.pop(point, None)
 
-    def tell_pending(
-        self,
-        point: Union[
-            Tuple[int, int],
-            Tuple[float, float],
-            Tuple[float, float],
-            Tuple[float, float],
-        ],
-    ) -> None:
+    def tell_pending(self, point: Tuple[float, float],) -> None:
         point = tuple(point)
         if not self.inside_bounds(point):
             return
@@ -622,7 +598,7 @@ def _fill_stack(
 
         return points_new, losses_new
 
-    def ask(self, n: int, tell_pending: bool = True) -> Any:
+    def ask(self, n: int, tell_pending: bool = True) -> Tuple[np.ndarray, np.ndarray]:
         # 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())