type fixes

Author: basnijholt

adaptive/learner/average_learner.py

@@ -35,7 +35,7 @@ class AverageLearner(BaseLearner):
 
     def __init__(
         self,
-        function: Callable[[BaseLearner], float],
+        function: Callable[["AverageLearner"], float],
         atol: Optional[float] = None,
         rtol: Optional[float] = None,
         min_npoints: int = 2,
@@ -49,7 +49,7 @@ def __init__(
 
         self.data = {}
         self.pending_points = set()
-        self.function = function
+        self.function = function  # type: ignore
         self.atol = atol
         self.rtol = rtol
         self.npoints = 0

adaptive/learner/balancing_learner.py

@@ -101,7 +101,7 @@ def __init__(
         # Naively we would make 'function' a method, but this causes problems
         # when using executors from 'concurrent.futures' because we have to
         # pickle the whole learner.
-        self.function = partial(dispatch, [l.function for l in self.learners])
+        self.function = partial(dispatch, [l.function for l in self.learners])  # type: ignore
 
         self._ask_cache = {}
         self._loss = {}

adaptive/learner/base_learner.py

@@ -82,8 +82,13 @@ class BaseLearner(metaclass=_RequireAttrsABCMeta):
     """
 
     data: dict
-    npoints: int
     pending_points: set
+    function: Callable
+
+    @property
+    @abc.abstractmethod
+    def npoints(self) -> int:
+        """Number of learned points."""
 
     def tell(self, x: Any, y) -> None:
         """Tell the learner about a single value.
@@ -149,7 +154,7 @@ def _get_data(self):
         pass
 
     @abc.abstractmethod
-    def _set_data(self):
+    def _set_data(self, data: Any):
         pass
 
     def copy_from(self, other):

adaptive/learner/integrator_learner.py

@@ -379,7 +379,7 @@ def __init__(self, function: Callable, bounds: Tuple[int, int], tol: float) -> N
         plot : hv.Scatter
             Plots all the points that are evaluated.
         """
-        self.function = function
+        self.function = function  # type: ignore
         self.bounds = bounds
         self.tol = tol
         self.max_ivals = 1000

adaptive/learner/learner1D.py

@@ -26,9 +26,11 @@
 from adaptive.notebook_integration import ensure_holoviews
 from adaptive.utils import cache_latest
 
+Point = Tuple[float, float]
+
 
 @uses_nth_neighbors(0)
-def uniform_loss(xs: Tuple[float, float], ys: Tuple[float, float]) -> float:
+def uniform_loss(xs: Point, ys: Any) -> float:
     """Loss function that samples the domain uniformly.
 
     Works with `~adaptive.Learner1D` only.
@@ -49,8 +51,8 @@ def uniform_loss(xs: Tuple[float, float], ys: Tuple[float, float]) -> float:
 
 @uses_nth_neighbors(0)
 def default_loss(
-    xs: Tuple[float, float],
-    ys: Union[Tuple[Iterable[float], Iterable[float]], Tuple[float, float]],
+    xs: Point,
+    ys: Union[Tuple[Iterable[float], Iterable[float]], Point],
 ) -> float:
     """Calculate loss on a single interval.
 
@@ -60,8 +62,8 @@ def default_loss(
     """
     dx = xs[1] - xs[0]
     if isinstance(ys[0], collections.abc.Iterable):
-        dy = [abs(a - b) for a, b in zip(*ys)]
-        return np.hypot(dx, dy).max()
+        dy_vec = [abs(a - b) for a, b in zip(*ys)]
+        return np.hypot(dx, dy_vec).max()
     else:
         dy = ys[1] - ys[0]
         return np.hypot(dx, dy)
@@ -200,7 +202,7 @@ def __init__(
         bounds: Tuple[float, float],
         loss_per_interval: Optional[Callable] = None,
     ) -> None:
-        self.function = function
+        self.function = function  # type: ignore
 
         if hasattr(loss_per_interval, "nth_neighbors"):
             self.nth_neighbors = loss_per_interval.nth_neighbors
@@ -238,7 +240,7 @@ def __init__(
 
         self.bounds = list(bounds)
 
-        self._vdim = None
+        self._vdim: Optional[int] = None
 
     @property
     def vdim(self) -> int:
@@ -565,7 +567,8 @@ def _ask_points_without_adding(self, n: int) -> Any:
         # Add bound intervals to quals if bounds were missing.
         if len(self.data) + len(self.pending_points) == 0:
             # We don't have any points, so return a linspace with 'n' points.
-            return np.linspace(*self.bounds, n).tolist(), [np.inf] * n
+            a, b = self.bounds
+            return np.linspace(a, b, n).tolist(), [np.inf] * n
 
         quals = loss_manager(self._scale[0])
         if len(missing_bounds) > 0:
@@ -601,7 +604,7 @@ def _ask_points_without_adding(self, n: int) -> Any:
                 quals[(*xs, n + 1)] = loss_qual * n / (n + 1)
 
         points = list(
-            itertools.chain.from_iterable(linspace(*ival, n) for (*ival, n) in quals)
+            itertools.chain.from_iterable(linspace(a, b, n) for ((a, b), n) in quals)
         )
 
         loss_improvements = list(
@@ -665,7 +668,8 @@ def _get_data(self) -> Dict[float, float]:
 
     def _set_data(self, data: Dict[float, float]) -> None:
         if data:
-            self.tell_many(*zip(*data.items()))
+            xs, ys = zip(*data.items())
+            self.tell_many(xs, ys)
 
     def __getstate__(self):
         return (

adaptive/learner/learner2D.py

@@ -372,7 +372,7 @@ def __init__(
 
         self._bounds_points = list(itertools.product(*bounds))
         self._stack.update({p: np.inf for p in self._bounds_points})
-        self.function = function
+        self.function = function  # type: ignore
         self._ip = self._ip_combined = None
 
         self.stack_size = 10
@@ -600,7 +600,7 @@ def _fill_stack(
 
     def ask(
         self, n: int, tell_pending: bool = True
-    ) -> Tuple[List[Tuple[float, float]], List[float]]:
+    ) -> Tuple[List[Union[Tuple[float, float], np.array]], List[float]]:
         # 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())

adaptive/learner/learnerND.py

@@ -305,7 +305,7 @@ class LearnerND(BaseLearner):
 
     def __init__(
         self,
-        func: Callable,
+        function: Callable,
         bounds: Union[Sequence[Tuple[float, float]], ConvexHull],
         loss_per_simplex: Optional[Callable] = None,
     ) -> None:
@@ -339,8 +339,8 @@ def __init__(
 
         self.ndim = len(self._bbox)
 
-        self.function = func
-        self._tri = None
+        self.function = function  # type: ignore
+        self._tri: Optional[Triangulation] = None
         self._losses: Dict[Simplex, float] = dict()
 
         self._pending_to_simplex: Dict[Point, Simplex] = dict()  # vertex → simplex
@@ -455,6 +455,7 @@ def tell(self, point: Tuple[float, ...], value: Union[float, np.ndarray]) -> Non
         self._update_range(value)
         if tri is not None:
             simplex = self._pending_to_simplex.get(point)
+            assert self.tri is not None
             if simplex is not None and not self._simplex_exists(simplex):
                 simplex = None
             to_delete, to_add = tri.add_point(point, simplex, transform=self._transform)

adaptive/learner/sequence_learner.py

@@ -18,7 +18,7 @@ class _IgnoreFirstArgument:
     """
 
     def __init__(self, function: Callable) -> None:
-        self.function = function
+        self.function = function  # type: ignore
 
     def __call__(
         self, index_point: Tuple[int, Union[float, np.ndarray]], *args, **kwargs
@@ -62,7 +62,7 @@ class SequenceLearner(BaseLearner):
 
     def __init__(self, function: Callable, sequence: Iterable) -> None:
         self._original_function = function
-        self.function = _IgnoreFirstArgument(function)
+        self.function = _IgnoreFirstArgument(function)  # type: ignore
         self._to_do_indices = SortedSet({i for i, _ in enumerate(sequence)})
         self._ntotal = len(sequence)
         self.sequence = copy(sequence)

adaptive/learner/skopt_learner.py

@@ -25,7 +25,7 @@ class SKOptLearner(Optimizer, BaseLearner):
     """
 
     def __init__(self, function: Callable, **kwargs) -> None:
-        self.function = function
+        self.function = function  # type: ignore
         self.pending_points = set()
         self.data = collections.OrderedDict()
         super().__init__(**kwargs)

adaptive/learner/triangulation.py

@@ -27,9 +27,9 @@
 from numpy.linalg import matrix_rank, norm, slogdet, solve
 
 SimplexPoints = Union[List[Tuple[float, ...]], ndarray]
-Simplex = Union[Iterable[numbers.Integral], ndarray]
+Simplex = Union[Sequence[numbers.Integral], ndarray]
 Point = Union[Tuple[float, ...], ndarray]
-Points = Union[Sequence[Point], ndarray]
+Points = Union[Sequence[Tuple[float, ...]], ndarray]
 
 
 def fast_norm(v: Union[Tuple[float, ...], ndarray]) -> float:
@@ -168,7 +168,7 @@ def fast_det(matrix: ndarray) -> float:
         return ndet(matrix)
 
 
-def circumsphere(pts: ndarray) -> Tuple[Tuple[float, ...], float]:
+def circumsphere(pts: Simplex) -> Tuple[Tuple[float, ...], float]:
     """Compute the center and radius of a N dimension sphere which touches each point in pts.
 
     Parameters