typing fixes

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -1,4 +1,5 @@
 import itertools
+import numbers
 from collections import defaultdict
 from collections.abc import Iterable
 from contextlib import suppress
@@ -210,12 +211,12 @@ def _ask_and_tell_based_on_loss(
         return points, loss_improvements
 
     def _ask_and_tell_based_on_npoints(
-        self, n: int
-    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
+        self, n: numbers.Integral
+    ) -> Tuple[List[Tuple[numbers.Integral, Any]], List[float]]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
-            index = int(np.argmin(total_points))
+            index = np.argmin(total_points)
             # Take the points from the cache
             if index not in self._ask_cache:
                 self._ask_cache[index] = self.learners[index].ask(n=1)
@@ -229,7 +230,7 @@ def _ask_and_tell_based_on_npoints(
 
     def _ask_and_tell_based_on_cycle(
         self, n: int
-    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
+    ) -> Tuple[List[Tuple[numbers.Integral, Any]], List[float]]:
         points, loss_improvements = [], []
         for _ in range(n):
             index = next(self._cycle)
@@ -242,7 +243,7 @@ def _ask_and_tell_based_on_cycle(
 
     def ask(
         self, n: int, tell_pending: bool = True
-    ) -> Tuple[List[Tuple[int, Any]], List[float]]:
+    ) -> Tuple[List[Tuple[numbers.Integral, Any]], List[float]]:
         """Chose points for learners."""
         if n == 0:
             return [], []
@@ -253,14 +254,14 @@ def ask(
         else:
             return self._ask_and_tell(n)
 
-    def tell(self, x: Tuple[int, Any], y: Any) -> None:
+    def tell(self, x: Tuple[numbers.Integral, Any], y: Any) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self._pending_loss.pop(index, None)
         self.learners[index].tell(x, y)
 
-    def tell_pending(self, x: Tuple[int, Any]) -> None:
+    def tell_pending(self, x: Tuple[numbers.Integral, Any]) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
         self._loss.pop(index, None)

adaptive/learner/learner1D.py

@@ -1,6 +1,7 @@
 import collections.abc
 import itertools
 import math
+import numbers
 from copy import deepcopy
 from typing import (
     Any,
@@ -414,7 +415,9 @@ def _update_scale(self, x: float, y: Union[float, np.ndarray]) -> None:
                 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: float, y: Union[float, Sequence[float], np.ndarray]) -> None:
+    def tell(
+        self, x: float, y: Union[float, Sequence[numbers.Number], np.ndarray]
+    ) -> None:
         if x in self.data:
             # The point is already evaluated before
             return

adaptive/learner/learnerND.py

@@ -3,7 +3,7 @@
 import random
 from collections import OrderedDict
 from collections.abc import Iterable
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import scipy.spatial
@@ -306,7 +306,7 @@ class LearnerND(BaseLearner):
     def __init__(
         self,
         func: Callable,
-        bounds: Union[Tuple[Tuple[float, float], ...], ConvexHull],
+        bounds: Union[Sequence[Tuple[float, float]], ConvexHull],
         loss_per_simplex: Optional[Callable] = None,
     ) -> None:
         self._vdim = None

adaptive/learner/triangulation.py

@@ -1,9 +1,11 @@
 import collections.abc
+import numbers
 from collections import Counter
 from itertools import chain, combinations
 from math import factorial, sqrt
 from typing import Any, Iterable, Iterator, List, Optional, Sequence, Set, Tuple, Union
 
+import numpy as np
 import scipy.spatial
 from numpy import abs as np_abs
 from numpy import (
@@ -25,7 +27,7 @@
 from numpy.linalg import matrix_rank, norm, slogdet, solve
 
 SimplexPoints = Union[List[Tuple[float, ...]], ndarray]  # XXX: check if this is correct
-Simplex = Tuple[int, ...]
+Simplex = Union[Tuple[numbers.Integral, ...], ndarray]
 Point = Union[Tuple[float, ...], ndarray]  # XXX: check if this is correct
 
 
@@ -49,7 +51,7 @@ def fast_norm(v: Union[Tuple[float, ...], ndarray]) -> float:
 
 def fast_2d_point_in_simplex(
     point: Point, simplex: SimplexPoints, eps: float = 1e-8
-) -> bool:
+) -> Union[np.bool_, bool]:
     (p0x, p0y), (p1x, p1y), (p2x, p2y) = simplex
     px, py = point
 
@@ -63,7 +65,9 @@ def fast_2d_point_in_simplex(
     return (t >= -eps) and (s + t <= 1 + eps)
 
 
-def point_in_simplex(point: Point, simplex: SimplexPoints, eps: float = 1e-8) -> bool:
+def point_in_simplex(
+    point: Point, simplex: SimplexPoints, eps: float = 1e-8
+) -> Union[np.bool_, bool]:
     if len(point) == 2:
         return fast_2d_point_in_simplex(point, simplex, eps)
 
@@ -322,7 +326,7 @@ class Triangulation:
         or more simplices in the
     """
 
-    def __init__(self, coords: Sequence[Point]) -> None:
+    def __init__(self, coords: Union[Sequence[Point], ndarray]) -> None:
         if not is_iterable_and_sized(coords):
             raise TypeError("Please provide a 2-dimensional list of points")
         coords = list(coords)
@@ -373,10 +377,12 @@ def add_simplex(self, simplex: Simplex) -> None:
         for vertex in simplex:
             self.vertex_to_simplices[vertex].add(simplex)
 
-    def get_vertices(self, indices: Sequence[int]) -> List[Optional[Point]]:
+    def get_vertices(
+        self, indices: Sequence[numbers.Integral]
+    ) -> List[Optional[Point]]:
         return [self.get_vertex(i) for i in indices]
 
-    def get_vertex(self, index: Optional[int]) -> Optional[Point]:
+    def get_vertex(self, index: Optional[numbers.Integral]) -> Optional[Point]:
         if index is None:
             return None
         return self.vertices[index]
@@ -410,7 +416,7 @@ def get_reduced_simplex(
 
     def point_in_simplex(
         self, point: Point, simplex: Simplex, eps: float = 1e-8
-    ) -> bool:
+    ) -> Union[np.bool_, bool]:
         vertices = self.get_vertices(simplex)
         return point_in_simplex(point, vertices, eps)
 
@@ -616,7 +622,7 @@ def add_point(
         point: Point,
         simplex: Optional[Simplex] = None,
         transform: Optional[ndarray] = None,
-    ) -> Any:
+    ) -> Tuple[Set[Simplex], Set[Simplex]]:
         """Add a new vertex and create simplices as appropriate.
 
         Parameters