perf: faster model build and IDAKLU set-up for large models

CHANGELOG.md

@@ -1,9 +1,14 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Breaking changes
+
+- `StateVector` and `StateVectorDot` raise `ShapeError` for stepped, negative, overlapping or out-of-order slices; pass sorted, disjoint, non-negative slices. ([#5994](https://github.com/pybamm-team/PyBaMM/pull/5994))
+
 ## Deprecated
 
 - `FiniteVolume2D` is deprecated in favour of `FiniteVolumeUnstructured` and raises a `DeprecationWarning` when created. Remaining gaps before removal are tracked in [#5938](https://github.com/pybamm-team/PyBaMM/issues/5938). ([#5939](https://github.com/pybamm-team/PyBaMM/pull/5939))
 - `pybamm.Magnitude` is deprecated and will be removed in 27.1; use `pybamm.Component(symbol, 0)` / `pybamm.Component(symbol, 1)` for the lr/tb components of a `VectorField`. ([#5943](https://github.com/pybamm-team/PyBaMM/pull/5943))
+- The `evaluation_array` argument of `StateVector` and `StateVectorDot` is deprecated and ignored; drop the argument. ([#5994](https://github.com/pybamm-team/PyBaMM/pull/5994))
 
 ## Features
 
@@ -36,6 +41,7 @@
 
 - Read the Docs pull request builds drop `linkcheck` and the unpublished PDF, install the released solver wheel instead of compiling it, and are skipped when nothing the docs render changed, saving about 14 of 20 minutes. ([#5986](https://github.com/pybamm-team/PyBaMM/pull/5986))
 - Pull request CI no longer installs TeXLive, or the compilers and OpenBLAS that the prebuilt solver wheels make redundant. The latexify test mocks `sympy.preview` and runs everywhere. TeXLive downloads took 12–50 minutes in 6 of 15 runs. ([#5985](https://github.com/pybamm-team/PyBaMM/pull/5985))
+- Faster model build and `IDAKLUSolver` set-up for large models. ([#5994](https://github.com/pybamm-team/PyBaMM/pull/5994))
 - Experiments with a different model per step no longer compile Jacobians of the state mappers between models, which were unused. ([#5918](https://github.com/pybamm-team/PyBaMM/pull/5918))
 
 # [v26.10.1.1](https://github.com/pybamm-team/PyBaMM/tree/pybamm-v26.10.1.1) - 2026-10-19

packages/pybamm/src/pybamm/expression_tree/array.py

@@ -7,7 +7,8 @@
 import numpy as np
 import numpy.typing as npt
 import sympy
-from scipy.sparse import csr_matrix, issparse
+from pybammsolvers import idaklu
+from scipy.sparse import csc_matrix, csr_matrix, issparse
 
 import pybamm
 from pybamm.type_definitions import AuxiliaryDomainType, DomainsType, DomainType
@@ -132,12 +133,12 @@ def _make_entries_string(self, value: tuple | None) -> tuple:
 
     def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols):
         """See :meth:`pybamm.Symbol._to_casadi()`."""
-        return casadi.MX(self.evaluate(t, y, y_dot, inputs))
+        return casadi.MX(to_casadi_dm(self.evaluate(t, y, y_dot, inputs)))
 
     def _jac(self, variable) -> pybamm.Matrix:
         """See :meth:`pybamm.Symbol._jac()`."""
         # Return zeros of correct size
-        jac = csr_matrix((self.size, variable.evaluation_array.count(True)))
+        jac = csr_matrix((self.size, variable.size))
         return pybamm.Matrix(jac)
 
     def create_copy(
@@ -217,3 +218,58 @@ def meshgrid(
     """
     [x_grid, y_grid] = np.meshgrid(x.entries, y.entries)
     return pybamm.Array(x_grid), pybamm.Array(y_grid)
+
+
+def to_casadi_dm(entries) -> casadi.DM:
+    """
+    Convert a numpy array or scipy sparse matrix to a casadi.DM.
+
+    When pybammsolvers links the Python casadi library, the buffers are copied
+    once in C++; casadi's own conversion copies element by element.
+
+    Parameters
+    ----------
+    entries : numpy.ndarray or scipy.sparse matrix
+        The values, at most two-dimensional; a 1D array becomes a column.
+
+    Returns
+    -------
+    casadi.DM
+        The values, sparse if ``entries`` is sparse.
+    """
+    if issparse(entries):
+        compressed_rows = entries.format == "csr"
+        compressed = entries if compressed_rows else csc_matrix(entries)
+        # casadi aborts the process on unsorted or duplicate entries
+        if not compressed.has_canonical_format:
+            compressed = compressed.copy()
+            compressed.sum_duplicates()
+        if not idaklu.shares_python_casadi:
+            return casadi.DM(csc_matrix(compressed))
+        matrix = casadi.DM()
+        idaklu.fill_casadi_dm_sparse(
+            # SWIG's address attribute is not in casadi's stubs
+            # pyrefly: ignore[missing-attribute]
+            int(matrix.this),
+            compressed.shape[0],
+            compressed.shape[1],
+            compressed.indptr.astype(np.int64),
+            compressed.indices.astype(np.int64),
+            compressed.data.astype(np.float64),
+            compressed_rows=compressed_rows,
+        )
+    else:
+        if not idaklu.shares_python_casadi:
+            return casadi.DM(entries)
+        matrix = casadi.DM()
+        values = np.asarray(entries, dtype=np.float64)
+        if values.ndim < 2:
+            values = values.reshape(-1, 1)
+        idaklu.fill_casadi_dm_dense(
+            # pyrefly: ignore[missing-attribute]
+            int(matrix.this),
+            values.shape[0],
+            values.shape[1],
+            values.ravel(order="F"),
+        )
+    return matrix

packages/pybamm/src/pybamm/expression_tree/concatenations.py

@@ -695,22 +695,18 @@ def simplified_domain_concatenation(
     # Create the DomainConcatenation to read domain and child domain
     concat = DomainConcatenation(children, mesh, copy_this=copy_this)
     # Simplify Concatenation of StateVectors to a single StateVector
-    # The sum of the evalation arrays of the StateVectors must be exactly 1
+    # The StateVectors must cover their combined span exactly once
     if all(isinstance(child, pybamm.StateVector) for child in children):
         sv_children: list[pybamm.StateVector] = children  # type: ignore[assignment]
-        longest_eval_array = len(sv_children[-1]._evaluation_array)
-        eval_arrays = {}
-        for child in sv_children:
-            eval_arrays[child] = np.concatenate(
-                [
-                    child.evaluation_array,
-                    np.zeros(longest_eval_array - len(child.evaluation_array)),
-                ]
-            )
         first_start = sv_children[0].y_slices[0].start
         last_stop = sv_children[-1].y_slices[-1].stop
-        if all(
-            sum(array for array in eval_arrays.values())[first_start:last_stop] == 1
+        ranges = sorted(
+            (y_slice.start or 0, y_slice.stop)
+            for child in sv_children
+            for y_slice in child.y_slices
+        )
+        if pybamm.expression_tree.state_vector.ranges_tile(
+            ranges, first_start, last_stop
         ):
             return pybamm.StateVector(
                 slice(first_start, last_stop), domains=concat._domains

packages/pybamm/src/pybamm/expression_tree/input_parameter.py

@@ -88,7 +88,7 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols):
 
     def _jac(self, variable: pybamm.StateVector) -> pybamm.Matrix:
         """See :meth:`pybamm.Symbol._jac()`."""
-        n_variable = variable.evaluation_array.count(True)
+        n_variable = variable.size
         nan_vector = self._evaluate_for_shape()
         if isinstance(nan_vector, numbers.Number):
             n_self = 1

packages/pybamm/src/pybamm/expression_tree/state_vector.py

@@ -3,6 +3,8 @@
 #
 from __future__ import annotations
 
+import warnings
+
 import casadi
 import numpy as np
 import numpy.typing as npt
@@ -12,7 +14,31 @@
 from pybamm.type_definitions import AuxiliaryDomainType, DomainsType, DomainType
 
 
+def ranges_tile(ranges, start: int, stop: int) -> bool:
+    """
+    Whether sorted ``(start, stop)`` ranges cover ``[start, stop)`` exactly once.
+
+    Parameters
+    ----------
+    ranges : iterable of tuple of int
+        Half-open index ranges, sorted by start.
+    start, stop : int
+        The interval to cover.
+
+    Returns
+    -------
+    bool
+        True if the ranges are contiguous, non-overlapping and span the interval.
+    """
+    position = start
+    for range_start, range_stop in ranges:
+        if range_start != position:
+            return False
+        position = range_stop
+    return position == stop
+
+
 class StateVectorBase(pybamm.Symbol):
     """
     Node in the expression tree that holds a slice to read from an external vector type.
 
@@ -33,13 +59,18 @@
         (not both). In future, the 'domain' and 'auxiliary_domains' arguments may be
         deprecated.
     evaluation_array : list, optional
-        List of boolean arrays representing slices. Default is None, in which case the
-        evaluation_array is computed from y_slices.
+        Deprecated and ignored: the entries read are always given by ``y_slices``.
+
+    Raises
+    ------
+    :class:`pybamm.ShapeError`
+        If a slice has a step or a negative bound, or the slices are not in
+        increasing order or overlap.
     """
 
-    __slots__ = ("_evaluation_array", "_first_point", "_last_point", "_y_slices")
-    # the evaluation array determines the slices and end points
-    _id_excluded_fields = ("_y_slices", "_first_point", "_last_point")
+    __slots__ = ("_first_point", "_index_ranges", "_last_point", "_y_slices")
+    # The merged index ranges define the entries read, however the slices split them
+    _id_excluded_fields = ("_first_point", "_last_point", "_y_slices")
 
     def __init__(
         self,
@@ -54,6 +85,13 @@
         for y_slice in y_slices:
             if not isinstance(y_slice, slice):
                 raise TypeError("all y_slices must be slice objects")
+        if evaluation_array is not None:
+            warnings.warn(
+                "The 'evaluation_array' argument of StateVector is deprecated and "
+                "ignored; the entries read are given by the y_slices.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         if name is None:
             if y_slices[0].start is None:
                 name = base_name + f"[0:{y_slice.stop:d}"
@@ -70,7 +108,7 @@
         self._y_slices = y_slices
         self._first_point = y_slices[0].start
         self._last_point = y_slices[-1].stop
-        self._evaluation_array = self._make_evaluation_array(y_slices, evaluation_array)
+        self._index_ranges = self._merge_slices(y_slices)
         super().__init__(
             name=name,
             domain=domain,
@@ -82,12 +120,7 @@
     def _from_json(cls, snippet: dict):
         y_slices = [slice(s["start"], s["stop"], s["step"]) for s in snippet["y_slice"]]
 
-        return cls(
-            *y_slices,
-            name=snippet["name"],
-            domains=snippet["domains"],
-            evaluation_array=snippet["evaluation_array"],
-        )
+        return cls(*y_slices, name=snippet["name"], domains=snippet["domains"])
 
     @property
     def y_slices(self):
@@ -103,35 +136,76 @@
 
     @property
     def evaluation_array(self):
-        """Array to use for evaluating."""
-        return self._evaluation_array
+        """Boolean mask, as long as ``last_point``, of the entries of y read."""
+        array = np.zeros(self._last_point, dtype=bool)
+        for start, stop in self._index_ranges:
+            array[start:stop] = True
+        return array.tolist()
 
     @property
     def size(self):
-        return self.evaluation_array.count(True)
+        return sum(stop - start for start, stop in self._index_ranges)
+
+    def _read(self, vector):
+        """The entries of ``vector`` this state vector reads, as a column."""
+        if len(self._index_ranges) <= 1:
+            start, stop = self._index_ranges[0] if self._index_ranges else (0, 0)
+            out = vector[start:stop].copy()
+        else:
+            out = np.concatenate(
+                [vector[start:stop] for start, stop in self._index_ranges]
+            )
+        if isinstance(out, np.ndarray) and out.ndim == 1:
+            out = out[:, np.newaxis]
+        return out
+
+    @staticmethod
+    def _merge_slices(y_slices) -> tuple[tuple[int, int], ...]:
+        """
+        The entries read by ``y_slices``, as ``(start, stop)`` ranges with adjacent
+        slices merged.
+
+        Raises
+        ------
+        :class:`pybamm.ShapeError`
+            If a slice has a step or a negative bound, or the slices are not
+            increasing or overlap.
+        """
+        ranges: list[list[int]] = []
+        for y_slice in y_slices:
+            if y_slice.step not in (None, 1):
+                raise pybamm.ShapeError("y_slices must not have a step")
+            start = y_slice.start or 0
+            if start < 0 or y_slice.stop < 0:
+                raise pybamm.ShapeError("y_slices must not have negative bounds")
+            if ranges and start < ranges[-1][1]:
+                raise pybamm.ShapeError(
+                    "y_slices must be in increasing order and must not overlap"
+                )
+            if ranges and start == ranges[-1][1]:
+                ranges[-1][1] = y_slice.stop
+            elif y_slice.stop > start:
+                ranges.append([start, y_slice.stop])
+        return tuple((start, stop) for start, stop in ranges)
+
+    def __setstate__(self, state):
+        if "_evaluation_array" in state:  # pickled when the mask was stored
+            state = {
+                key: value for key, value in state.items() if key != "_evaluation_array"
+            }
+            state["_index_ranges"] = self._merge_slices(state["_y_slices"])
+        super().__setstate__(state)
 
     def set_evaluation_array(self, y_slices, evaluation_array):
-        """Set evaluation array using slices (deprecated)."""
+        """Set the slices to read (deprecated); ``evaluation_array`` is ignored."""
         pybamm.expression_tree.symbol._warn_mutation(
             "set_evaluation_array",
             "Construct a new StateVector with the desired slices.",
         )
-        object.__setattr__(
-            self,
-            "_evaluation_array",
-            self._make_evaluation_array(y_slices, evaluation_array),
-        )
-
-    @staticmethod
-    def _make_evaluation_array(y_slices, evaluation_array) -> list[bool]:
-        """The evaluation array: ``evaluation_array`` if given, else computed from
-        the slices."""
-        if evaluation_array is not None and pybamm.settings.debug_mode is False:
-            return evaluation_array
-        array = np.zeros(y_slices[-1].stop)
-        for y_slice in y_slices:
-            array[y_slice] = True
-        return [bool(x) for x in array]
+        object.__setattr__(self, "_y_slices", tuple(y_slices))
+        object.__setattr__(self, "_first_point", y_slices[0].start)
+        object.__setattr__(self, "_last_point", y_slices[-1].stop)
+        object.__setattr__(self, "_index_ranges", self._merge_slices(y_slices))
 
     def _jac_diff_vector(self, variable: pybamm.StateVectorBase):
         """
@@ -209,7 +283,6 @@
             *self.y_slices,
             name=self.name,
             domains=self._domains,
-            evaluation_array=self.evaluation_array,
         )
 
     def _evaluate_for_shape(self):
@@ -298,15 +371,11 @@
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
         if y is None:
             raise TypeError("StateVector cannot evaluate input 'y=None'")
-        if y.shape[0] < len(self.evaluation_array):
+        if y.shape[0] < self._last_point:
             raise ValueError(
                 "y is too short, so value with slice is smaller than expected"
             )
-
-        out = (y[: len(self._evaluation_array)])[self._evaluation_array]
-        if isinstance(out, np.ndarray) and out.ndim == 1:
-            out = out[:, np.newaxis]
-        return out
+        return self._read(y)
 
     def diff(self, variable: pybamm.Symbol):
         if variable == self:
@@ -316,7 +385,6 @@
                 *self._y_slices,
                 name=self.name + "'",
                 domains=self._domains,
-                evaluation_array=self.evaluation_array,
             )
         else:
             return pybamm.Scalar(0)
@@ -390,15 +458,11 @@
         """See :meth:`pybamm.Symbol._base_evaluate()`."""
         if y_dot is None:
             raise TypeError("StateVectorDot cannot evaluate input 'y_dot=None'")
-        if y_dot.shape[0] < len(self.evaluation_array):
+        if y_dot.shape[0] < self._last_point:
             raise ValueError(
                 "y_dot is too short, so value with slice is smaller than expected"
             )
-
-        out = (y_dot[: len(self._evaluation_array)])[self._evaluation_array]
-        if isinstance(out, np.ndarray) and out.ndim == 1:
-            out = out[:, np.newaxis]
-        return out
+        return self._read(y_dot)
 
     def diff(self, variable: pybamm.Symbol):
         if variable == self: