Updated np.ndarray to numpy.typing.NDArray everywhere by writing a script to do the changes automatically

benchmarks/different_model_options.py

@@ -33,8 +33,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: numpy.typing.NDArray
+    t_interp: numpy.typing.NDArray | None
 
     def solve_setup(self, parameter, model_, option, value, solver_class):
         import importlib

benchmarks/time_solve_models.py

@@ -30,8 +30,8 @@ class TimeSolveSPM:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
-    t_interp: np.ndarray | None
+    t_eval: numpy.typing.NDArray
+    t_interp: numpy.typing.NDArray | None
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()
@@ -96,7 +96,7 @@ class TimeSolveSPMe:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
+    t_eval: numpy.typing.NDArray
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()
@@ -160,7 +160,7 @@ class TimeSolveDFN:
     )
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
-    t_eval: np.ndarray
+    t_eval: numpy.typing.NDArray
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()

src/pybamm/discretisations/discretisation.py

@@ -1012,7 +1012,7 @@ def check_initial_conditions(self, model):
         # Individual
         for var, eqn in model.initial_conditions.items():
             ic_eval = eqn.evaluate(t=0, inputs="shape test")
-            if not isinstance(ic_eval, np.ndarray):
+            if not isinstance(ic_eval, numpy.typing.NDArray):
                 raise pybamm.ModelError(
                     "initial conditions must be numpy array after discretisation but "
                     f"they are {type(ic_eval)} for variable '{var}'."

src/pybamm/experiment/step/base_step.py

@@ -74,7 +74,7 @@ def __init__(
         self.input_duration = duration
         self.input_value = value
         # Check if drive cycle
-        is_drive_cycle = isinstance(value, np.ndarray)
+        is_drive_cycle = isinstance(value, numpy.typing.NDArray)
         is_python_function = callable(value)
         if is_drive_cycle:
             if value.ndim != 2 or value.shape[1] != 2:
@@ -260,7 +260,7 @@ def default_duration(self, value):
         Default duration for the step is one day (24 hours) or the duration of the
         drive cycle
         """
-        if isinstance(value, np.ndarray):
+        if isinstance(value, numpy.typing.NDArray):
             t = value[:, 0]
             return t[-1]
         else:

src/pybamm/expression_tree/array.py

@@ -38,7 +38,7 @@ class Array(pybamm.Symbol):
 
     def __init__(
         self,
-        entries: np.ndarray | list[float] | csr_matrix,
+        entries: numpy.typing.NDArray | list[float] | csr_matrix,
         name: str | None = None,
         domain: DomainType = None,
         auxiliary_domains: AuxiliaryDomainType = None,
@@ -144,8 +144,8 @@ def create_copy(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
@@ -165,7 +165,7 @@ def to_json(self):
         Method to serialise an Array object into JSON.
         """
 
-        if isinstance(self.entries, np.ndarray):
+        if isinstance(self.entries, numpy.typing.NDArray):
             matrix = self.entries.tolist()
         elif isinstance(self.entries, csr_matrix):
             matrix = {

src/pybamm/expression_tree/binary_operators.py

@@ -22,13 +22,13 @@ def _preprocess_binary(
 ) -> tuple[pybamm.Symbol, pybamm.Symbol]:
     if isinstance(left, (float, int, np.number)):
         left = pybamm.Scalar(left)
-    elif isinstance(left, np.ndarray):
+    elif isinstance(left, numpy.typing.NDArray):
         if left.ndim > 1:
             raise ValueError("left must be a 1D array")
         left = pybamm.Vector(left)
     if isinstance(right, (float, int, np.number)):
         right = pybamm.Scalar(right)
-    elif isinstance(right, np.ndarray):
+    elif isinstance(right, numpy.typing.NDArray):
         if right.ndim > 1:
             raise ValueError("right must be a 1D array")
         right = pybamm.Vector(right)
@@ -152,8 +152,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""
@@ -558,7 +558,7 @@ def _binary_jac(self, left_jac, right_jac):
     def _binary_evaluate(self, left, right):
         """See :meth:`pybamm.BinaryOperator._binary_evaluate()`."""
         # numpy 1.25 deprecation warning: extract value from numpy arrays
-        if isinstance(right, np.ndarray):
+        if isinstance(right, numpy.typing.NDArray):
             return int(left == right.item())
         else:
             return int(left == right)

src/pybamm/expression_tree/concatenations.py

@@ -111,7 +111,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[numpy.typing.NDArray]):
         """See :meth:`Concatenation._concatenation_evaluate()`."""
         if len(children_eval) == 0:
             return np.array([])
@@ -121,8 +121,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""
@@ -366,7 +366,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[numpy.typing.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,4 @@
 import pybamm
-import numpy as np
 
 
 class DiscreteTimeData(pybamm.Interpolant):
@@ -19,7 +18,9 @@ class DiscreteTimeData(pybamm.Interpolant):
 
     """
 
-    def __init__(self, time_points: np.ndarray, data: np.ndarray, name: str):
+    def __init__(
+        self, time_points: numpy.typing.NDArray, data: numpy.typing.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

@@ -122,8 +122,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""

src/pybamm/expression_tree/independent_variable.py

@@ -3,7 +3,6 @@
 #
 from __future__ import annotations
 import sympy
-import numpy as np
 
 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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/input_parameter.py

@@ -88,8 +88,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         # inputs should be a dictionary

src/pybamm/expression_tree/interpolant.py

@@ -43,8 +43,8 @@ class Interpolant(pybamm.Function):
 
     def __init__(
         self,
-        x: np.ndarray | Sequence[np.ndarray],
-        y: np.ndarray,
+        x: numpy.typing.NDArray | Sequence[numpy.typing.NDArray],
+        y: numpy.typing.NDArray,
         children: Sequence[pybamm.Symbol] | pybamm.Time,
         name: str | None = None,
         interpolator: str | None = "linear",
@@ -96,7 +96,7 @@ def __init__(
                 x1 = x[0]
             else:
                 x1 = x
-                x: list[np.ndarray] = [x]  # type: ignore[no-redef]
+                x: list[numpy.typing.NDArray] = [x]  # type: ignore[no-redef]
             x2 = None
             if x1.shape[0] != y.shape[0]:
                 raise ValueError(
@@ -269,7 +269,7 @@ def create_copy(self, new_children=None, perform_simplifications=True):
     def _function_evaluate(self, evaluated_children):
         children_eval_flat = []
         for child in evaluated_children:
-            if isinstance(child, np.ndarray):
+            if isinstance(child, numpy.typing.NDArray):
                 children_eval_flat.append(child.flatten())
             else:
                 children_eval_flat.append(child)

src/pybamm/expression_tree/matrix.py

@@ -16,7 +16,7 @@ class Matrix(pybamm.Array):
 
     def __init__(
         self,
-        entries: np.ndarray | list[float] | csr_matrix,
+        entries: numpy.typing.NDArray | list[float] | csr_matrix,
         name: str | None = None,
         domain: DomainType = None,
         auxiliary_domains: AuxiliaryDomainType = None,

src/pybamm/expression_tree/operations/evaluate_python.py

@@ -539,7 +539,7 @@ def __init__(self, symbol: pybamm.Symbol):
 
         # convert all numpy constants to device vectors
         for symbol_id in constants:
-            if isinstance(constants[symbol_id], np.ndarray):
+            if isinstance(constants[symbol_id], numpy.typing.NDArray):
                 constants[symbol_id] = jax.device_put(constants[symbol_id])
 
         # get a list of constant arguments to input to the function
@@ -614,7 +614,7 @@ def _demote_64_to_32(cls, c):
             c = jax.numpy.int32(c)
         if isinstance(c, np.int64):
             c = c.astype(jax.numpy.int32)
-        if isinstance(c, np.ndarray):
+        if isinstance(c, numpy.typing.NDArray):
             if c.dtype == np.float64:
                 c = c.astype(jax.numpy.float32)
             if c.dtype == np.int64:

src/pybamm/expression_tree/scalar.py

@@ -53,7 +53,7 @@ def value(self):
     def value(self, value):
         self._value = (
             np.float64(value.item())
-            if isinstance(value, np.ndarray)
+            if isinstance(value, numpy.typing.NDArray)
             else np.float64(value)
         )
 
@@ -66,8 +66,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""

src/pybamm/expression_tree/state_vector.py

@@ -281,8 +281,8 @@ def __init__(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
@@ -294,7 +294,7 @@ def _base_evaluate(
             )
 
         out = (y[: len(self._evaluation_array)])[self._evaluation_array]
-        if isinstance(out, np.ndarray) and out.ndim == 1:
+        if isinstance(out, numpy.typing.NDArray) and out.ndim == 1:
             out = out[:, np.newaxis]
         return out
 
@@ -365,8 +365,8 @@ def __init__(
     def _base_evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
@@ -378,7 +378,7 @@ def _base_evaluate(
             )
 
         out = (y_dot[: len(self._evaluation_array)])[self._evaluation_array]
-        if isinstance(out, np.ndarray) and out.ndim == 1:
+        if isinstance(out, numpy.typing.NDArray) and out.ndim == 1:
             out = out[:, np.newaxis]
         return out
 

src/pybamm/expression_tree/symbol.py

@@ -132,7 +132,7 @@ def is_matrix_x(expr: Symbol, x: int):
                     and np.all(result.__dict__["data"] == x)
                 )
             )
-        ) or (isinstance(result, np.ndarray) and np.all(result == x))
+        ) or (isinstance(result, numpy.typing.NDArray) and np.all(result == x))
     else:
         return False
 
@@ -168,18 +168,18 @@ def simplify_if_constant(symbol: pybamm.Symbol):
         if result is not None:
             if (
                 isinstance(result, numbers.Number)
-                or (isinstance(result, np.ndarray) and result.ndim == 0)
+                or (isinstance(result, numpy.typing.NDArray) and result.ndim == 0)
                 or isinstance(result, np.bool_)
             ):
                 # type-narrow for Scalar
                 new_result = cast(float, result)
                 return pybamm.Scalar(new_result)
-            elif isinstance(result, np.ndarray) or issparse(result):
+            elif isinstance(result, numpy.typing.NDArray) or issparse(result):
                 if result.ndim == 1 or result.shape[1] == 1:
                     return pybamm.Vector(result, domains=symbol.domains)
                 else:
                     # Turn matrix of zeros into sparse matrix
-                    if isinstance(result, np.ndarray) and np.all(result == 0):
+                    if isinstance(result, numpy.typing.NDArray) and np.all(result == 0):
                         result = csr_matrix(result)
                     return pybamm.Matrix(result, domains=symbol.domains)
 
@@ -759,8 +759,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """
@@ -791,8 +791,8 @@ def _base_evaluate(
     def evaluate(
         self,
         t: float | None = None,
-        y: np.ndarray | None = None,
-        y_dot: np.ndarray | None = None,
+        y: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.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

@@ -93,8 +93,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: numpy.typing.NDArray | None = None,
+        y_dot: numpy.typing.NDArray | None = None,
         inputs: dict | str | None = None,
     ):
         """See :meth:`pybamm.Symbol.evaluate()`."""