Replace np.ndarray with npt.NDArray for improved type hinting

benchmarks/different_model_options.py

@@ -1,6 +1,7 @@
 import pybamm
 from benchmarks.benchmark_utils import set_random_seed
 import numpy as np
+import numpy.typing as npt
 
 
 def compute_discretisation(model, param):
@@ -33,8 +34,8 @@ def build_model(parameter, model_, option, value):
 class SolveModel:
     solver: pybamm.BaseSolver
     model: pybamm.BaseModel
-    t_eval: np.ndarray
-    t_interp: np.ndarray | None
+    t_eval: npt.NDArray
+    t_interp: npt.NDArray | None
 
     def solve_setup(self, parameter, model_, option, value, solver_class):
         self.solver = solver_class()

benchmarks/time_solve_models.py

@@ -4,6 +4,7 @@
 import pybamm
 from benchmarks.benchmark_utils import set_random_seed
 import numpy as np
+import numpy.typing as npt
 
 
 def solve_model_once(model, solver, t_eval, t_interp):
@@ -30,8 +31,8 @@ class TimeSolveSPM:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
-    t_interp: np.ndarray | None
+    t_eval: npt.NDArray
+    t_interp: npt.NDArray | None
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()
@@ -96,7 +97,7 @@ class TimeSolveSPMe:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
+    t_eval: npt.NDArray
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()
@@ -160,7 +161,7 @@ class TimeSolveDFN:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
+    t_eval: npt.NDArray
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()

src/pybamm/expression_tree/array.py

@@ -3,6 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
+import numpy.typing as npt
 from scipy.sparse import csr_matrix, issparse
 
 import pybamm
@@ -38,7 +39,7 @@ class Array(pybamm.Symbol):
 
     def __init__(
         self,
-        entries: np.ndarray | list[float] | csr_matrix,
+        entries: npt.NDArray | list[float] | csr_matrix,
         name: str | None = None,
         domain: DomainType = None,
         auxiliary_domains: AuxiliaryDomainType = None,
@@ -144,8 +145,8 @@ def create_copy(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/binary_operators.py

@@ -5,6 +5,7 @@
 import numbers
 
 import numpy as np
+import numpy.typing as npt
 import sympy
 from scipy.sparse import csr_matrix, issparse
 import functools
@@ -152,8 +153,8 @@ def _binary_new_copy(self, left: ChildSymbol, right: ChildSymbol):
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""

src/pybamm/expression_tree/concatenations.py

@@ -7,6 +7,7 @@
 from typing import Optional
 
 import numpy as np
+import numpy.typing as npt
 import sympy
 from scipy.sparse import issparse, vstack
 from collections.abc import Sequence
@@ -112,7 +113,7 @@ def get_children_domains(self, children: Sequence[pybamm.Symbol]):
 
         return domains
 
-    def _concatenation_evaluate(self, children_eval: list[np.ndarray]):
+    def _concatenation_evaluate(self, children_eval: list[npt.NDArray]):
         """See :meth:`Concatenation._concatenation_evaluate()`."""
         if len(children_eval) == 0:
             return np.array([])
@@ -122,8 +123,8 @@ def _concatenation_evaluate(self, children_eval: list[np.ndarray]):
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""
@@ -367,7 +368,7 @@ def create_slices(self, node: pybamm.Symbol) -> defaultdict:
                 start = end
         return slices
 
-    def _concatenation_evaluate(self, children_eval: list[np.ndarray]):
+    def _concatenation_evaluate(self, children_eval: list[npt.NDArray]):
         """See :meth:`Concatenation._concatenation_evaluate()`."""
         # preallocate vector
         vector = np.empty((self._size, 1))

src/pybamm/expression_tree/discrete_time_sum.py

@@ -1,5 +1,5 @@
 import pybamm
-import numpy as np
+import numpy.typing as npt
 
 
 class DiscreteTimeData(pybamm.Interpolant):
@@ -19,7 +19,7 @@ class DiscreteTimeData(pybamm.Interpolant):
 
     """
 
-    def __init__(self, time_points: np.ndarray, data: np.ndarray, name: str):
+    def __init__(self, time_points: npt.NDArray, data: npt.NDArray, name: str):
         super().__init__(time_points, data, pybamm.t, name)
 
     def create_copy(self, new_children=None, perform_simplifications=True):

src/pybamm/expression_tree/functions.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import numpy as np
+import numpy.typing as npt
 from scipy import special
 import sympy
 from typing import Callable
@@ -122,8 +123,8 @@ def _function_jac(self, children_jacs):
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""

src/pybamm/expression_tree/independent_variable.py

@@ -3,8 +3,7 @@
 #
 from __future__ import annotations
 import sympy
-import numpy as np
-
+import numpy.typing as npt
 import pybamm
 from pybamm.type_definitions import DomainType, AuxiliaryDomainType, DomainsType
 
@@ -94,8 +93,8 @@ def create_copy(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/input_parameter.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 import numbers
 import numpy as np
+import numpy.typing as npt
 import scipy.sparse
 import pybamm
 
@@ -88,8 +89,8 @@ def _jac(self, variable: pybamm.StateVector) -> pybamm.Matrix:
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         # inputs should be a dictionary

src/pybamm/expression_tree/interpolant.py

@@ -3,6 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
+import numpy.typing as npt
 from scipy import interpolate
 from collections.abc import Sequence
 import numbers
@@ -43,8 +44,8 @@ class Interpolant(pybamm.Function):
 
     def __init__(
         self,
-        x: np.ndarray | Sequence[np.ndarray],
-        y: np.ndarray,
+        x: npt.NDArray | Sequence[npt.NDArray],
+        y: npt.NDArray,
         children: Sequence[pybamm.Symbol] | pybamm.Time,
         name: str | None = None,
         interpolator: str | None = "linear",
@@ -96,7 +97,7 @@ def __init__(
                 x1 = x[0]
             else:
                 x1 = x
-                x: list[np.ndarray] = [x]  # type: ignore[no-redef]
+                x: list[npt.NDArray] = [x]  # type: ignore[no-redef]
             x2 = None
             if x1.shape[0] != y.shape[0]:
                 raise ValueError(

src/pybamm/expression_tree/matrix.py

@@ -3,6 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
+import numpy.typing as npt
 from scipy.sparse import csr_matrix, issparse
 
 import pybamm
@@ -16,7 +17,7 @@ class Matrix(pybamm.Array):
 
     def __init__(
         self,
-        entries: np.ndarray | list[float] | csr_matrix,
+        entries: npt.NDArray | list[float] | csr_matrix,
         name: str | None = None,
         domain: DomainType = None,
         auxiliary_domains: AuxiliaryDomainType = None,

src/pybamm/expression_tree/scalar.py

@@ -3,6 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
+import numpy.typing as npt
 import sympy
 from typing import Literal
 
@@ -66,8 +67,8 @@ def set_id(self):
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/state_vector.py

@@ -3,6 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
+import numpy.typing as npt
 from scipy.sparse import csr_matrix, vstack
 
 import pybamm
@@ -281,8 +282,8 @@ def __init__(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
@@ -365,8 +366,8 @@ def __init__(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/symbol.py

@@ -6,6 +6,7 @@
 import warnings
 
 import numpy as np
+import numpy.typing as npt
 import sympy
 from scipy.sparse import csr_matrix, issparse
 from functools import cached_property
@@ -769,8 +770,8 @@ def _jac(self, variable):
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """
@@ -801,8 +802,8 @@ def _base_evaluate(
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ) -> ChildValue:
         """Evaluate expression tree (wrapper to allow using dict of known values).

src/pybamm/expression_tree/unary_operators.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import numpy as np
+import numpy.typing as npt
 from scipy.sparse import csr_matrix, issparse
 import sympy
 import pybamm
@@ -93,8 +94,8 @@ def _unary_evaluate(self, child):
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""

src/pybamm/expression_tree/vector.py

@@ -3,7 +3,7 @@
 #
 from __future__ import annotations
 import numpy as np
-
+import numpy.typing as npt
 import pybamm
 from pybamm.type_definitions import DomainType, AuxiliaryDomainType, DomainsType
 
@@ -15,7 +15,7 @@ class Vector(pybamm.Array):
 
     def __init__(
         self,
-        entries: np.ndarray | list[float] | np.matrix,
+        entries: npt.NDArray | list[float] | np.matrix,
         name: str | None = None,
         domain: DomainType = None,
         auxiliary_domains: AuxiliaryDomainType = None,

src/pybamm/models/event.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-import numpy as np
-
+import numpy.typing as npt
 from typing import TypeVar
 
 
@@ -75,8 +74,8 @@ def _from_json(cls: type[E], snippet: dict) -> E:
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: npt.NDArray | None = None,
+        y_dot: npt.NDArray | None = None,
         inputs: dict | None = None,
     ):
         """