type hint fixes for adaptive/learner/learnerND.py

Author: basnijholt

adaptive/learner/learnerND.py

@@ -29,15 +29,7 @@ def to_list(inp: float) -> List[float]:
     return [inp]
 
 
-def volume(
-    simplex: Union[
-        List[Tuple[float, float]],
-        List[Tuple[float, float]],
-        List[Tuple[float, float]],
-        np.ndarray,
-    ],
-    ys: None = None,
-) -> float:
+def volume(simplex: List[Tuple[float, float]], ys: None = None,) -> float:
     # Notice the parameter ys is there so you can use this volume method as
     # as loss function
     matrix = np.subtract(simplex[:-1], simplex[-1], dtype=float)
@@ -207,13 +199,7 @@ def curvature_loss(simplex, values, value_scale, neighbors, neighbor_values):
 
 
 def choose_point_in_simplex(
-    simplex: Union[
-        List[Union[Tuple[int, int], Tuple[float, float]]],
-        List[Union[Tuple[float, float, float], Tuple[int, int, int]]],
-        List[Tuple[float, float, float]],
-        List[Tuple[float, float]],
-    ],
-    transform: Optional[np.ndarray] = None,
+    simplex: np.ndarray, transform: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     """Choose a new point in inside a simplex.
 
@@ -318,13 +304,7 @@ class LearnerND(BaseLearner):
     def __init__(
         self,
         func: Callable,
-        bounds: Union[
-            Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]],
-            np.ndarray,
-            Tuple[Tuple[int, int], Tuple[int, int]],
-            List[Tuple[int, int]],
-            ConvexHull,
-        ],
+        bounds: Union[Tuple[Tuple[float, float], ...], ConvexHull],
         loss_per_simplex: Optional[Callable] = None,
     ) -> None:
         self._vdim = None
@@ -452,17 +432,7 @@ def points(self) -> np.ndarray:
         """Get the points from `data` as a numpy array."""
         return np.array(list(self.data.keys()), dtype=float)
 
-    def tell(
-        self,
-        point: Union[
-            Tuple[float, float],
-            Tuple[int, int],
-            Tuple[int, int, int],
-            Tuple[float, float, float],
-            Tuple[float, float, float],
-        ],
-        value: Union[List[int], float, float, np.ndarray],
-    ) -> None:
+    def tell(self, point: Tuple[float, ...], value: Union[float, np.ndarray],) -> None:
         point = tuple(point)
 
         if point in self.data:
@@ -486,20 +456,11 @@ def tell(
             to_delete, to_add = tri.add_point(point, simplex, transform=self._transform)
             self._update_losses(to_delete, to_add)
 
-    def _simplex_exists(self, simplex: Any) -> bool:
+    def _simplex_exists(self, simplex: Any) -> bool:  # XXX: specify simplex: Any
         simplex = tuple(sorted(simplex))
         return simplex in self.tri.simplices
 
-    def inside_bounds(
-        self,
-        point: Union[
-            Tuple[float, float],
-            Tuple[float, float, float],
-            Tuple[int, int, int],
-            Tuple[int, int],
-            Tuple[float, float, float],
-        ],
-    ) -> Union[bool, np.bool_]:
+    def inside_bounds(self, point: Tuple[float, ...],) -> Union[bool, np.bool_]:
         """Check whether a point is inside the bounds."""
         if hasattr(self, "_interior"):
             return self._interior.find_simplex(point, tol=1e-8) >= 0
@@ -509,17 +470,7 @@ def inside_bounds(
                 (mn - eps) <= p <= (mx + eps) for p, (mn, mx) in zip(point, self._bbox)
             )
 
-    def tell_pending(
-        self,
-        point: Union[
-            Tuple[int, int],
-            Tuple[float, float, float],
-            Tuple[float, float],
-            Tuple[int, int, int],
-        ],
-        *,
-        simplex=None,
-    ) -> None:
+    def tell_pending(self, point: Tuple[float, ...], *, simplex=None,) -> None:
         point = tuple(point)
         if not self.inside_bounds(point):
             return
@@ -547,9 +498,7 @@ def tell_pending(
             self._update_subsimplex_losses(simpl, to_add)
 
     def _try_adding_pending_point_to_simplex(
-        self,
-        point: Union[Tuple[float, float, float], Tuple[float, float]],
-        simplex: Any,
+        self, point: Tuple[float, ...], simplex: Any,  # XXX: specify simplex: Any
     ) -> Any:
         # try to insert it
         if not self.tri.point_in_simplex(point, simplex):
@@ -562,7 +511,9 @@ def _try_adding_pending_point_to_simplex(
         self._pending_to_simplex[point] = simplex
         return self._subtriangulations[simplex].add_point(point)
 
-    def _update_subsimplex_losses(self, simplex: Any, new_subsimplices: Any) -> None:
+    def _update_subsimplex_losses(
+        self, simplex: Any, new_subsimplices: Any
+    ) -> None:  # XXX: specify simplex: Any
         loss = self._losses[simplex]
 
         loss_density = loss / self.tri.volume(simplex)
@@ -583,14 +534,7 @@ def ask(self, n: int, tell_pending: bool = True) -> Any:
         else:
             return self._ask_and_tell_pending(n)
 
-    def _ask_bound_point(
-        self,
-    ) -> Union[
-        Tuple[Tuple[int, int, int], float],
-        Tuple[Tuple[int, int], float],
-        Tuple[Tuple[float, float], float],
-        Tuple[Tuple[float, float, float], float],
-    ]:
+    def _ask_bound_point(self,) -> Tuple[Tuple[float, ...], float]:
         # get the next bound point that is still available
         new_point = next(
             p
@@ -600,11 +544,7 @@ def _ask_bound_point(
         self.tell_pending(new_point)
         return new_point, np.inf
 
-    def _ask_point_without_known_simplices(
-        self,
-    ) -> Union[
-        Tuple[Tuple[float, float], float], Tuple[Tuple[float, float, float], float],
-    ]:
+    def _ask_point_without_known_simplices(self,) -> Tuple[Tuple[float, ...], float]:
         assert not self._bounds_available
         # pick a random point inside the bounds
         # XXX: change this into picking a point based on volume loss
@@ -645,11 +585,7 @@ def _pop_highest_existing_simplex(self) -> Any:
             "  be a simplex available if LearnerND.tri() is not None."
         )
 
-    def _ask_best_point(
-        self,
-    ) -> Union[
-        Tuple[Tuple[float, float], float], Tuple[Tuple[float, float, float], float],
-    ]:
+    def _ask_best_point(self,) -> Tuple[Tuple[float, ...], float]:
         assert self.tri is not None
 
         loss, simplex, subsimplex = self._pop_highest_existing_simplex()
@@ -676,14 +612,7 @@ def _bounds_available(self) -> bool:
             for p in self._bounds_points
         )
 
-    def _ask(
-        self,
-    ) -> Union[
-        Tuple[Tuple[int, int, int], float],
-        Tuple[Tuple[float, float, float], float],
-        Tuple[Tuple[float, float], float],
-        Tuple[Tuple[int, int], float],
-    ]:
+    def _ask(self,) -> Tuple[Tuple[float, ...], float]:
         if self._bounds_available:
             return self._ask_bound_point()  # O(1)
 
@@ -695,7 +624,7 @@ def _ask(
 
         return self._ask_best_point()  # O(log N)
 
-    def _compute_loss(self, simplex: Any) -> float:
+    def _compute_loss(self, simplex: Any) -> float:  # XXX: specify simplex: Any
         # get the loss
         vertices = self.tri.get_vertices(simplex)
         values = [self.data[tuple(v)] for v in vertices]