Merge pull request #724 from zwicker-group/typing

Defined more precise array type

Author: david-zwicker

pde/fields/base.py

@@ -17,7 +17,7 @@
 
 from ..grids.base import GridBase
 from ..tools.plotting import napari_add_layers, napari_viewer
-from ..tools.typing import NumberOrArray
+from ..tools.typing import NumberOrArray, NumericArray
 
 if TYPE_CHECKING:
     from .scalar import ScalarField
@@ -38,15 +38,15 @@ class FieldBase(metaclass=ABCMeta):
 
     _subclasses: dict[str, type[FieldBase]] = {}  # all classes inheriting from this
     _grid: GridBase  # the grid on which the field is defined
-    __data_full: np.ndarray  # the data on the grid including ghost points
-    _data_valid: np.ndarray  # the valid data without ghost points
+    __data_full: NumericArray  # the data on the grid including ghost points
+    _data_valid: NumericArray  # the valid data without ghost points
     _label: str | None  # name of the field
     _logger: logging.Logger  # logger instance to output information
 
     def __init__(
         self,
         grid: GridBase,
-        data: np.ndarray,
+        data: NumericArray,
         *,
         label: str | None = None,
     ):
@@ -82,7 +82,7 @@ def __getstate__(self) -> dict[str, Any]:
         return state
 
     @property
-    def data(self) -> np.ndarray:
+    def data(self) -> NumericArray:
         """:class:`~numpy.ndarray`: discretized data at the support points."""
         return self._data_valid
 
@@ -109,7 +109,7 @@ def _idx_valid(self) -> tuple[slice, ...]:
         return idx_comp + self.grid._idx_valid
 
     @property
-    def _data_full(self) -> np.ndarray:
+    def _data_full(self) -> NumericArray:
         """:class:`~numpy.ndarray`: the full data including ghost cells."""
         return self.__data_full
 
@@ -146,17 +146,17 @@ def _data_full(self, value: NumberOrArray) -> None:
         self._data_valid = self.__data_full[self._idx_valid]
 
     @property
-    def _data_flat(self) -> np.ndarray:
+    def _data_flat(self) -> NumericArray:
         """:class:`~numpy.ndarray`: flat version of discretized data with ghost
         cells."""
         # flatten the first dimension of the internal data by creating a view and then
         # setting the new shape. This disallows accidental copying of the data
         data_flat = self._data_full.view()
         data_flat.shape = (-1, *self.grid._shape_full)
-        return data_flat
+        return data_flat  # type: ignore
 
     @_data_flat.setter
-    def _data_flat(self, value: np.ndarray) -> None:
+    def _data_flat(self, value: NumericArray) -> None:
         """Set the full data including ghost cells from a flattened array."""
         # simply set the data -> this might need to be overwritten
         self._data_full = value
@@ -187,7 +187,7 @@ def label(self, value: str | None = None):
 
     @classmethod
     def from_state(
-        cls, attributes: dict[str, Any], data: np.ndarray | None = None
+        cls, attributes: dict[str, Any], data: NumericArray | None = None
     ) -> FieldBase:
         """Create a field from given state.
 

pde/fields/collection.py

@@ -26,7 +26,7 @@
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import Number, number_array
 from ..tools.plotting import PlotReference, plot_on_axes, plot_on_figure
-from ..tools.typing import NumberOrArray
+from ..tools.typing import NumberOrArray, NumericArray
 from .base import FieldBase
 from .datafield_base import DataFieldBase
 from .scalar import ScalarField
@@ -103,7 +103,7 @@ def __init__(
             self._fields = fields  # type: ignore
 
         # extract data from individual fields
-        fields_data: list[np.ndarray] = []
+        fields_data: list[NumericArray] = []
         self._slices: list[slice] = []
         dof = 0  # count local degrees of freedom
         for field in self.fields:
@@ -253,7 +253,7 @@ def __eq__(self, other):
 
     @classmethod
     def from_state(
-        cls, attributes: dict[str, Any], data: np.ndarray | None = None
+        cls, attributes: dict[str, Any], data: NumericArray | None = None
     ) -> FieldCollection:
         """Create a field collection from given state.
 
@@ -286,7 +286,7 @@ def from_data(
         cls,
         field_classes,
         grid: GridBase,
-        data: np.ndarray,
+        data: NumericArray,
         *,
         with_ghost_cells: bool = True,
         label: str | None = None,
@@ -701,7 +701,7 @@ def averages(self) -> list:
         return [field.average for field in self]
 
     @property
-    def magnitudes(self) -> np.ndarray:
+    def magnitudes(self) -> NumericArray:
         """:class:`~numpy.ndarray`: scalar magnitudes of all fields."""
         return np.array([field.magnitude for field in self])  # type: ignore
 
@@ -807,7 +807,7 @@ def _get_merged_image_data(
         transpose: bool = False,
         vmin: float | list[float | None] | None = None,
         vmax: float | list[float | None] | None = None,
-    ) -> tuple[np.ndarray, dict[str, Any]]:
+    ) -> tuple[NumericArray, dict[str, Any]]:
         """Obtain data required for a merged plot.
 
         Args:

pde/fields/datafield_base.py

@@ -25,7 +25,7 @@
 from ..tools.numba import get_common_numba_dtype, jit, make_array_constructor
 from ..tools.plotting import PlotReference, plot_on_axes
 from ..tools.spectral import CorrelationType, make_correlated_noise
-from ..tools.typing import ArrayLike, NumberOrArray
+from ..tools.typing import ArrayLike, NumberOrArray, NumericArray
 from .base import FieldBase, RankError
 
 if TYPE_CHECKING:
@@ -182,7 +182,7 @@ def random_uniform(
             # create complex random numbers for the field
             real_part = rng.uniform(np.real(vmin), np.real(vmax), size=shape)
             imag_part = rng.uniform(np.imag(vmin), np.imag(vmax), size=shape)
-            data: np.ndarray = real_part + 1j * imag_part
+            data: NumericArray = real_part + 1j * imag_part
         else:
             # create real random numbers for the field
             data = rng.uniform(vmin, vmax, size=shape)
@@ -293,7 +293,7 @@ def random_normal(
         else:
             tensor_shape = (grid.dim,) * cls.rank
 
-            def make_random_field() -> np.ndarray:
+            def make_random_field() -> NumericArray:
                 """Helper function that creates a single tensor field."""
                 out = np.empty(tensor_shape + grid.shape)
                 print(out.shape, tensor_shape, grid)
@@ -313,7 +313,7 @@ def make_random_field() -> np.ndarray:
             # create complex random numbers for the field
             real_part = np.real(mean) + np.real(std) * scale * make_random_field()
             imag_part = np.imag(mean) + np.imag(std) * scale * make_random_field()
-            data: np.ndarray = real_part + 1j * imag_part
+            data: NumericArray = real_part + 1j * imag_part
         else:
             # create real random numbers for the field
             data = mean + std * scale * make_random_field()
@@ -481,7 +481,7 @@ def get_class_by_rank(cls, rank: int) -> type[DataFieldBase]:
     def from_state(
         cls: type[TDataField],
         attributes: dict[str, Any],
-        data: np.ndarray | None = None,
+        data: NumericArray | None = None,
     ) -> TDataField:
         """Create a field from given state.
 
@@ -593,7 +593,7 @@ def make_interpolator(
         *,
         fill: Number | None = None,
         with_ghost_cells: bool = False,
-    ) -> Callable[[np.ndarray, np.ndarray], NumberOrArray]:
+    ) -> Callable[[NumericArray, NumericArray], NumberOrArray]:
         r"""Returns a function that can be used to interpolate values.
 
         Args:
@@ -617,7 +617,7 @@ def make_interpolator(
         # convert `fill` to dtype of data
         if fill is not None:
             if self.rank == 0:
-                fill = self.data.dtype.type(fill)
+                fill = self.data.dtype.type(fill)  # type: ignore
             else:
                 fill = np.broadcast_to(fill, self.data_shape).astype(self.data.dtype)
 
@@ -636,8 +636,8 @@ def make_interpolator(
 
         @jit
         def interpolator(
-            point: np.ndarray, data: np.ndarray | None = None
-        ) -> np.ndarray:
+            point: NumericArray, data: NumericArray | None = None
+        ) -> NumericArray:
             """Return the interpolated value at the position `point`
 
             Args:
@@ -680,11 +680,11 @@ def interpolator(
     @fill_in_docstring
     def interpolate(
         self,
-        point: np.ndarray,
+        point: NumericArray,
         *,
         bc: BoundariesData | None = None,
         fill: Number | None = None,
-    ) -> np.ndarray:
+    ) -> NumericArray:
         r"""Interpolate the field to points between support points.
 
         Args:
@@ -748,7 +748,7 @@ def interpolate_to_grid(
         """
         raise NotImplementedError(f"Cannot interpolate {self.__class__.__name__}")
 
-    def insert(self, point: np.ndarray, amount: ArrayLike) -> None:
+    def insert(self, point: NumericArray, amount: ArrayLike) -> None:
         """Adds an (integrated) value to the field at an interpolated position.
 
         Args:
@@ -965,7 +965,7 @@ def apply_operator(
 
     def make_dot_operator(
         self, backend: Literal["numpy", "numba"] = "numba", *, conjugate: bool = True
-    ) -> Callable[[np.ndarray, np.ndarray, np.ndarray | None], np.ndarray]:
+    ) -> Callable[[NumericArray, NumericArray, NumericArray | None], NumericArray]:
         """Return operator calculating the dot product between two fields.
 
         This supports both products between two vectors as well as products
@@ -986,13 +986,13 @@ def make_dot_operator(
         num_axes = self.grid.num_axes
 
         @register_jitable
-        def maybe_conj(arr: np.ndarray) -> np.ndarray:
+        def maybe_conj(arr: NumericArray) -> NumericArray:
             """Helper function implementing optional conjugation."""
             return arr.conjugate() if conjugate else arr
 
         def dot(
-            a: np.ndarray, b: np.ndarray, out: np.ndarray | None = None
-        ) -> np.ndarray:
+            a: NumericArray, b: NumericArray, out: NumericArray | None = None
+        ) -> NumericArray:
             """Numpy implementation to calculate dot product between two fields."""
             rank_a = a.ndim - num_axes
             rank_b = b.ndim - num_axes
@@ -1040,8 +1040,8 @@ def get_rank(arr: nb.types.Type | nb.types.Optional) -> int:
 
             @overload(dot, inline="always")
             def dot_ol(
-                a: np.ndarray, b: np.ndarray, out: np.ndarray | None = None
-            ) -> np.ndarray:
+                a: NumericArray, b: NumericArray, out: NumericArray | None = None
+            ) -> NumericArray:
                 """Numba implementation to calculate dot product between two fields."""
                 # get (and check) rank of the input arrays
                 rank_a = get_rank(a)
@@ -1050,15 +1050,19 @@ def dot_ol(
                 if rank_a == 1 and rank_b == 1:  # result is scalar field
 
                     @register_jitable
-                    def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
+                    def calc(
+                        a: NumericArray, b: NumericArray, out: NumericArray
+                    ) -> None:
                         out[:] = a[0] * maybe_conj(b[0])
                         for j in range(1, dim):
                             out[:] += a[j] * maybe_conj(b[j])
 
                 elif rank_a == 2 and rank_b == 1:  # result is vector field
 
                     @register_jitable
-                    def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
+                    def calc(
+                        a: NumericArray, b: NumericArray, out: NumericArray
+                    ) -> None:
                         for i in range(dim):
                             out[i] = a[i, 0] * maybe_conj(b[0])
                             for j in range(1, dim):
@@ -1067,7 +1071,9 @@ def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
                 elif rank_a == 1 and rank_b == 2:  # result is vector field
 
                     @register_jitable
-                    def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
+                    def calc(
+                        a: NumericArray, b: NumericArray, out: NumericArray
+                    ) -> None:
                         for i in range(dim):
                             out[i] = a[0] * maybe_conj(b[0, i])
                             for j in range(1, dim):
@@ -1076,7 +1082,9 @@ def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
                 elif rank_a == 2 and rank_b == 2:  # result is tensor-2 field
 
                     @register_jitable
-                    def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
+                    def calc(
+                        a: NumericArray, b: NumericArray, out: NumericArray
+                    ) -> None:
                         for i in range(dim):
                             for j in range(dim):
                                 out[i, j] = a[i, 0] * maybe_conj(b[0, j])
@@ -1095,8 +1103,10 @@ def calc(a: np.ndarray, b: np.ndarray, out: np.ndarray) -> None:
                     dtype = get_common_numba_dtype(a, b)
 
                     def dot_impl(
-                        a: np.ndarray, b: np.ndarray, out: np.ndarray | None = None
-                    ) -> np.ndarray:
+                        a: NumericArray,
+                        b: NumericArray,
+                        out: NumericArray | None = None,
+                    ) -> NumericArray:
                         """Helper function allocating output array."""
                         assert a.shape == a_shape
                         assert b.shape == b_shape
@@ -1108,8 +1118,10 @@ def dot_impl(
                     # function is called with `out` argument -> reuse `out` array
 
                     def dot_impl(
-                        a: np.ndarray, b: np.ndarray, out: np.ndarray | None = None
-                    ) -> np.ndarray:
+                        a: NumericArray,
+                        b: NumericArray,
+                        out: NumericArray | None = None,
+                    ) -> NumericArray:
                         """Helper function without allocating output array."""
                         assert a.shape == a_shape
                         assert b.shape == b_shape
@@ -1121,8 +1133,8 @@ def dot_impl(
 
             @jit
             def dot_compiled(
-                a: np.ndarray, b: np.ndarray, out: np.ndarray | None = None
-            ) -> np.ndarray:
+                a: NumericArray, b: NumericArray, out: NumericArray | None = None
+            ) -> NumericArray:
                 """Numba implementation to calculate dot product between two fields."""
                 return dot(a, b, out)
 

pde/fields/scalar.py

@@ -18,7 +18,7 @@
 from ..grids.boundaries.axes import BoundariesData
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import Number
-from ..tools.typing import NumberOrArray
+from ..tools.typing import NumberOrArray, NumericArray
 from .datafield_base import DataFieldBase
 
 if TYPE_CHECKING:

pde/fields/tensorial.py

@@ -15,7 +15,7 @@
 from ..tools.docstrings import fill_in_docstring
 from ..tools.misc import get_common_dtype
 from ..tools.plotting import PlotReference, plot_on_figure
-from ..tools.typing import NumberOrArray
+from ..tools.typing import NumberOrArray, NumericArray
 from .datafield_base import DataFieldBase
 from .scalar import ScalarField
 from .vectorial import VectorField
@@ -102,7 +102,7 @@ def from_expression(
         points = [grid.cell_coords[..., i] for i in range(grid.num_axes)]
 
         # evaluate all vector components at all points
-        data: list[list[np.ndarray]] = [[None] * grid.dim for _ in range(grid.dim)]  # type: ignore
+        data: list[list[NumericArray]] = [[None] * grid.dim for _ in range(grid.dim)]  # type: ignore
         for i in range(grid.dim):
             for j in range(grid.dim):
                 expr = ScalarExpression(
@@ -245,7 +245,7 @@ def divergence(
         return self.apply_operator("tensor_divergence", bc=bc, out=out, **kwargs)  # type: ignore
 
     @property
-    def integral(self) -> np.ndarray:
+    def integral(self) -> NumericArray:
         """:class:`~numpy.ndarray`: integral of each component over space."""
         return self.grid.integrate(self.data)  # type: ignore