type hint fixes for adaptive/learner/learner1D.py

Author: basnijholt

adaptive/learner/learner1D.py

@@ -2,12 +2,9 @@
 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 sortedcollections.recipes import ItemSortedDict
 from sortedcontainers.sorteddict import SortedDict
 
@@ -19,7 +16,7 @@
 
 
 @uses_nth_neighbors(0)
-def uniform_loss(xs: Tuple[float, float], ys: Tuple[float, float],) -> float:
+def uniform_loss(xs: Tuple[float, float], ys: Tuple[float, float]) -> float:
     """Loss function that samples the domain uniformly.
 
     Works with `~adaptive.Learner1D` only.
@@ -59,7 +56,7 @@ def default_loss(
 
 
 @uses_nth_neighbors(1)
-def triangle_loss(xs: Any, ys: Any) -> float:
+def triangle_loss(xs: Tuple[float], ys: Tuple[Union[float, np.ndarray]]) -> float:
     xs = [x for x in xs if x is not None]
     ys = [y for y in ys if y is not None]
 
@@ -77,7 +74,7 @@ def triangle_loss(xs: Any, ys: Any) -> float:
 
 
 def curvature_loss_function(
-    area_factor: int = 1, euclid_factor: float = 0.02, horizontal_factor: float = 0.02
+    area_factor: float = 1, euclid_factor: float = 0.02, horizontal_factor: float = 0.02
 ) -> Callable:
     # XXX: add a doc-string
     @uses_nth_neighbors(1)
@@ -97,9 +94,7 @@ def curvature_loss(xs, ys):
     return curvature_loss
 
 
-def linspace(
-    x_left: Union[int, float], x_right: Union[int, float], n: int,
-) -> Union[List[float], List[float]]:
+def linspace(x_left: float, x_right: float, n: int,) -> List[float]:
     """This is equivalent to
     'np.linspace(x_left, x_right, n, endpoint=False)[1:]',
     but it is 15-30 times faster for small 'n'."""
@@ -118,12 +113,10 @@ def _get_neighbors_from_list(xs: np.ndarray) -> SortedDict:
     xs_left[0] = None
     xs_right[-1] = None
     neighbors = {x: [x_L, x_R] for x, x_L, x_R in zip(xs, xs_left, xs_right)}
-    return sortedcontainers.SortedDict(neighbors)
+    return SortedDict(neighbors)
 
 
-def _get_intervals(
-    x: Union[int, float], neighbors: SortedDict, nth_neighbors: int
-) -> Any:
+def _get_intervals(x: float, neighbors: SortedDict, nth_neighbors: int) -> Any:
     nn = nth_neighbors
     i = neighbors.index(x)
     start = max(0, i - nn - 1)
@@ -178,8 +171,8 @@ class Learner1D(BaseLearner):
 
     def __init__(
         self,
-        function: Union[Callable, partial],
-        bounds: Union[Tuple[int, int], Tuple[float, float], np.ndarray],
+        function: Callable,
+        bounds: Tuple[float, float],
         loss_per_interval: Optional[Callable] = None,
     ) -> None:
         self.function = function
@@ -201,8 +194,8 @@ def __init__(
 
         # A dict {x_n: [x_{n-1}, x_{n+1}]} for quick checking of local
         # properties.
-        self.neighbors = sortedcontainers.SortedDict()
-        self.neighbors_combined = sortedcontainers.SortedDict()
+        self.neighbors = SortedDict()
+        self.neighbors_combined = SortedDict()
 
         # Bounding box [[minx, maxx], [miny, maxy]].
         self._bbox = [list(bounds), [np.inf, -np.inf]]
@@ -248,34 +241,32 @@ def npoints(self) -> int:
         return len(self.data)
 
     @cache_latest
-    def loss(self, real: bool = True) -> Union[int, float]:
+    def loss(self, real: bool = True) -> 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: Optional[Union[float, int]]) -> Optional[float]:
+    def _scale_x(self, x: Optional[float]) -> Optional[float]:
         if x is None:
             return None
         return x / self._scale[0]
 
     def _scale_y(
-        self, y: Optional[Union[int, np.ndarray, float, float]]
+        self, y: Optional[Union[float, np.ndarray]]
     ) -> Optional[Union[float, np.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: int) -> Optional[Union[int, float, float]]:
+    def _get_point_by_index(self, ind: int) -> Optional[float]:
         if ind < 0 or ind >= len(self.neighbors):
             return None
         return self.neighbors.keys()[ind]
 
-    def _get_loss_in_interval(
-        self, x_left: Union[int, float], x_right: Union[int, float],
-    ) -> Union[int, float]:
+    def _get_loss_in_interval(self, x_left: float, x_right: float,) -> float:
         assert x_left is not None and x_right is not None
 
         if x_right - x_left < self._dx_eps:
@@ -296,7 +287,7 @@ def _get_loss_in_interval(
         return self.loss_per_interval(xs_scaled, ys_scaled)
 
     def _update_interpolated_loss_in_interval(
-        self, x_left: Union[int, float], x_right: Union[int, float],
+        self, x_left: float, x_right: float,
     ) -> None:
         if x_left is None or x_right is None:
             return
@@ -313,7 +304,7 @@ def _update_interpolated_loss_in_interval(
             self.losses_combined[a, b] = (b - a) * loss / dx
             a = b
 
-    def _update_losses(self, x: Union[int, float], real: bool = True) -> None:
+    def _update_losses(self, x: 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'.
@@ -356,7 +347,7 @@ def _update_losses(self, x: Union[int, float], real: bool = True) -> None:
             self.losses_combined[x, b] = float("inf")
 
     @staticmethod
-    def _find_neighbors(x: Union[int, float], neighbors: SortedDict) -> Any:
+    def _find_neighbors(x: float, neighbors: SortedDict) -> Any:
         if x in neighbors:
             return neighbors[x]
         pos = neighbors.bisect_left(x)
@@ -365,16 +356,14 @@ def _find_neighbors(x: Union[int, float], neighbors: SortedDict) -> Any:
         x_right = keys[pos] if pos != len(neighbors) else None
         return x_left, x_right
 
-    def _update_neighbors(self, x: Union[int, float], neighbors: SortedDict) -> None:
+    def _update_neighbors(self, x: 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: Union[int, float], y: Union[float, int, float, np.ndarray],
-    ) -> None:
+    def _update_scale(self, x: float, y: Union[float, np.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
@@ -401,7 +390,7 @@ def _update_scale(
                 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: Union[int, float], y: Any) -> None:
+    def tell(self, x: float, y: Union[float, np.ndarray]) -> None:
         if x in self.data:
             # The point is already evaluated before
             return
@@ -436,15 +425,15 @@ def tell(self, x: Union[int, float], y: Any) -> None:
 
             self._oldscale = deepcopy(self._scale)
 
-    def tell_pending(self, x: Union[int, float]) -> None:
+    def tell_pending(self, x: 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: Any, ys: Any, *, force=False) -> None:
+    def tell_many(self, xs: List[float], ys: List[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
@@ -644,24 +633,24 @@ def remove_unfinished(self) -> None:
         self.losses_combined = deepcopy(self.losses)
         self.neighbors_combined = deepcopy(self.neighbors)
 
-    def _get_data(self) -> Dict[Union[int, float], float]:
+    def _get_data(self) -> Dict[float, float]:
         return self.data
 
-    def _set_data(self, data: Dict[Union[int, float], float]) -> None:
+    def _set_data(self, data: Dict[float, float]) -> None:
         if data:
             self.tell_many(*zip(*data.items()))
 
 
-def loss_manager(x_scale: Union[int, float]) -> ItemSortedDict:
+def loss_manager(x_scale: float) -> ItemSortedDict:
     def sort_key(ival, loss):
         loss, ival = finite_loss(ival, loss, x_scale)
         return -loss, ival
 
-    sorted_dict = sortedcollections.ItemSortedDict(sort_key)
+    sorted_dict = ItemSortedDict(sort_key)
     return sorted_dict
 
 
-def finite_loss(ival: Any, loss: Union[int, float], x_scale: Union[int, float],) -> Any:
+def finite_loss(ival: Any, loss: float, x_scale: 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