diff --git a/farms_core/analysis/plot.py b/farms_core/analysis/plot.py index 139f584..ad2ff55 100644 --- a/farms_core/analysis/plot.py +++ b/farms_core/analysis/plot.py @@ -5,7 +5,7 @@ from typing import List import numpy as np -from nptyping import NDArray, Shape, Float +from ..array.types import NDARRAY from cycler import cycler import matplotlib as mpl import matplotlib.pyplot as plt @@ -201,7 +201,7 @@ def column(cls, *args, **kwargs): def plot_matrix( - matrix: NDArray[Shape['Any, Any'], Float], + matrix: NDARRAY, fig_name: str, labels: List[List[str]], clabel: str, diff --git a/farms_core/array/types.py b/farms_core/array/types.py index 741b1ca..b8555a6 100644 --- a/farms_core/array/types.py +++ b/farms_core/array/types.py @@ -1,49 +1,40 @@ """Types""" -from typing import Any -from importlib.metadata import version +from typing import TYPE_CHECKING, Any import numpy as np -from nptyping import NDArray -from ..sensors.sensor_convention import sc +import numpy.typing as npt +from ..sensors.sensor_convention import sc # pylint: disable=no-name-in-module -# pylint: disable=invalid-name + +if TYPE_CHECKING: + from typing import Annotated, TypeVar + DType = TypeVar("DType", bound=np.generic) # pylint: disable=invalid-name + def shaped_array(dtype: type[DType], shape: tuple): + """Type-level helper for shape-annotated NDArray.""" + return Annotated[npt.NDArray[dtype], shape] +else: + def shaped_array(_dtype, _shape): # type: ignore + """Runtime stub for shaped_array (returns NDArray[Any]).""" + return npt.NDArray[Any] SC_LINK_SIZE = sc.link_size -if version('nptyping') < '2.0.0': - NDARRAY = NDArray[Any, float] - NDARRAY_3 = NDArray[(3,), float] - NDARRAY_4 = NDArray[(4,), float] - NDARRAY_6 = NDArray[(6,), float] - NDARRAY_V1 = NDArray[(Any,), float] - NDARRAY_3_D = NDArray[(3,), np.double] - NDARRAY_4_D = NDArray[(4,), np.double] - NDARRAY_V1_I = NDArray[(Any,), np.uintc] - NDARRAY_V1_D = NDArray[(Any,), np.double] - NDARRAY_33 = NDArray[(3, 3), float] - NDARRAY_44 = NDArray[(4, 4), float] - NDARRAY_V2_D = NDArray[(Any, Any), np.double] - NDARRAY_X3_D = NDArray[(Any, 3), np.double] - NDARRAY_V3_D = NDArray[(Any, Any, Any), np.double] - NDARRAY_XX3_D = NDArray[(Any, Any, 3), np.double] - NDARRAY_XX4_D = NDArray[(Any, Any, 4), np.double] - NDARRAY_LINKS_ARRAY = NDArray[(Any, Any, SC_LINK_SIZE), np.double] -else: - from nptyping import Shape, Float - NDARRAY = NDArray[Any, Float] - NDARRAY_3 = NDArray[Shape['3'], Float] - NDARRAY_4 = NDArray[Shape['4'], Float] - NDARRAY_6 = NDArray[Shape['6'], Float] - NDARRAY_V1 = NDArray[Shape['*'], Float] - NDARRAY_3_D = NDArray[Shape['3'], np.double] - NDARRAY_4_D = NDArray[Shape['4'], np.double] - NDARRAY_V1_I = NDArray[Shape['*'], np.uintc] - NDARRAY_V1_D = NDArray[Shape['*'], np.double] - NDARRAY_33 = NDArray[Shape['3, 3'], Float] - NDARRAY_44 = NDArray[Shape['4, 4'], Float] - NDARRAY_V2_D = NDArray[Shape['*, *'], np.double] - NDARRAY_X3_D = NDArray[Shape['*, 3'], np.double] - NDARRAY_V3_D = NDArray[Shape['*, *, *'], np.double] - NDARRAY_XX3_D = NDArray[Shape['*, *, 3'], np.double] - NDARRAY_XX4_D = NDArray[Shape['*, *, 4'], np.double] - NDARRAY_LINKS_ARRAY = NDArray[Shape['*, *, SC_LINK_SIZE'], np.double] +NDARRAY = shaped_array(np.float32, ...) +NDARRAY_3 = shaped_array(np.float32, (3,)) +NDARRAY_4 = shaped_array(np.float32, (4,)) +NDARRAY_6 = shaped_array(np.float32, (6,)) +NDARRAY_V1 = shaped_array(np.float32, (...,)) +NDARRAY_V2 = shaped_array(np.float32, (..., ...)) +NDARRAY_3_D = shaped_array(np.double, (3,)) +NDARRAY_4_D = shaped_array(np.double, (4,)) +NDARRAY_V1_I = shaped_array(np.int64, (...,)) +NDARRAY_V1_D = shaped_array(np.double, (...,)) +NDARRAY_33 = shaped_array(np.float32, (3, 3)) +NDARRAY_44 = shaped_array(np.float32, (4, 4)) +NDARRAY_V2_D = shaped_array(np.double, (..., ...)) +NDARRAY_X3_D = shaped_array(np.double, (..., 3)) +NDARRAY_V3_D = shaped_array(np.double, (..., ..., ...)) +NDARRAY_XX3_D = shaped_array(np.double, (..., ..., 3)) +NDARRAY_XX4_D = shaped_array(np.double, (..., ..., 4)) +NDARRAY_XXXXX_UI8 = shaped_array(np.uint8, (..., ..., ..., ..., ...)) +NDARRAY_LINKS_ARRAY = shaped_array(np.double, (..., ..., SC_LINK_SIZE)) diff --git a/farms_core/model/control.py b/farms_core/model/control.py index 934e8a2..b6922b9 100644 --- a/farms_core/model/control.py +++ b/farms_core/model/control.py @@ -3,7 +3,7 @@ from typing import List, Tuple, Dict from enum import IntEnum import numpy as np -from nptyping import NDArray +from numpy.typing import NDArray from ..array.types import NDARRAY_V1 diff --git a/requirements.txt b/requirements.txt index 35d2b32..dabca83 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ numpy -nptyping matplotlib scipy cython