use data types from np namespace instead importing it separately

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -7,14 +7,15 @@
 from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import numpy as np
-from numpy import float64, int64
 
 from adaptive.learner.base_learner import BaseLearner
 from adaptive.notebook_integration import ensure_holoviews
 from adaptive.utils import cache_latest, named_product, restore
 
 
-def dispatch(child_functions: List[Callable], arg: Any,) -> Union[int, float64, float]:
+def dispatch(
+    child_functions: List[Callable], arg: Any,
+) -> Union[int, np.float64, float]:
     index, x = arg
     return child_functions[index](x)
 
@@ -165,7 +166,7 @@ def _ask_and_tell_based_on_loss_improvements(self, n: int) -> Any:
     def _ask_and_tell_based_on_loss(
         self, n: int
     ) -> Union[
-        Tuple[List[Tuple[int, float]], List[float64]],
+        Tuple[List[Tuple[int, float]], List[np.float64]],
         Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
         Tuple[List[Tuple[int, int]], List[float]],
     ]:
@@ -192,9 +193,9 @@ def _ask_and_tell_based_on_loss(
     def _ask_and_tell_based_on_npoints(
         self, n: int
     ) -> Union[
-        Tuple[List[Union[Tuple[int64, int], Tuple[int64, float]]], List[float]],
-        Tuple[List[Tuple[int64, float]], List[float64]],
-        Tuple[List[Tuple[int64, int]], List[float]],
+        Tuple[List[Union[Tuple[np.int64, int], Tuple[np.int64, float]]], List[float]],
+        Tuple[List[Tuple[np.int64, float]], List[np.float64]],
+        Tuple[List[Tuple[np.int64, int]], List[float]],
     ]:
         selected = []  # tuples ((learner_index, point), loss_improvement)
         total_points = [l.npoints + len(l.pending_points) for l in self.learners]
@@ -214,7 +215,7 @@ def _ask_and_tell_based_on_npoints(
     def _ask_and_tell_based_on_cycle(
         self, n: int
     ) -> Union[
-        Tuple[List[Tuple[int, float]], List[float64]],
+        Tuple[List[Tuple[int, float]], List[np.float64]],
         Tuple[List[Union[Tuple[int, int], Tuple[int, float]]], List[float]],
         Tuple[List[Tuple[int, int]], List[float]],
     ]:
@@ -240,7 +241,9 @@ def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[Any], List[float]
             return self._ask_and_tell(n)
 
     def tell(
-        self, x: Any, y: Union[int, float64, float, Tuple[int, int], Tuple[int64, int]]
+        self,
+        x: Any,
+        y: Union[int, np.float64, float, Tuple[int, int], Tuple[np.int64, int]],
     ) -> None:
         index, x = x
         self._ask_cache.pop(index, None)
@@ -256,7 +259,7 @@ def tell_pending(self, x: Any) -> None:
 
     def _losses(
         self, real: bool = True
-    ) -> Union[List[float], List[float64], List[Union[float, float64]]]:
+    ) -> Union[List[float], List[np.float64], List[Union[float, np.float64]]]:
         losses = []
         loss_dict = self._loss if real else self._pending_loss
 
@@ -268,7 +271,7 @@ def _losses(
         return losses
 
     @cache_latest
-    def loss(self, real: bool = True) -> Union[float64, float]:
+    def loss(self, real: bool = True) -> Union[np.float64, float]:
         losses = self._losses(real)
         return max(losses)
 

adaptive/learner/integrator_coeffs.py

@@ -6,7 +6,6 @@
 
 import numpy as np
 import scipy.linalg
-from numpy import ndarray
 
 
 def legendre(n: int) -> List[List[Fraction]]:
@@ -30,7 +29,7 @@ def legendre(n: int) -> List[List[Fraction]]:
     return result
 
 
-def newton(n: int) -> ndarray:
+def newton(n: int) -> np.ndarray:
     """Compute the monomial coefficients of the Newton polynomial over the
     nodes of the n-point Clenshaw-Curtis quadrature rule.
     """
@@ -108,7 +107,7 @@ def scalar_product(a: List[Fraction], b: List[Fraction]) -> Fraction:
     return 2 * sum(c[i] / (i + 1) for i in range(0, lc, 2))
 
 
-def calc_bdef(ns: Tuple[int, int, int, int]) -> List[ndarray]:
+def calc_bdef(ns: Tuple[int, int, int, int]) -> List[np.ndarray]:
     """Calculate the decompositions of Newton polynomials (over the nodes
     of the n-point Clenshaw-Curtis quadrature rule) in terms of
     Legandre polynomials.
@@ -134,7 +133,7 @@ def calc_bdef(ns: Tuple[int, int, int, int]) -> List[ndarray]:
     return result
 
 
-def calc_V(x: ndarray, n: int) -> ndarray:
+def calc_V(x: np.ndarray, n: int) -> np.ndarray:
     V = [np.ones(x.shape), x.copy()]
     for i in range(2, n):
         V.append((2 * i - 1) / i * x * V[-1] - (i - 1) / i * V[-2])

adaptive/learner/integrator_learner.py

@@ -8,7 +8,7 @@
 from typing import Any, Callable, List, Optional, Set, Tuple, Union
 
 import numpy as np
-from numpy import bool_, float64, ndarray, ufunc
+from numpy import ufunc
 from scipy.linalg import norm
 from sortedcontainers import SortedSet
 
@@ -33,7 +33,7 @@
 )
 
 
-def _downdate(c: ndarray, nans: List[int], depth: int) -> ndarray:
+def _downdate(c: np.ndarray, nans: List[int], depth: int) -> np.ndarray:
     # This is algorithm 5 from the thesis of Pedro Gonnet.
     b = b_def[depth].copy()
     m = ns[depth] - 1
@@ -51,7 +51,7 @@ def _downdate(c: ndarray, nans: List[int], depth: int) -> ndarray:
     return c
 
 
-def _zero_nans(fx: ndarray) -> List[int]:
+def _zero_nans(fx: np.ndarray) -> List[int]:
     """Caution: this function modifies fx."""
     nans = []
     for i in range(len(fx)):
@@ -61,7 +61,7 @@ def _zero_nans(fx: ndarray) -> List[int]:
     return nans
 
 
-def _calc_coeffs(fx: ndarray, depth: int) -> ndarray:
+def _calc_coeffs(fx: np.ndarray, depth: int) -> np.ndarray:
     """Caution: this function modifies fx."""
     nans = _zero_nans(fx)
     c_new = V_inv[depth] @ fx
@@ -142,7 +142,11 @@ class _Interval:
     ]
 
     def __init__(
-        self, a: Union[int, float64], b: Union[int, float64], depth: int, rdepth: int
+        self,
+        a: Union[int, np.float64],
+        b: Union[int, np.float64],
+        depth: int,
+        rdepth: int,
     ) -> None:
         self.children = []
         self.data = {}
@@ -163,7 +167,7 @@ def make_first(cls, a: int, b: int, depth: int = 2) -> "_Interval":
         return ival
 
     @property
-    def T(self) -> ndarray:
+    def T(self) -> np.ndarray:
         """Get the correct shift matrix.
 
         Should only be called on children of a split interval.
@@ -180,7 +184,7 @@ def refinement_complete(self, depth: int) -> bool:
             return False
         return all(p in self.data for p in self.points(depth))
 
-    def points(self, depth: Optional[int] = None) -> ndarray:
+    def points(self, depth: Optional[int] = None) -> np.ndarray:
         if depth is None:
             depth = self.depth
         a = self.a
@@ -209,7 +213,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[float64, float]) -> None:
+    def update_heuristic_err(self, value: Union[np.float64, 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
@@ -222,7 +226,7 @@ def update_heuristic_err(self, value: Union[float64, float]) -> None:
                 continue
             child.update_heuristic_err(value / 2)
 
-    def calc_err(self, c_old: ndarray) -> float:
+    def calc_err(self, c_old: np.ndarray) -> float:
         c_new = self.c
         c_diff = np.zeros(max(len(c_old), len(c_new)))
         c_diff[: len(c_old)] = c_old
@@ -255,7 +259,7 @@ def update_ndiv_recursively(self) -> None:
 
     def complete_process(
         self, depth: int
-    ) -> Union[Tuple[bool, bool], Tuple[bool, bool_]]:
+    ) -> Union[Tuple[bool, bool], Tuple[bool, np.bool_]]:
         """Calculate the integral contribution and error from this interval,
         and update the done leaves of all ancestor intervals."""
         assert self.depth_complete is None or self.depth_complete == depth - 1
@@ -399,7 +403,7 @@ def __init__(
     def approximating_intervals(self) -> Set["_Interval"]:
         return self.first_ival.done_leaves
 
-    def tell(self, point: float64, value: float64) -> None:
+    def tell(self, point: np.float64, value: np.float64) -> None:
         if point not in self.x_mapping:
             raise ValueError(f"Point {point} doesn't belong to any interval")
         self.data[point] = value
@@ -458,7 +462,9 @@ def add_ival(self, ival: "_Interval") -> None:
 
     def ask(
         self, n: int, tell_pending: bool = True
-    ) -> Union[Tuple[List[float64], List[float64]], Tuple[List[float64], List[float]]]:
+    ) -> Union[
+        Tuple[List[np.float64], List[np.float64]], Tuple[List[np.float64], List[float]]
+    ]:
         """Choose points for learners."""
         if not tell_pending:
             with restore(self):
@@ -468,7 +474,9 @@ def ask(
 
     def _ask_and_tell_pending(
         self, n: int
-    ) -> Union[Tuple[List[float64], List[float64]], Tuple[List[float64], List[float]]]:
+    ) -> Union[
+        Tuple[List[np.float64], List[np.float64]], Tuple[List[np.float64], List[float]]
+    ]:
         points, loss_improvements = self.pop_from_stack(n)
         n_left = n - len(points)
         while n_left > 0:
@@ -487,9 +495,9 @@ def _ask_and_tell_pending(
     def pop_from_stack(
         self, n: int
     ) -> Union[
-        Tuple[List[float64], List[float64]],
+        Tuple[List[np.float64], List[np.float64]],
         Tuple[List[Any], List[Any]],
-        Tuple[List[float64], List[float]],
+        Tuple[List[np.float64], List[float]],
     ]:
         points = self._stack[:n]
         self._stack = self._stack[n:]
@@ -501,7 +509,7 @@ def pop_from_stack(
     def remove_unfinished(self):
         pass
 
-    def _fill_stack(self) -> List[float64]:
+    def _fill_stack(self) -> List[np.float64]:
         # 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)
@@ -542,11 +550,11 @@ def npoints(self) -> int:
         return len(self.data)
 
     @property
-    def igral(self) -> float64:
+    def igral(self) -> np.float64:
         return sum(i.igral for i in self.approximating_intervals)
 
     @property
-    def err(self) -> float64:
+    def err(self) -> np.float64:
         if self.approximating_intervals:
             err = sum(i.err for i in self.approximating_intervals)
             if err > sys.float_info.max: