feat: added an option for multiple initial conditions in IDAKLU solver (#4981)

* initial conditions options and few tests

* added to changelog

* fixed tests, coverage, changelog and other minor changes

* changes in changelog

* made changes according to comments

* minor changes

* fixed tests and tolerances

* minor change

* style: pre-commit fixes

* Update CHANGELOG.md

* fixed docs

* minor change in test

* style: pre-commit fixes

* Update src/pybamm/solvers/idaklu_solver.py

Co-authored-by: Eric G. Kratz <[email protected]>

---------

Co-authored-by: Martin Robinson <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric G. Kratz <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -26,6 +26,7 @@
 - Creates `BaseProcessedVariable` to enable object combination when adding solutions together ([#5076](https://github.com/pybamm-team/PyBaMM/pull/5076))
 - Added a `Constant` symbol for named constants. This is a subclass of `Scalar` and is used to represent named constants such as the gas constant. This avoids constants being simplified out when constructing expressions. ([#5070](https://github.com/pybamm-team/PyBaMM/pull/5070))
 - Generalise `pybamm.DiscreteTimeSum` to allow it to be embedded in other expressions ([#5044](https://github.com/pybamm-team/PyBaMM/pull/5044))
+- Added an option for multiple initial conditions in IDAKLU solver ([#4981](https://github.com/pybamm-team/PyBaMM/pull/4981))
 - Adds `all` key-value pair to `output_variables` sensitivity dictionaries, accessible through `solution[var].sensitivities['all']`. Aligns shape with conventional solution sensitivities object. ([#5067](https://github.com/pybamm-team/PyBaMM/pull/5067))
 - Added a new `BaseHysteresisOpenCircuitPotential` class that sets variables for the lithiation and delithiation OCP and the hysteresis voltage (`H = U_lith - U_delith`). Allow the initial hysteresis state to be a function of position through the electrode. Allow the hysteresis decay rates of the Axen and Wycisk models to be a function of stoichiometry and temperature. Added a heat source term in each active material phase `Q_hys = i_vol * (U - U_eq)` where `i_vol` is the volumetric interfacial current density, `U` is the OCP (i.e. includes hysteresis), and `U_eq` is the "equilibrium OCP". Renamed the open-circuit potential models to be more descriptive. The options "Axen" and "Wycisk" are now "one-state hysteresis" and "one-state differential capacity hysteresis". The old option names still work but will raise a warning. ([#4893](https://github.com/pybamm-team/PyBaMM/pull/4893))
 - Add support for `output_variables` to `pybamm.DiscreteTimeSum` and `pybamm.ExplicitTimeIntegral` expressions. ([#5071](https://github.com/pybamm-team/PyBaMM/pull/5071))

src/pybamm/simulation.py

@@ -368,6 +368,7 @@ def solve(
         showprogress=False,
         inputs=None,
         t_interp=None,
+        initial_conditions=None,
         **kwargs,
     ):
         """
@@ -521,9 +522,13 @@ def solve(
                             pybamm.SolverWarning,
                             stacklevel=2,
                         )
-
             self._solution = solver.solve(
-                self._built_model, t_eval, inputs=inputs, t_interp=t_interp, **kwargs
+                self._built_model,
+                t_eval,
+                inputs=inputs,
+                t_interp=t_interp,
+                **kwargs,
+                initial_conditions=initial_conditions,
             )
 
         elif self.operating_mode == "with experiment":

src/pybamm/solvers/base_solver.py

@@ -676,6 +676,7 @@ def solve(
         nproc=None,
         calculate_sensitivities=False,
         t_interp=None,
+        initial_conditions=None,
     ):
         """
         Execute the solver setup and calculate the solution of the model at
@@ -706,7 +707,14 @@ def solve(
         t_interp : None, list or ndarray, optional
             The times (in seconds) at which to interpolate the solution. Defaults to None.
             Only valid for solvers that support intra-solve interpolation (`IDAKLUSolver`).
+        initial_conditions : dict, numpy.ndarray, or list, optional
+            Override the model’s default `y0`.  Can be:
 
+            - a dict mapping variable names → values
+            - a 1D array of length `n_states`
+            - a list of such overrides (one per parallel solve)
+
+            Only valid for IDAKLU solver.
         Returns
         -------
         :class:`pybamm.Solution` or list of :class:`pybamm.Solution` objects.
@@ -878,6 +886,7 @@ def solve(
                     t_eval[start_index:end_index],
                     model_inputs_list,
                     t_interp,
+                    initial_conditions,
                 )
             else:
                 ninputs = len(model_inputs_list)

src/pybamm/solvers/idaklu_solver.py

@@ -532,7 +532,44 @@ def __setstate__(self, d):
     def supports_parallel_solve(self):
         return True
 
-    def _integrate(self, model, t_eval, inputs_list=None, t_interp=None):
+    def _apply_solver_initial_conditions(self, model, initial_conditions):
+        """
+        Apply custom initial conditions to a model by overriding model.y0.
+
+        Parameters
+        ----------
+        model : pybamm.BaseModel
+            A model with a precomputed y0 vector.
+        initial_conditions : dict or numpy.ndarray
+            Either a mapping from variable names to values (scalar or array),
+            or a flat numpy array matching the length of model.y0.
+        """
+        if isinstance(initial_conditions, dict):
+            y0_np = (
+                model.y0.full() if isinstance(model.y0, casadi.DM) else model.y0.copy()
+            )
+
+            for var_name, value in initial_conditions.items():
+                found = False
+                for symbol, slice_info in model.y_slices.items():
+                    if symbol.name == var_name:
+                        var_slice = slice_info[0]
+                        y0_np[var_slice] = value
+                        found = True
+                        break
+                if not found:
+                    raise ValueError(f"Variable '{var_name}' not found in model")
+
+            model.y0 = casadi.DM(y0_np)
+
+        elif isinstance(initial_conditions, np.ndarray):
+            model.y0 = casadi.DM(initial_conditions)
+        else:
+            raise TypeError("Initial conditions must be dict or numpy array")
+
+    def _integrate(
+        self, model, t_eval, inputs_list=None, t_interp=None, initial_conditions=None
+    ):
         """
         Solve a DAE model defined by residuals with initial conditions y0.
 
@@ -547,6 +584,13 @@ def _integrate(self, model, t_eval, inputs_list=None, t_interp=None):
         t_interp : None, list or ndarray, optional
             The times (in seconds) at which to interpolate the solution. Defaults to `None`,
             which returns the adaptive time-stepping times.
+        initial_conditions : dict, numpy.ndarray, or list, optional
+            Override the model’s default `y0`.  Can be:
+
+            - a dict mapping variable names → values
+            - a 1D array of length `n_states`
+            - a list of such overrides (one per parallel solve)
+
         """
         if model.convert_to_format != "casadi":  # pragma: no cover
             # Shouldn't ever reach this point
@@ -565,11 +609,36 @@ def _integrate(self, model, t_eval, inputs_list=None, t_interp=None):
         else:
             inputs = np.array([[]] * len(inputs_list))
 
-        # stack y0full and ydot0full so they are a 2D array of shape (number_of_inputs, number_of_states + number_of_parameters * number_of_states)
-        # note that y0full and ydot0full are currently 1D arrays (i.e. independent of inputs), but in the future we will support
-        # different initial conditions for different inputs (see https://github.com/pybamm-team/PyBaMM/pull/4260). For now we just repeat the same initial conditions for each input
-        y0full = np.vstack([model.y0full] * len(inputs_list))
-        ydot0full = np.vstack([model.ydot0full] * len(inputs_list))
+        if initial_conditions is not None:
+            if isinstance(initial_conditions, list):
+                if len(initial_conditions) != len(inputs_list):
+                    raise ValueError(
+                        "Number of initial conditions must match number of input sets"
+                    )
+
+                y0_list = []
+
+                model_copy = model.new_copy()
+                for ic in initial_conditions:
+                    self._apply_solver_initial_conditions(model_copy, ic)
+                    y0_list.append(model_copy.y0.full().flatten())
+
+                y0full = np.vstack(y0_list)
+                ydot0full = np.zeros_like(y0full)
+
+            else:
+                self._apply_solver_initial_conditions(model, initial_conditions)
+
+                y0_np = model.y0.full()
+
+                y0full = np.vstack([y0_np for _ in range(len(inputs_list))])
+                ydot0full = np.zeros_like(y0full)
+        else:
+            # stack y0full and ydot0full so they are a 2D array of shape (number_of_inputs, number_of_states + number_of_parameters * number_of_states)
+            # note that y0full and ydot0full are currently 1D arrays (i.e. independent of inputs), but in the future we will support
+            # different initial conditions for different inputs. For now we just repeat the same initial conditions for each input
+            y0full = np.vstack([model.y0full] * len(inputs_list))
+            ydot0full = np.vstack([model.ydot0full] * len(inputs_list))
 
         atol = getattr(model, "atol", self.atol)
         atol = self._check_atol_type(atol, y0full.size)

src/pybamm/solvers/jax_solver.py

@@ -200,7 +200,9 @@ def supports_parallel_solve(self):
     def requires_explicit_sensitivities(self):
         return False
 
-    def _integrate(self, model, t_eval, inputs=None, t_interp=None):
+    def _integrate(
+        self, model, t_eval, inputs=None, t_interp=None, intial_conditions=None
+    ):
         """
         Solve a model defined by dydt with initial conditions y0.
 
@@ -220,6 +222,10 @@ def _integrate(self, model, t_eval, inputs=None, t_interp=None):
             various diagnostic messages.
 
         """
+        if intial_conditions is not None:  # pragma: no cover
+            raise NotImplementedError(
+                "Setting initial conditions is not yet implemented for the JAX IDAKLU solver"
+            )
         if isinstance(inputs, dict):
             inputs = [inputs]
         timer = pybamm.Timer()

tests/integration/test_solvers/test_idaklu.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 import pybamm
 
@@ -210,6 +211,71 @@ def test_with_experiments(self):
             sols[1].cycles[-1]["Current [A]"].data,
         )
 
+    @pytest.mark.parametrize(
+        "model_cls, make_ics",
+        [
+            (pybamm.lithium_ion.SPM, lambda y0: [y0, 2 * y0]),
+            (
+                pybamm.lithium_ion.DFN,
+                lambda y0: [y0, y0 * (1 + 0.01 * np.ones_like(y0))],
+            ),
+        ],
+    )
+    def test_multiple_initial_conditions_against_independent_solves(
+        self, model_cls, make_ics
+    ):
+        model = model_cls()
+        geom = model.default_geometry
+        pv = model.default_parameter_values
+        pv.process_model(model)
+        pv.process_geometry(geom)
+        mesh = pybamm.Mesh(geom, model.default_submesh_types, model.default_var_pts)
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+
+        t_eval = np.array([0, 1])
+        solver = pybamm.IDAKLUSolver()
+
+        base_sol = solver.solve(model, t_eval)
+        y0_base = base_sol.y[:, 0]
+
+        ics = make_ics(y0_base)
+        inputs = [{}] * len(ics)
+
+        multi_sols = solver.solve(
+            model,
+            t_eval,
+            inputs=inputs,
+            initial_conditions=ics,
+        )
+        assert isinstance(multi_sols, list) and len(multi_sols) == 2
+
+        indep_sols = []
+        for ic in ics:
+            sol_indep = solver.solve(
+                model, t_eval, inputs=[{}], initial_conditions=[ic]
+            )
+            if isinstance(sol_indep, list):
+                sol_indep = sol_indep[0]
+            indep_sols.append(sol_indep)
+
+        if model_cls is pybamm.lithium_ion.SPM:
+            rtol, atol = 1e-8, 1e-10
+        else:
+            rtol, atol = 1e-6, 1e-8
+
+        for idx in (0, 1):
+            sol_vec = multi_sols[idx]
+            sol_ind = indep_sols[idx]
+
+            np.testing.assert_allclose(sol_vec.t, sol_ind.t, rtol=1e-12, atol=0)
+            np.testing.assert_allclose(sol_vec.y, sol_ind.y, rtol=rtol, atol=atol)
+
+            if model_cls is pybamm.lithium_ion.SPM:
+                np.testing.assert_allclose(
+                    sol_vec.y[:, 0], ics[idx], rtol=1e-8, atol=1e-10
+                )
+
     def test_outvars_with_experiments_multi_simulation(self):
         model = pybamm.lithium_ion.SPM()