type hint fixes for adaptive/learner/integrator_learner.py

Author: basnijholt

adaptive/learner/integrator_learner.py

@@ -2,13 +2,11 @@
 
 import sys
 from collections import defaultdict
-from functools import partial
 from math import sqrt
 from operator import attrgetter
-from typing import Any, Callable, List, Optional, Set, Tuple, Union
+from typing import Callable, List, Optional, Set, Tuple, Union
 
 import numpy as np
-from numpy import ufunc
 from scipy.linalg import norm
 from sortedcontainers import SortedSet
 
@@ -142,11 +140,7 @@ class _Interval:
     ]
 
     def __init__(
-        self,
-        a: Union[int, np.float64],
-        b: Union[int, np.float64],
-        depth: int,
-        rdepth: int,
+        self, a: Union[int, float], b: Union[int, float], depth: int, rdepth: int,
     ) -> None:
         self.children = []
         self.data = {}
@@ -213,7 +207,7 @@ def split(self) -> List["_Interval"]:
     def calc_igral(self) -> None:
         self.igral = (self.b - self.a) * self.c[0] / sqrt(2)
 
-    def update_heuristic_err(self, value: Union[np.float64, float]) -> None:
+    def update_heuristic_err(self, value: float) -> None:
         """Sets the error of an interval using a heuristic (half the error of
         the parent) when the actual error cannot be calculated due to its
         parents not being finished yet. This error is propagated down to its
@@ -347,10 +341,7 @@ def __repr__(self) -> str:
 
 class IntegratorLearner(BaseLearner):
     def __init__(
-        self,
-        function: Union[partial, ufunc, Callable],
-        bounds: Tuple[int, int],
-        tol: float,
+        self, function: Callable, bounds: Tuple[int, int], tol: float,
     ) -> None:
         """
         Parameters
@@ -403,7 +394,7 @@ def __init__(
     def approximating_intervals(self) -> Set["_Interval"]:
         return self.first_ival.done_leaves
 
-    def tell(self, point: np.float64, value: np.float64) -> None:
+    def tell(self, point: float, value: float) -> None:
         if point not in self.x_mapping:
             raise ValueError(f"Point {point} doesn't belong to any interval")
         self.data[point] = value
@@ -460,23 +451,15 @@ def add_ival(self, ival: "_Interval") -> None:
                 self._stack.append(x)
         self.ivals.add(ival)
 
-    def ask(
-        self, n: int, tell_pending: bool = True
-    ) -> Union[
-        Tuple[List[np.float64], List[np.float64]], Tuple[List[np.float64], List[float]]
-    ]:
+    def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[float], List[float]]:
         """Choose points for learners."""
         if not tell_pending:
             with restore(self):
                 return self._ask_and_tell_pending(n)
         else:
             return self._ask_and_tell_pending(n)
 
-    def _ask_and_tell_pending(
-        self, n: int
-    ) -> Union[
-        Tuple[List[np.float64], List[np.float64]], Tuple[List[np.float64], List[float]]
-    ]:
+    def _ask_and_tell_pending(self, n: int) -> Tuple[List[float], List[float]]:
         points, loss_improvements = self.pop_from_stack(n)
         n_left = n - len(points)
         while n_left > 0:
@@ -492,13 +475,7 @@ def _ask_and_tell_pending(
 
         return points, loss_improvements
 
-    def pop_from_stack(
-        self, n: int
-    ) -> Union[
-        Tuple[List[np.float64], List[np.float64]],
-        Tuple[List[Any], List[Any]],
-        Tuple[List[np.float64], List[float]],
-    ]:
+    def pop_from_stack(self, n: int) -> Tuple[List[float], List[float]]:
         points = self._stack[:n]
         self._stack = self._stack[n:]
         loss_improvements = [
@@ -509,7 +486,7 @@ def pop_from_stack(
     def remove_unfinished(self):
         pass
 
-    def _fill_stack(self) -> List[np.float64]:
+    def _fill_stack(self) -> List[float]:
         # XXX: to-do if all the ivals have err=inf, take the interval
         # with the lowest rdepth and no children.
         force_split = bool(self.priority_split)
@@ -550,11 +527,11 @@ def npoints(self) -> int:
         return len(self.data)
 
     @property
-    def igral(self) -> np.float64:
+    def igral(self) -> float:
         return sum(i.igral for i in self.approximating_intervals)
 
     @property
-    def err(self) -> np.float64:
+    def err(self) -> float:
         if self.approximating_intervals:
             err = sum(i.err for i in self.approximating_intervals)
             if err > sys.float_info.max: