From 5a7a2a982a7a7308029ac1aee96d7a36f29d50d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Tue, 10 Dec 2024 16:42:05 +0100 Subject: [PATCH 1/9] TYP: fix invalid np.ndarray type annotations --- src/gpgi/_data_types.py | 10 +++--- src/gpgi/_typing.py | 74 ++++++++++++++++++++--------------------- 2 files changed, 41 insertions(+), 43 deletions(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index 37bb15f..cfaf213 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -380,10 +380,10 @@ def _validate(self) -> None: f"- from particles: {self.particles.dtype}\n" ) - def _get_padded_cell_edges(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + def _get_padded_cell_edges(self) -> tuple[RealArray, RealArray, RealArray]: edges = iter(self.grid.cell_edges.values()) - def pad(a: np.ndarray) -> np.ndarray: + def pad(a: RealArray) -> RealArray: dx = a[1] - a[0] return np.concatenate([[a[0] - dx], a, [a[-1] + dx]]) @@ -399,7 +399,7 @@ def pad(a: np.ndarray) -> np.ndarray: return cell_edges_x1, cell_edges_x2, cell_edges_x3 - def _get_3D_particle_coordinates(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + def _get_3D_particle_coordinates(self) -> tuple[RealArray, RealArray, RealArray]: particle_coords = iter(self.particles.coordinates.values()) particles_x1 = next(particle_coords) DTYPE = particles_x1.dtype @@ -477,9 +477,7 @@ def _validate_sort_axes(self, axes: tuple[int, ...]) -> None: if any(axis > self.grid.ndim - 1 for axis in axes): raise ValueError(f"Expected all axes to be <{self.grid.ndim}, got {axes!r}") - def _get_sort_key( - self, axes: tuple[int, ...] - ) -> np.ndarray[Any, np.dtype[np.uint16]]: + def _get_sort_key(self, axes: tuple[int, ...]) -> NDArray[np.uint16]: self._validate_sort_axes(axes) hci = self.host_cell_index diff --git a/src/gpgi/_typing.py b/src/gpgi/_typing.py index 5cfdc94..2d0145a 100644 --- a/src/gpgi/_typing.py +++ b/src/gpgi/_typing.py @@ -1,61 +1,61 @@ -from typing import NotRequired, TypedDict, TypeVar +from typing import Generic, NotRequired, TypedDict, TypeVar import numpy as np -import numpy.typing as npt +from numpy.typing import NDArray -Real = TypeVar("Real", np.float32, np.float64) -RealArray = npt.NDArray[Real] -HCIArray = npt.NDArray[np.uint16] +RealT = TypeVar("RealT", np.float32, np.float64) +RealArray = NDArray[RealT] +HCIArray = NDArray[np.uint16] Name = str -FieldMap = dict[str, np.ndarray] +FieldMap = dict[str, NDArray[RealT]] -class CartesianCoordinates(TypedDict): - x: np.ndarray - y: NotRequired[np.ndarray] - z: NotRequired[np.ndarray] +class CartesianCoordinates(TypedDict, Generic[RealT]): + x: NDArray[RealT] + y: NotRequired[NDArray[RealT]] + z: NotRequired[NDArray[RealT]] -class CylindricalCoordinates(TypedDict): - radius: np.ndarray - azimuth: NotRequired[np.ndarray] - z: NotRequired[np.ndarray] +class CylindricalCoordinates(TypedDict, Generic[RealT]): + radius: NDArray[RealT] + azimuth: NotRequired[NDArray[RealT]] + z: NotRequired[NDArray[RealT]] -class PolarCoordinates(TypedDict): - radius: np.ndarray - z: NotRequired[np.ndarray] - azimuth: NotRequired[np.ndarray] +class PolarCoordinates(TypedDict, Generic[RealT]): + radius: NDArray[RealT] + z: NotRequired[NDArray[RealT]] + azimuth: NotRequired[NDArray[RealT]] -class SphericalCoordinates(TypedDict): - colatitude: np.ndarray - radius: NotRequired[np.ndarray] - azimuth: NotRequired[np.ndarray] +class SphericalCoordinates(TypedDict, Generic[RealT]): + colatitude: NDArray[RealT] + radius: NotRequired[NDArray[RealT]] + azimuth: NotRequired[NDArray[RealT]] -class EquatorialCoordinates(TypedDict): - radius: np.ndarray - latitude: NotRequired[np.ndarray] - azimuth: NotRequired[np.ndarray] +class EquatorialCoordinates(TypedDict, Generic[RealT]): + radius: NDArray[RealT] + latitude: NotRequired[NDArray[RealT]] + azimuth: NotRequired[NDArray[RealT]] CoordMap = ( - CartesianCoordinates - | CylindricalCoordinates - | PolarCoordinates - | SphericalCoordinates - | EquatorialCoordinates + CartesianCoordinates[RealT] + | CylindricalCoordinates[RealT] + | PolarCoordinates[RealT] + | SphericalCoordinates[RealT] + | EquatorialCoordinates[RealT] ) -class GridDict(TypedDict): - cell_edges: CoordMap - fields: NotRequired[FieldMap] +class GridDict(TypedDict, Generic[RealT]): + cell_edges: CoordMap[RealT] + fields: NotRequired[FieldMap[RealT]] -class ParticleSetDict(TypedDict): - coordinates: CoordMap - fields: NotRequired[FieldMap] +class ParticleSetDict(TypedDict, Generic[RealT]): + coordinates: CoordMap[RealT] + fields: NotRequired[FieldMap[RealT]] From 8564d0361506e16b8327e2ed1a921b167816721b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Tue, 10 Dec 2024 17:26:38 +0100 Subject: [PATCH 2/9] fixup ? --- src/gpgi/_data_types.py | 27 ++++++++++++++++----------- src/gpgi/_spatial_data.py | 8 ++++---- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index cfaf213..7a78283 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -45,9 +45,13 @@ from _thread import LockType if TYPE_CHECKING: - from typing import Any, Self + from typing import Any, Self, TypeVar - from gpgi._typing import FieldMap, HCIArray, Name, RealArray + from numpy.typing import NDArray + + from gpgi._typing import FieldMap, HCIArray, Name, RealArray, RealT + + _FloatingT = TypeVar("_FloatingT", bound=np.floating) BoundarySpec = tuple[tuple[str, str, str], ...] @@ -116,8 +120,8 @@ def __init__( self, *, geometry: Geometry, - cell_edges: FieldMap, - fields: FieldMap | None = None, + cell_edges: FieldMap[RealT], + fields: FieldMap[RealT] | None = None, ) -> None: r""" Define a Grid from cell left-edges and data fields. @@ -144,6 +148,7 @@ def __init__( self.dtype = self.coordinates[self.axes[0]].dtype self._dx = np.full((3,), -1, dtype=self.coordinates[self.axes[0]].dtype) + for i, ax in enumerate(self.axes): if self.size == 1 or np.diff(self.coordinates[ax]).std() < 1e-16: # got a constant step in this direction, store it @@ -172,17 +177,17 @@ def __repr__(self) -> str: ) @property - def cell_edges(self) -> FieldMap: + def cell_edges(self) -> FieldMap[RealT]: r"""An alias for self.coordinates.""" return self.coordinates @cached_property - def cell_centers(self) -> FieldMap: + def cell_centers(self) -> FieldMap[RealT]: r"""The positions of cell centers in each direction.""" return {ax: 0.5 * (arr[1:] + arr[:-1]) for ax, arr in self.coordinates.items()} @cached_property - def cell_widths(self) -> FieldMap: + def cell_widths(self) -> FieldMap[RealT]: r"""The width of cells, expressed as the difference between consecutive left edges.""" return {ax: np.diff(arr) for ax, arr in self.coordinates.items()} @@ -206,7 +211,7 @@ def ndim(self) -> int: return len(self.axes) @property - def cell_volumes(self) -> RealArray: + def cell_volumes(self) -> RealArray[RealT]: r""" The generalized ND-volume of grid cells. @@ -241,8 +246,8 @@ def __init__( self, *, geometry: Geometry, - coordinates: FieldMap, - fields: FieldMap | None = None, + coordinates: FieldMap[RealT], + fields: FieldMap[RealT] | None = None, ) -> None: r""" Define a ParticleSet from point positions and data fields. @@ -261,7 +266,7 @@ def __init__( if fields is None: fields = {} - self.fields: FieldMap = fields + self.fields = fields self.axes = tuple(self.coordinates.keys()) self._validate() diff --git a/src/gpgi/_spatial_data.py b/src/gpgi/_spatial_data.py index 60be47e..44160c7 100644 --- a/src/gpgi/_spatial_data.py +++ b/src/gpgi/_spatial_data.py @@ -7,7 +7,7 @@ import numpy as np from numpy.typing import NDArray -from gpgi._typing import FieldMap, Name, Real +from gpgi._typing import FieldMap, Name, RealT class Geometry(StrEnum): @@ -154,7 +154,7 @@ def check( @staticmethod def _validate_shape_equality( name: str, - data: NDArray[Real], + data: NDArray[RealT], ref_arr: NamedArray | None, ) -> NamedArray: if ref_arr is not None and data.shape != ref_arr.data.shape: @@ -165,7 +165,7 @@ def _validate_shape_equality( return ref_arr or NamedArray(name, data) @staticmethod - def _validate_sorted_state(name: str, data: NDArray[Real]) -> None: + def _validate_sorted_state(name: str, data: NDArray[RealT]) -> None: a = data[0] for i, b in enumerate(data[1:], start=1): if a > b: @@ -178,7 +178,7 @@ def _validate_sorted_state(name: str, data: NDArray[Real]) -> None: @staticmethod def _validate_required_attributes( name: str, - data: NDArray[Real], + data: NDArray[RealT], required_attrs: dict[str, Any], ) -> None: for attr, expected in required_attrs.items(): From aa2483da785d1e850b0a1efdeab214c6416f971b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 17:58:31 +0100 Subject: [PATCH 3/9] fixup 2 --- src/gpgi/_data_types.py | 42 ++++++++++++------------ src/gpgi/_spatial_data.py | 8 ++--- src/gpgi/_typing.py | 68 +++++++++++++++++++-------------------- 3 files changed, 60 insertions(+), 58 deletions(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index 7a78283..8a080d3 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -11,7 +11,7 @@ from textwrap import indent from threading import Lock from time import monotonic_ns -from typing import TYPE_CHECKING, Literal, cast, final +from typing import TYPE_CHECKING, Generic, Literal, cast, final import numpy as np @@ -36,7 +36,7 @@ GeometryValidator, Validator, ) -from gpgi._typing import FieldMap, Name +from gpgi._typing import FieldMap, FloatT, Name from gpgi.typing import DepositionMethodT, DepositionMethodWithMetadataT if sys.version_info >= (3, 13): @@ -47,9 +47,10 @@ if TYPE_CHECKING: from typing import Any, Self, TypeVar + from numpy import dtype from numpy.typing import NDArray - from gpgi._typing import FieldMap, HCIArray, Name, RealArray, RealT + from gpgi._typing import FieldMap, HCIArray, Name, RealArray _FloatingT = TypeVar("_FloatingT", bound=np.floating) @@ -115,13 +116,13 @@ def check(cls, data: Grid) -> None: @final -class Grid: +class Grid(Generic[FloatT]): def __init__( self, *, geometry: Geometry, - cell_edges: FieldMap[RealT], - fields: FieldMap[RealT] | None = None, + cell_edges: FieldMap[FloatT], + fields: FieldMap[FloatT] | None = None, ) -> None: r""" Define a Grid from cell left-edges and data fields. @@ -137,7 +138,7 @@ def __init__( fields (keyword-only, optional): gpgi.typing.FieldMap """ self.geometry = geometry - self.coordinates = cell_edges + self.coordinates: FieldMap[FloatT] = cell_edges if fields is None: fields = {} @@ -145,10 +146,11 @@ def __init__( self.axes = tuple(self.coordinates.keys()) self._validate() - self.dtype = self.coordinates[self.axes[0]].dtype - - self._dx = np.full((3,), -1, dtype=self.coordinates[self.axes[0]].dtype) + self.dtype: dtype[FloatT] = self.coordinates[self.axes[0]].dtype + self._dx: NDArray[FloatT] = np.full( + (3,), -1, dtype=self.coordinates[self.axes[0]].dtype + ) for i, ax in enumerate(self.axes): if self.size == 1 or np.diff(self.coordinates[ax]).std() < 1e-16: # got a constant step in this direction, store it @@ -177,17 +179,17 @@ def __repr__(self) -> str: ) @property - def cell_edges(self) -> FieldMap[RealT]: + def cell_edges(self) -> FieldMap[FloatT]: r"""An alias for self.coordinates.""" return self.coordinates @cached_property - def cell_centers(self) -> FieldMap[RealT]: + def cell_centers(self) -> FieldMap[FloatT]: r"""The positions of cell centers in each direction.""" return {ax: 0.5 * (arr[1:] + arr[:-1]) for ax, arr in self.coordinates.items()} @cached_property - def cell_widths(self) -> FieldMap[RealT]: + def cell_widths(self) -> FieldMap[FloatT]: r"""The width of cells, expressed as the difference between consecutive left edges.""" return {ax: np.diff(arr) for ax, arr in self.coordinates.items()} @@ -211,7 +213,7 @@ def ndim(self) -> int: return len(self.axes) @property - def cell_volumes(self) -> RealArray[RealT]: + def cell_volumes(self) -> RealArray[FloatT]: r""" The generalized ND-volume of grid cells. @@ -241,13 +243,13 @@ def check(cls, data: ParticleSet) -> None: @final -class ParticleSet: +class ParticleSet(Generic[FloatT]): def __init__( self, *, geometry: Geometry, - coordinates: FieldMap[RealT], - fields: FieldMap[RealT] | None = None, + coordinates: FieldMap[FloatT], + fields: FieldMap[FloatT] | None = None, ) -> None: r""" Define a ParticleSet from point positions and data fields. @@ -262,15 +264,15 @@ def __init__( fields (keyword-only, optional): gpgi.typing.FieldMap """ self.geometry = geometry - self.coordinates = coordinates + self.coordinates: FieldMap[FloatT] = coordinates if fields is None: fields = {} - self.fields = fields + self.fields: FieldMap[FloatT] = fields self.axes = tuple(self.coordinates.keys()) self._validate() - self.dtype = self.coordinates[self.axes[0]].dtype + self.dtype: dtype[FloatT] = self.coordinates[self.axes[0]].dtype _validators: list[type[Validator[ParticleSet]]] = [ GeometryValidator, diff --git a/src/gpgi/_spatial_data.py b/src/gpgi/_spatial_data.py index 44160c7..0cc8d53 100644 --- a/src/gpgi/_spatial_data.py +++ b/src/gpgi/_spatial_data.py @@ -7,7 +7,7 @@ import numpy as np from numpy.typing import NDArray -from gpgi._typing import FieldMap, Name, RealT +from gpgi._typing import FieldMap, FloatT, Name class Geometry(StrEnum): @@ -154,7 +154,7 @@ def check( @staticmethod def _validate_shape_equality( name: str, - data: NDArray[RealT], + data: NDArray[FloatT], ref_arr: NamedArray | None, ) -> NamedArray: if ref_arr is not None and data.shape != ref_arr.data.shape: @@ -165,7 +165,7 @@ def _validate_shape_equality( return ref_arr or NamedArray(name, data) @staticmethod - def _validate_sorted_state(name: str, data: NDArray[RealT]) -> None: + def _validate_sorted_state(name: str, data: NDArray[FloatT]) -> None: a = data[0] for i, b in enumerate(data[1:], start=1): if a > b: @@ -178,7 +178,7 @@ def _validate_sorted_state(name: str, data: NDArray[RealT]) -> None: @staticmethod def _validate_required_attributes( name: str, - data: NDArray[RealT], + data: NDArray[FloatT], required_attrs: dict[str, Any], ) -> None: for attr, expected in required_attrs.items(): diff --git a/src/gpgi/_typing.py b/src/gpgi/_typing.py index 2d0145a..1a16589 100644 --- a/src/gpgi/_typing.py +++ b/src/gpgi/_typing.py @@ -3,59 +3,59 @@ import numpy as np from numpy.typing import NDArray -RealT = TypeVar("RealT", np.float32, np.float64) -RealArray = NDArray[RealT] +FloatT = TypeVar("FloatT", np.float32, np.float64) +RealArray = NDArray[FloatT] HCIArray = NDArray[np.uint16] Name = str -FieldMap = dict[str, NDArray[RealT]] +FieldMap = dict[str, NDArray[FloatT]] -class CartesianCoordinates(TypedDict, Generic[RealT]): - x: NDArray[RealT] - y: NotRequired[NDArray[RealT]] - z: NotRequired[NDArray[RealT]] +class CartesianCoordinates(TypedDict, Generic[FloatT]): + x: NDArray[FloatT] + y: NotRequired[NDArray[FloatT]] + z: NotRequired[NDArray[FloatT]] -class CylindricalCoordinates(TypedDict, Generic[RealT]): - radius: NDArray[RealT] - azimuth: NotRequired[NDArray[RealT]] - z: NotRequired[NDArray[RealT]] +class CylindricalCoordinates(TypedDict, Generic[FloatT]): + radius: NDArray[FloatT] + azimuth: NotRequired[NDArray[FloatT]] + z: NotRequired[NDArray[FloatT]] -class PolarCoordinates(TypedDict, Generic[RealT]): - radius: NDArray[RealT] - z: NotRequired[NDArray[RealT]] - azimuth: NotRequired[NDArray[RealT]] +class PolarCoordinates(TypedDict, Generic[FloatT]): + radius: NDArray[FloatT] + z: NotRequired[NDArray[FloatT]] + azimuth: NotRequired[NDArray[FloatT]] -class SphericalCoordinates(TypedDict, Generic[RealT]): - colatitude: NDArray[RealT] - radius: NotRequired[NDArray[RealT]] - azimuth: NotRequired[NDArray[RealT]] +class SphericalCoordinates(TypedDict, Generic[FloatT]): + colatitude: NDArray[FloatT] + radius: NotRequired[NDArray[FloatT]] + azimuth: NotRequired[NDArray[FloatT]] -class EquatorialCoordinates(TypedDict, Generic[RealT]): - radius: NDArray[RealT] - latitude: NotRequired[NDArray[RealT]] - azimuth: NotRequired[NDArray[RealT]] +class EquatorialCoordinates(TypedDict, Generic[FloatT]): + radius: NDArray[FloatT] + latitude: NotRequired[NDArray[FloatT]] + azimuth: NotRequired[NDArray[FloatT]] CoordMap = ( - CartesianCoordinates[RealT] - | CylindricalCoordinates[RealT] - | PolarCoordinates[RealT] - | SphericalCoordinates[RealT] - | EquatorialCoordinates[RealT] + CartesianCoordinates[FloatT] + | CylindricalCoordinates[FloatT] + | PolarCoordinates[FloatT] + | SphericalCoordinates[FloatT] + | EquatorialCoordinates[FloatT] ) -class GridDict(TypedDict, Generic[RealT]): - cell_edges: CoordMap[RealT] - fields: NotRequired[FieldMap[RealT]] +class GridDict(TypedDict, Generic[FloatT]): + cell_edges: CoordMap[FloatT] + fields: NotRequired[FieldMap[FloatT]] -class ParticleSetDict(TypedDict, Generic[RealT]): - coordinates: CoordMap[RealT] - fields: NotRequired[FieldMap[RealT]] +class ParticleSetDict(TypedDict, Generic[FloatT]): + coordinates: CoordMap[FloatT] + fields: NotRequired[FieldMap[FloatT]] From 95b81dba87d17a4ae80d87873ce4d445c336d65b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 17:59:31 +0100 Subject: [PATCH 4/9] fixup! fixup 2 --- src/gpgi/_data_types.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index 8a080d3..2c7cace 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -47,7 +47,6 @@ if TYPE_CHECKING: from typing import Any, Self, TypeVar - from numpy import dtype from numpy.typing import NDArray from gpgi._typing import FieldMap, HCIArray, Name, RealArray @@ -146,7 +145,7 @@ def __init__( self.axes = tuple(self.coordinates.keys()) self._validate() - self.dtype: dtype[FloatT] = self.coordinates[self.axes[0]].dtype + self.dtype: np.dtype[FloatT] = self.coordinates[self.axes[0]].dtype self._dx: NDArray[FloatT] = np.full( (3,), -1, dtype=self.coordinates[self.axes[0]].dtype @@ -272,7 +271,7 @@ def __init__( self.axes = tuple(self.coordinates.keys()) self._validate() - self.dtype: dtype[FloatT] = self.coordinates[self.axes[0]].dtype + self.dtype: np.dtype[FloatT] = self.coordinates[self.axes[0]].dtype _validators: list[type[Validator[ParticleSet]]] = [ GeometryValidator, From f41e091d6fa57acdef65ae65e97df91e09964a0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 18:07:26 +0100 Subject: [PATCH 5/9] fixup 3 --- src/gpgi/_data_types.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index 2c7cace..d5bf5cf 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -49,7 +49,7 @@ from numpy.typing import NDArray - from gpgi._typing import FieldMap, HCIArray, Name, RealArray + from gpgi._typing import FieldMap, HCIArray, Name _FloatingT = TypeVar("_FloatingT", bound=np.floating) @@ -212,7 +212,7 @@ def ndim(self) -> int: return len(self.axes) @property - def cell_volumes(self) -> RealArray[FloatT]: + def cell_volumes(self) -> NDArray[FloatT]: r""" The generalized ND-volume of grid cells. @@ -223,7 +223,7 @@ def cell_volumes(self) -> RealArray[FloatT]: widths = list(self.cell_widths.values()) if self.geometry is Geometry.CARTESIAN: raw = np.prod(np.meshgrid(*widths), axis=0) - return cast("RealArray", np.swapaxes(raw, 0, 1)) + return np.swapaxes(raw, 0, 1) else: raise NotImplementedError( f"cell_volumes property is not implemented for {self.geometry} geometry" @@ -386,10 +386,12 @@ def _validate(self) -> None: f"- from particles: {self.particles.dtype}\n" ) - def _get_padded_cell_edges(self) -> tuple[RealArray, RealArray, RealArray]: + def _get_padded_cell_edges( + self, + ) -> tuple[NDArray[FloatT], NDArray[FloatT], NDArray[FloatT]]: edges = iter(self.grid.cell_edges.values()) - def pad(a: RealArray) -> RealArray: + def pad(a: NDArray[FloatT]) -> NDArray[FloatT]: dx = a[1] - a[0] return np.concatenate([[a[0] - dx], a, [a[-1] + dx]]) @@ -397,7 +399,12 @@ def pad(a: RealArray) -> RealArray: cell_edges_x1 = pad(x1) DTYPE = cell_edges_x1.dtype - cell_edges_x2 = cell_edges_x3 = np.empty(0, DTYPE) + cell_edges_x2: np.ndarray[tuple[int, ...], np.dtype[FloatT]] = np.empty( + 0, DTYPE + ) + cell_edges_x3: np.ndarray[tuple[int, ...], np.dtype[FloatT]] = np.empty( + 0, DTYPE + ) if self.grid.ndim >= 2: cell_edges_x2 = pad(next(edges)) if self.grid.ndim == 3: @@ -405,7 +412,9 @@ def pad(a: RealArray) -> RealArray: return cell_edges_x1, cell_edges_x2, cell_edges_x3 - def _get_3D_particle_coordinates(self) -> tuple[RealArray, RealArray, RealArray]: + def _get_3D_particle_coordinates( + self, + ) -> tuple[NDArray[FloatT], NDArray[FloatT], NDArray[FloatT]]: particle_coords = iter(self.particles.coordinates.values()) particles_x1 = next(particle_coords) DTYPE = particles_x1.dtype @@ -784,9 +793,9 @@ def _sanitize_boundaries(self, boundaries: dict[Name, tuple[Name, Name]]) -> Non def _apply_boundary_conditions( self, - array: RealArray, + array: NDArray[FloatT], boundaries: dict[Name, tuple[Name, Name]], - weight_array: RealArray | None, + weight_array: NDArray[FloatT] | None, ) -> None: axes = list(self.grid.axes) for ax, bv in boundaries.items(): From 892b8ceb0998fc1502fddfbbd46394b532275800 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 18:07:58 +0100 Subject: [PATCH 6/9] fixup! fixup 3 --- src/gpgi/_typing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/gpgi/_typing.py b/src/gpgi/_typing.py index 1a16589..6951327 100644 --- a/src/gpgi/_typing.py +++ b/src/gpgi/_typing.py @@ -4,7 +4,6 @@ from numpy.typing import NDArray FloatT = TypeVar("FloatT", np.float32, np.float64) -RealArray = NDArray[FloatT] HCIArray = NDArray[np.uint16] From 4a465f7d332ee1e3bea4f977fe33de9e3f8ee54e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 22:19:36 +0100 Subject: [PATCH 7/9] fixup! fixup! fixup 3 --- src/gpgi/_boundaries.py | 103 ++++++++++++++++++++-------------------- src/gpgi/typing.py | 40 ++++++++-------- 2 files changed, 73 insertions(+), 70 deletions(-) diff --git a/src/gpgi/_boundaries.py b/src/gpgi/_boundaries.py index 403765e..dcc93f8 100644 --- a/src/gpgi/_boundaries.py +++ b/src/gpgi/_boundaries.py @@ -2,25 +2,26 @@ from collections.abc import Callable from threading import Lock -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import Any, Literal, cast -if TYPE_CHECKING: - from gpgi._typing import RealArray +from numpy.typing import NDArray + +from gpgi._typing import FloatT BoundaryRecipeT = Callable[ [ - "RealArray", - "RealArray", - "RealArray", - "RealArray", - "RealArray", - "RealArray", - "RealArray", - "RealArray", + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], + NDArray[FloatT], Literal["left", "right"], dict[str, Any], ], - "RealArray", + NDArray[FloatT], ] @@ -126,64 +127,64 @@ def __contains__(self, key: str) -> bool: # basic recipes def open_boundary( - same_side_active_layer: RealArray, - same_side_ghost_layer: RealArray, - opposite_side_active_layer: RealArray, - opposite_side_ghost_layer: RealArray, - weight_same_side_active_layer: RealArray, - weight_same_side_ghost_layer: RealArray, - weight_opposite_side_active_layer: RealArray, - weight_opposite_side_ghost_layer: RealArray, + same_side_active_layer: NDArray[FloatT], + same_side_ghost_layer: NDArray[FloatT], + opposite_side_active_layer: NDArray[FloatT], + opposite_side_ghost_layer: NDArray[FloatT], + weight_same_side_active_layer: NDArray[FloatT], + weight_same_side_ghost_layer: NDArray[FloatT], + weight_opposite_side_active_layer: NDArray[FloatT], + weight_opposite_side_ghost_layer: NDArray[FloatT], side: Literal["left", "right"], metadata: dict[str, Any], -) -> RealArray: +) -> NDArray[FloatT]: # return the active layer unchanged return same_side_active_layer def wall_boundary( - same_side_active_layer: RealArray, - same_side_ghost_layer: RealArray, - opposite_side_active_layer: RealArray, - opposite_side_ghost_layer: RealArray, - weight_same_side_active_layer: RealArray, - weight_same_side_ghost_layer: RealArray, - weight_opposite_side_active_layer: RealArray, - weight_opposite_side_ghost_layer: RealArray, + same_side_active_layer: NDArray[FloatT], + same_side_ghost_layer: NDArray[FloatT], + opposite_side_active_layer: NDArray[FloatT], + opposite_side_ghost_layer: NDArray[FloatT], + weight_same_side_active_layer: NDArray[FloatT], + weight_same_side_ghost_layer: NDArray[FloatT], + weight_opposite_side_active_layer: NDArray[FloatT], + weight_opposite_side_ghost_layer: NDArray[FloatT], side: Literal["left", "right"], metadata: dict[str, Any], -) -> RealArray: - return cast("RealArray", same_side_active_layer + same_side_ghost_layer) +) -> NDArray[FloatT]: + return cast(NDArray[FloatT], same_side_active_layer + same_side_ghost_layer) def antisymmetric_boundary( - same_side_active_layer: RealArray, - same_side_ghost_layer: RealArray, - opposite_side_active_layer: RealArray, - opposite_side_ghost_layer: RealArray, - weight_same_side_active_layer: RealArray, - weight_same_side_ghost_layer: RealArray, - weight_opposite_side_active_layer: RealArray, - weight_opposite_side_ghost_layer: RealArray, + same_side_active_layer: NDArray[FloatT], + same_side_ghost_layer: NDArray[FloatT], + opposite_side_active_layer: NDArray[FloatT], + opposite_side_ghost_layer: NDArray[FloatT], + weight_same_side_active_layer: NDArray[FloatT], + weight_same_side_ghost_layer: NDArray[FloatT], + weight_opposite_side_active_layer: NDArray[FloatT], + weight_opposite_side_ghost_layer: NDArray[FloatT], side: Literal["left", "right"], metadata: dict[str, Any], -) -> RealArray: - return cast("RealArray", same_side_active_layer - same_side_ghost_layer) +) -> NDArray[FloatT]: + return cast(NDArray[FloatT], same_side_active_layer - same_side_ghost_layer) def periodic_boundary( - same_side_active_layer: RealArray, - same_side_ghost_layer: RealArray, - opposite_side_active_layer: RealArray, - opposite_side_ghost_layer: RealArray, - weight_same_side_active_layer: RealArray, - weight_same_side_ghost_layer: RealArray, - weight_opposite_side_active_layer: RealArray, - weight_opposite_side_ghost_layer: RealArray, + same_side_active_layer: NDArray[FloatT], + same_side_ghost_layer: NDArray[FloatT], + opposite_side_active_layer: NDArray[FloatT], + opposite_side_ghost_layer: NDArray[FloatT], + weight_same_side_active_layer: NDArray[FloatT], + weight_same_side_ghost_layer: NDArray[FloatT], + weight_opposite_side_active_layer: NDArray[FloatT], + weight_opposite_side_ghost_layer: NDArray[FloatT], side: Literal["left", "right"], metadata: dict[str, Any], -) -> RealArray: - return cast("RealArray", same_side_active_layer + opposite_side_ghost_layer) +) -> NDArray[FloatT]: + return cast(NDArray[FloatT], same_side_active_layer + opposite_side_ghost_layer) _base_registry: dict[str, BoundaryRecipeT] = { diff --git a/src/gpgi/typing.py b/src/gpgi/typing.py index 1b431db..222c220 100644 --- a/src/gpgi/typing.py +++ b/src/gpgi/typing.py @@ -8,7 +8,9 @@ from gpgi._typing import FieldMap if TYPE_CHECKING: - from gpgi._typing import HCIArray, RealArray + from numpy.typing import NDArray + + from gpgi._typing import FloatT, HCIArray __all__ = [ @@ -21,32 +23,32 @@ class DepositionMethodT(Protocol): def __call__( # noqa D102 self, - cell_edges_x1: RealArray, - cell_edges_x2: RealArray, - cell_edges_x3: RealArray, - particles_x1: RealArray, - particles_x2: RealArray, - particles_x3: RealArray, - field: RealArray, - weight_field: RealArray, + cell_edges_x1: NDArray[FloatT], + cell_edges_x2: NDArray[FloatT], + cell_edges_x3: NDArray[FloatT], + particles_x1: NDArray[FloatT], + particles_x2: NDArray[FloatT], + particles_x3: NDArray[FloatT], + field: NDArray[FloatT], + weight_field: NDArray[FloatT], hci: HCIArray, - out: RealArray, + out: NDArray[FloatT], ) -> None: ... class DepositionMethodWithMetadataT(Protocol): def __call__( # noqa D102 self, - cell_edges_x1: RealArray, - cell_edges_x2: RealArray, - cell_edges_x3: RealArray, - particles_x1: RealArray, - particles_x2: RealArray, - particles_x3: RealArray, - field: RealArray, - weight_field: RealArray, + cell_edges_x1: NDArray[FloatT], + cell_edges_x2: NDArray[FloatT], + cell_edges_x3: NDArray[FloatT], + particles_x1: NDArray[FloatT], + particles_x2: NDArray[FloatT], + particles_x3: NDArray[FloatT], + field: NDArray[FloatT], + weight_field: NDArray[FloatT], hci: HCIArray, - out: RealArray, + out: NDArray[FloatT], *, metadata: dict[str, Any], ) -> None: ... From 5363e38a833cd1a4f07ce0affb6eb0b4390b7f82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 22:19:53 +0100 Subject: [PATCH 8/9] fixup! fixup! fixup! fixup 3 --- src/gpgi/_boundaries.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/gpgi/_boundaries.py b/src/gpgi/_boundaries.py index dcc93f8..f20cadc 100644 --- a/src/gpgi/_boundaries.py +++ b/src/gpgi/_boundaries.py @@ -1,5 +1,3 @@ -from __future__ import annotations - from collections.abc import Callable from threading import Lock from typing import Any, Literal, cast From cb16119af7ec4dee5babbbd4a3063a681b8f67c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 25 Jan 2025 22:33:19 +0100 Subject: [PATCH 9/9] ignore spurious type check error --- src/gpgi/_data_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gpgi/_data_types.py b/src/gpgi/_data_types.py index d5bf5cf..24a8f22 100644 --- a/src/gpgi/_data_types.py +++ b/src/gpgi/_data_types.py @@ -185,7 +185,7 @@ def cell_edges(self) -> FieldMap[FloatT]: @cached_property def cell_centers(self) -> FieldMap[FloatT]: r"""The positions of cell centers in each direction.""" - return {ax: 0.5 * (arr[1:] + arr[:-1]) for ax, arr in self.coordinates.items()} + return {ax: 0.5 * (arr[1:] + arr[:-1]) for ax, arr in self.coordinates.items()} # type: ignore [misc] @cached_property def cell_widths(self) -> FieldMap[FloatT]: