tighten NDArray annotations

Author: amorison

src/stagpy/_helpers.py

@@ -6,6 +6,7 @@
 from inspect import getdoc
 
 import matplotlib.pyplot as plt
+import numpy as np
 
 if typing.TYPE_CHECKING:
     from matplotlib.figure import Figure
@@ -101,7 +102,9 @@ def baredoc(obj: object) -> str:
     return doc.rstrip(" .").lstrip()
 
 
-def find_in_sorted_arr(value: float, array: NDArray, after: bool = False) -> int:
+def find_in_sorted_arr(
+    value: float, array: NDArray[np.float64], after: bool = False
+) -> int:
     """Return position of element in a sorted array.
 
     Returns:

src/stagpy/datatypes.py

@@ -5,6 +5,8 @@
 import typing
 from dataclasses import dataclass
 
+import numpy as np
+
 if typing.TYPE_CHECKING:
     from numpy.typing import NDArray
 
@@ -23,7 +25,7 @@ class Varf:
 class Field:
     """Scalar field and associated metadata."""
 
-    values: NDArray
+    values: NDArray[np.float64]
     """values of field."""
     meta: Varf
     """metadata."""
@@ -45,9 +47,9 @@ class Varr:
 class Rprof:
     """Radial profile with associated radius and metadata."""
 
-    values: NDArray
+    values: NDArray[np.float64]
     """values of profile."""
-    rad: NDArray
+    rad: NDArray[np.float64]
     """radial position of profile."""
     meta: Varr
     """metadata."""
@@ -69,9 +71,9 @@ class Vart:
 class Tseries:
     """A time series with associated time and metadata."""
 
-    values: NDArray
+    values: NDArray[np.float64]
     """values of time series."""
-    time: NDArray
+    time: NDArray[np.float64]
     """time position of time series."""
     meta: Vart
     """metadata."""

src/stagpy/dimensions.py

@@ -4,18 +4,18 @@
 from dataclasses import dataclass
 from functools import cached_property
 
+import numpy as np
+
 from . import phyvars
 from .config import Scaling
 
 if typing.TYPE_CHECKING:
-    from typing import TypeVar
-
     from numpy.typing import NDArray
 
     from .parfile import StagyyPar
     from .stagyydata import StagyyData
 
-    T = TypeVar("T", float, NDArray)
+    T = float | NDArray[np.float64]
 
 
 @dataclass(frozen=True)

src/stagpy/field.py

@@ -37,7 +37,7 @@
 
 def _threed_extract(
     conf: Config, step: Step, var: str, walls: bool = False
-) -> tuple[tuple[NDArray, NDArray], NDArray]:
+) -> tuple[tuple[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]]:
     """Return suitable slices and coords for 3D fields."""
     is_vector = not valid_field_var(var)
     hwalls = is_vector or walls
@@ -91,7 +91,7 @@ def valid_field_var(var: str) -> bool:
 
 def get_meshes_fld(
     conf: Config, step: Step, var: str, walls: bool = False
-) -> tuple[NDArray, NDArray, NDArray, Varf]:
+) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], Varf]:
     """Return scalar field along with coordinates meshes.
 
     Only works properly in 2D geometry and 3D cartesian.
@@ -134,7 +134,9 @@ def get_meshes_fld(
 
 def get_meshes_vec(
     conf: Config, step: Step, var: str
-) -> tuple[NDArray, NDArray, NDArray, NDArray]:
+) -> tuple[
+    NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]
+]:
     """Return vector field components along with coordinates meshes.
 
     Only works properly in 2D geometry and 3D cartesian.
@@ -178,7 +180,7 @@ def get_meshes_vec(
 def plot_scalar(
     step: Step,
     var: str,
-    field: NDArray | None = None,
+    field: NDArray[np.float64] | None = None,
     axis: Axes | None = None,
     conf: Config | None = None,
     **extra: Any,
@@ -272,7 +274,7 @@ def plot_iso(
     axis: Axes,
     step: Step,
     var: str,
-    field: NDArray | None = None,
+    field: NDArray[np.float64] | None = None,
     conf: Config | None = None,
     **extra: Any,
 ) -> None:

src/stagpy/plates.py

@@ -27,7 +27,7 @@
     from .step import Geometry, Step
 
 
-def _vzcheck(iphis: Sequence[int], snap: Step, vz_thres: float) -> NDArray:
+def _vzcheck(iphis: Sequence[int], snap: Step, vz_thres: float) -> NDArray[np.int32]:
     """Remove positions where vz is below threshold."""
     # verifying vertical velocity
     vzabs = np.abs(snap.fields["v3"].values[0, ..., 0])
@@ -40,7 +40,9 @@ def _vzcheck(iphis: Sequence[int], snap: Step, vz_thres: float) -> NDArray:
 
 
 @lru_cache
-def detect_plates(snap: Step, vz_thres_ratio: float = 0) -> tuple[NDArray, NDArray]:
+def detect_plates(
+    snap: Step, vz_thres_ratio: float = 0
+) -> tuple[NDArray[np.int32], NDArray[np.int32]]:
     """Detect plate limits using derivative of horizontal velocity.
 
     This function is cached for convenience.
@@ -98,7 +100,9 @@ def detect_plates(snap: Step, vz_thres_ratio: float = 0) -> tuple[NDArray, NDArr
     return itrenches, iridges
 
 
-def _plot_plate_limits(axis: Axes, trenches: NDArray, ridges: NDArray) -> None:
+def _plot_plate_limits(
+    axis: Axes, trenches: NDArray[np.float64], ridges: NDArray[np.float64]
+) -> None:
     """Plot lines designating ridges and trenches."""
     for trench in trenches:
         axis.axvline(x=trench, color="red", ls="dashed", alpha=0.4)
@@ -183,7 +187,7 @@ def _surf_diag(snap: Step, name: str) -> Field:
     raise error.UnknownVarError(name)
 
 
-def _continents_location(snap: Step, at_surface: bool = True) -> NDArray:
+def _continents_location(snap: Step, at_surface: bool = True) -> NDArray[np.bool]:
     """Location of continents as a boolean array.
 
     If at_surface is True, it is evaluated only at the surface, otherwise it is

src/stagpy/processing.py

@@ -119,7 +119,9 @@ def delta_r(step: Step) -> Rprof:
     return Rprof((edges[1:] - edges[:-1]), step.rprofs.centers, meta)
 
 
-def _scale_prof(step: Step, rprof: NDArray, rad: NDArray | None = None) -> NDArray:
+def _scale_prof(
+    step: Step, rprof: NDArray[np.float64], rad: NDArray[np.float64] | None = None
+) -> NDArray[np.float64]:
     """Scale profile to take sphericity into account."""
     rbot, rtop = step.rprofs.bounds
     if rbot == 0:  # not spherical

src/stagpy/stagyydata.py

@@ -212,12 +212,12 @@ def tslice(
         )
 
     @property
-    def time(self) -> NDArray:
+    def time(self) -> NDArray[np.float64]:
         """Time vector."""
         return self._tseries["t"].to_numpy()
 
     @property
-    def isteps(self) -> NDArray:
+    def isteps(self) -> NDArray[np.float64]:
         """Step indices.
 
         This is such that `time[istep]` is at step `isteps[istep]`.
@@ -282,11 +282,11 @@ def _first_rprofs(self) -> step.RprofsInstant:
         return first_step.rprofs
 
     @property
-    def centers(self) -> NDArray:
+    def centers(self) -> NDArray[np.float64]:
         return self._first_rprofs.centers
 
     @property
-    def walls(self) -> NDArray:
+    def walls(self) -> NDArray[np.float64]:
         return self._first_rprofs.walls
 
     @property

src/stagpy/stagyyparsers.py

@@ -477,7 +477,7 @@ def field_header(fieldfile: Path) -> dict[str, Any] | None:
     return hdr.header
 
 
-def fields(fieldfile: Path) -> tuple[dict[str, Any], NDArray] | None:
+def fields(fieldfile: Path) -> tuple[dict[str, Any], NDArray[np.float64]] | None:
     """Extract fields data.
 
     Args:
@@ -555,7 +555,7 @@ def fields(fieldfile: Path) -> tuple[dict[str, Any], NDArray] | None:
     return header, flds
 
 
-def tracers(tracersfile: Path) -> dict[str, list[NDArray]] | None:
+def tracers(tracersfile: Path) -> dict[str, list[NDArray[np.float64]]] | None:
     """Extract tracers data.
 
     Args:
@@ -566,7 +566,7 @@ def tracers(tracersfile: Path) -> dict[str, list[NDArray]] | None:
     """
     if not tracersfile.is_file():
         return None
-    tra: dict[str, list[NDArray]] = {}
+    tra: dict[str, list[NDArray[np.float64]]] = {}
     with tracersfile.open("rb") as fid:
         readbin = partial(_readbin, fid)
         magic = readbin()
@@ -603,7 +603,7 @@ def tracers(tracersfile: Path) -> dict[str, list[NDArray]] | None:
     return tra
 
 
-def _read_group_h5(filename: Path, groupname: str) -> NDArray:
+def _read_group_h5(filename: Path, groupname: str) -> NDArray[np.float64]:
     """Return group content.
 
     Args:
@@ -623,7 +623,7 @@ def _read_group_h5(filename: Path, groupname: str) -> NDArray:
     return data  # need to be reshaped
 
 
-def _make_3d(field: NDArray, twod: str | None) -> NDArray:
+def _make_3d(field: NDArray[np.float64], twod: str | None) -> NDArray[np.float64]:
     """Add a dimension to field if necessary.
 
     Args:
@@ -641,7 +641,9 @@ def _make_3d(field: NDArray, twod: str | None) -> NDArray:
     return field.reshape(shp)
 
 
-def _ncores(meshes: list[dict[str, NDArray]], twod: str | None) -> NDArray:
+def _ncores(
+    meshes: list[dict[str, NDArray[np.float64]]], twod: str | None
+) -> NDArray[np.float64]:
     """Compute number of nodes in each direction."""
     nnpb = len(meshes)  # number of nodes per block
     nns = [1, 1, 1]  # number of nodes in x, y, z directions
@@ -672,8 +674,8 @@ def _ncores(meshes: list[dict[str, NDArray]], twod: str | None) -> NDArray:
 
 
 def _conglomerate_meshes(
-    meshin: list[dict[str, NDArray]], header: dict[str, Any]
-) -> dict[str, NDArray]:
+    meshin: list[dict[str, NDArray[np.float64]]], header: dict[str, Any]
+) -> dict[str, NDArray[np.float64]]:
     """Conglomerate meshes from several cores into one."""
     meshout = {}
     npc = header["nts"] // header["ncs"]
@@ -870,7 +872,7 @@ def read_geom_h5(xdmf: FieldXmf, snapshot: int) -> dict[str, Any]:
     header["mo_thick_sol"] = entry.mo_thick_sol
     header["ntb"] = 2 if entry.yin_yang else 1
 
-    all_meshes: list[dict[str, NDArray]] = []
+    all_meshes: list[dict[str, NDArray[np.float64]]] = []
     for h5file in entry.coord_files_yin(xdmf.path.parent):
         all_meshes.append({})
         with h5py.File(h5file, "r") as h5f:
@@ -927,7 +929,9 @@ def read_geom_h5(xdmf: FieldXmf, snapshot: int) -> dict[str, Any]:
     return header
 
 
-def _to_spherical(flds: NDArray, header: dict[str, Any]) -> NDArray:
+def _to_spherical(
+    flds: NDArray[np.float64], header: dict[str, Any]
+) -> NDArray[np.float64]:
     """Convert vector field to spherical."""
     cth = np.cos(header["t_mesh"][:, :, :-1])
     sth = np.sin(header["t_mesh"][:, :, :-1])
@@ -962,7 +966,9 @@ def _flds_shape(fieldname: str, header: dict[str, Any]) -> list[int]:
     return shp
 
 
-def _post_read_flds(flds: NDArray, header: dict[str, Any]) -> NDArray:
+def _post_read_flds(
+    flds: NDArray[np.float64], header: dict[str, Any]
+) -> NDArray[np.float64]:
     """Process flds to handle sphericity."""
     if flds.shape[0] >= 3 and header["rcmb"] > 0:
         # spherical vector
@@ -982,7 +988,7 @@ def read_field_h5(
     fieldname: str,
     snapshot: int,
     header: dict[str, Any] | None = None,
-) -> tuple[dict[str, Any], NDArray] | None:
+) -> tuple[dict[str, Any], NDArray[np.float64]] | None:
     """Extract field data from hdf5 files.
 
     Args:
@@ -1146,7 +1152,9 @@ def __getitem__(self, isnap: int) -> XmfTracersEntry:
             raise ParsingError(self.path, f"no data for snapshot {isnap}")
 
 
-def read_tracers_h5(xdmf: TracersXmf, infoname: str, snapshot: int) -> list[NDArray]:
+def read_tracers_h5(
+    xdmf: TracersXmf, infoname: str, snapshot: int
+) -> list[NDArray[np.float64]]:
     """Extract tracers data from hdf5 files.
 
     Args:
@@ -1157,11 +1165,11 @@ def read_tracers_h5(xdmf: TracersXmf, infoname: str, snapshot: int) -> list[NDAr
     Returns:
         Tracers data organized by attribute and block.
     """
-    tra: list[list[NDArray]] = [[], []]  # [block][core]
+    tra: list[list[NDArray[np.float64]]] = [[], []]  # [block][core]
     for tsub in xdmf[snapshot].tra_subdomains(xdmf.path.parent, infoname):
         tra[tsub.iblock].append(_read_group_h5(tsub.file, tsub.dataset))
 
-    tra_concat: list[NDArray] = []
+    tra_concat: list[NDArray[np.float64]] = []
     for trab in tra:
         if trab:
             tra_concat.append(np.concatenate(trab))