Add IDAKLU upsampling_factor for scalar output resolution

CHANGELOG.md

@@ -7,6 +7,7 @@
 
 ## Features
 
+- Added the `IDAKLUSolver` option `upsampling_factor`: a value above 1.0 saves extra points inside each step until linear interpolation of the voltage (or current, for voltage control) is within the solver tolerances divided by `factor - 1`. ([#5998](https://github.com/pybamm-team/PyBaMM/pull/5998))
 - Replaced Codacy with open checks that also run locally. Ruff's `flake8-bandit` (`S`) rules replace Bandit, and CodeQL scans the Python and C++ code. `nox -s type-check` runs Pyrefly over `pybamm`, which now type-checks cleanly, and a new type error fails CI. clang-tidy reports solver findings to code scanning as an advisory check. ([#5974](https://github.com/pybamm-team/PyBaMM/pull/5974))
 
 ## Bug fixes

packages/pybamm/src/pybamm/experiment/step/base_step.py

@@ -296,6 +296,9 @@ def __hash__(self):
     #: their target value collapse to one unified-mode branch.
     control_kind = None
 
+    #: The scalar output the solver's ``upsampling_factor`` resolves.
+    upsampling_variable = "Voltage [V]"
+
     def get_control_residual(self, variables, target=None):
         """Control residual. ``target`` is the per-step input that collapsible steps
         share; when it is ``None`` the step falls back to its own ``_default_target``."""
@@ -457,11 +460,13 @@ def process_model(
             new_parameter_values["Ambient temperature [K]"] = self.temperature
 
         # Parameterise the model
-        return new_parameter_values.process_model(
+        processed_model = new_parameter_values.process_model(
             new_model,
             inplace=inplace,
             delayed_variable_processing=delayed_variable_processing,
         )
+        processed_model.upsampling_variable = self.upsampling_variable
+        return processed_model
 
     def get_termination_events(self, variables):
         events = []

packages/pybamm/src/pybamm/experiment/step/steps.py

@@ -194,6 +194,7 @@ class Voltage(BaseStepImplicit):
     """
 
     control_kind = ControlKind.VOLTAGE
+    upsampling_variable = "Current [A]"
 
     def get_parameter_values(self, variables):
         return {"Voltage function [V]": self.value}

packages/pybamm/src/pybamm/solvers/idaklu_solver.py

@@ -13,7 +13,7 @@
 
 import pybamm
 from pybamm.codegen.compilation import aot_compile
-from pybamm.solvers.base_solver import flatten_inputs, stack_inputs
+from pybamm.solvers.base_solver import flatten_inputs, process, stack_inputs
 from pybamm.solvers.observation import OutputAssembly
 
 _UNSET = object()
@@ -85,6 +85,7 @@ def _sensitivity_scales(inputs_dict: dict, sensitivity_names: list[str]) -> np.n
     "rootfn",
     "alg_res",
     "alg_jac",
+    "upsampling_fcn",
 )
 _SETUP_FCN_LIST_KEYS = (
     "var_idaklu_fcns",
@@ -241,6 +242,11 @@ class IDAKLUSolver(pybamm.BaseSolver):
                 # only active if `hermite_interpolation` is True and sensitivities
                 # are disabled.
                 "hermite_reduction_factor": 1.0,
+                # Setting upsampling_factor > 1.0 adds points within each step
+                # until linear interpolation of the model's `upsampling_variable`
+                # (default: the voltage) is within the solver tolerances divided
+                # by (factor - 1). Ignored with t_interp.
+                "upsampling_factor": 1.0,
                 ## Initial conditions calculation
                 # Positive constant in the Newton iteration convergence test within the
                 # initial condition calculation
@@ -370,6 +376,7 @@ def _combine_options(self, user_options: dict | None) -> dict:
             "suppress_algebraic_error": False,
             "hermite_interpolation": True,
             "hermite_reduction_factor": 1.0,
+            "upsampling_factor": 1.0,
             "nonlinear_convergence_coefficient_ic": 0.0033,
             "max_num_steps_ic": 50,
             "max_num_jacobians_ic": 40,
@@ -407,6 +414,15 @@ def _check_options(self, options: dict):
             if hermite_reduction_factor < 1.0:
                 raise pybamm.SolverError("hermite_reduction_factor must be >= 1.0.")
 
+        upsampling_factor = options["upsampling_factor"]
+        if upsampling_factor < 1.0:
+            raise pybamm.SolverError("upsampling_factor must be >= 1.0.")
+        if upsampling_factor > 1.0 and hermite_reduction_factor > 1.0:
+            raise pybamm.SolverError(
+                "upsampling_factor and hermite_reduction_factor are "
+                "mutually exclusive; at most one can be > 1.0."
+            )
+
         if not isinstance(options["compile"], bool):
             raise pybamm.SolverError("compile must be a bool")
 
@@ -642,6 +658,8 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
                 stacklevel=2,
             )
 
+        upsampling_fcn = self._upsampling_fcn(model, inputs_dict)
+
         # Collect all casadi functions for AOT compilation. With compile=True,
         # all functions are bundled into a single shared library via one gcc
         # invocation; each serialized Function then points at its entry point.
@@ -658,6 +676,7 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
                     rootfn,
                     alg_res_fn,
                     jac_alg_fn,
+                    upsampling_fcn,
                 ),
                 strict=True,
             )
@@ -743,6 +762,7 @@ def to_idaklu(fn):
             options=self._options,
             alg_res=self._setup["alg_res"],
             alg_jac=self._setup["alg_jac"],
+            upsampling_fcn=self._setup["upsampling_fcn"],
         )
 
         # C++ solver group now owns each casadi::Function via shared_ptr;
@@ -760,6 +780,34 @@ def to_idaklu(fn):
 
         return base_set_up_return
 
+    def _upsampling_fcn(self, model, inputs_dict):
+        """``g(t, y, p)``: the model's scalar ``upsampling_variable``.
+
+        Empty if upsampling is off, or the variable is missing or an algebraic
+        state, which the integrator's error test already resolves.
+        """
+        empty = casadi.Function("empty_upsampling", [], [])
+        name = getattr(model, "upsampling_variable", "Voltage [V]")
+        if self._options["upsampling_factor"] <= 1.0 or name not in model.variables:
+            return empty
+        # Model copies keep a state's name but not its id
+        variable = model.variables[name]
+        if isinstance(variable, pybamm.Variable) and variable.name in {
+            state.name for state in model.algebraic
+        }:
+            return empty
+        upsampling_fcn = process(
+            model.get_processed_variable_or_event(name),
+            pybamm.BaseSolver._wrangle_name(name),
+            self._get_vars_for_processing(model, inputs_dict),
+            use_jacobian=False,
+        )[0]
+        if upsampling_fcn.numel_out() != 1:
+            raise pybamm.SolverError(
+                f"upsampling_factor needs a scalar variable, got '{name}'"
+            )
+        return upsampling_fcn
+
     def __getstate__(self):
         # Drop the rebuildable state (C++ solver + casadi.Function graphs)
         # so the next solve() rebuilds from the model rather than shipping
@@ -771,6 +819,8 @@ def __getstate__(self):
 
     def __setstate__(self, d):
         self.__dict__.update(d)
+        # Pickles from older versions lack newer options
+        self._options = self._combine_options(self._options)
         # Restore the empty defaults BaseSolver.__init__ would set, so the
         # solver reads as "not yet set up" until the next solve().
         self._model_set_up = {}

packages/pybamm/tests/unit/test_experiments/test_base_step.py

@@ -160,3 +160,18 @@ def test_parse_termination_requires_operator_for_input_parameter():
         pybamm.experiment.step.base_step._parse_termination(
             "2A", pybamm.InputParameter("I_app")
         )
+
+
+@pytest.mark.parametrize(
+    "step, variable",
+    [
+        (pybamm.step.current(1), "Voltage [V]"),
+        (pybamm.step.c_rate(1), "Voltage [V]"),
+        (pybamm.step.rest(), "Voltage [V]"),
+        (pybamm.step.power(1), "Voltage [V]"),
+        (pybamm.step.resistance(1), "Voltage [V]"),
+        (pybamm.step.voltage(4.1), "Current [A]"),
+    ],
+)
+def test_upsampling_variable(step, variable):
+    assert step.upsampling_variable == variable

packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py

@@ -2733,3 +2733,162 @@ def test_iterative_linear_solver_statistics(
         stats = solution.solver_statistics
         assert stats.number_of_steps > 0
         assert stats.number_of_nonlinear_solver_iterations >= stats.number_of_steps
+
+
+class TestUpsampling:
+    @staticmethod
+    def _solve(options=None, **solver_kwargs):
+        parameter_values = pybamm.ParameterValues("Chen2020")
+        parameter_values["Current function [A]"] = 10
+        solver = pybamm.IDAKLUSolver(options=options, **solver_kwargs)
+        sim = pybamm.Simulation(
+            pybamm.lithium_ion.SPM(), parameter_values=parameter_values, solver=solver
+        )
+        return sim.solve([0, 3600])
+
+    def test_adds_points_in_every_storage_mode(self):
+        baseline = self._solve()
+        upsampled = self._solve({"upsampling_factor": 2.0})
+        assert upsampled.t.size > baseline.t.size
+        assert np.all(np.isin(baseline.t, upsampled.t))
+
+        # Upsampling uses the integrator's own interpolant, not stored derivatives
+        no_hermite = self._solve(
+            {"upsampling_factor": 2.0, "hermite_interpolation": False}
+        )
+        outputs_only = self._solve(
+            {"upsampling_factor": 2.0}, output_variables=["Voltage [V]"]
+        )
+        np.testing.assert_array_equal(no_hermite.t, upsampled.t)
+        np.testing.assert_array_equal(outputs_only.t, upsampled.t)
+        np.testing.assert_allclose(
+            outputs_only["Voltage [V]"].entries,
+            upsampled["Voltage [V]"].entries,
+            rtol=1e-12,
+        )
+
+    def test_resolves_voltage_linear_interpolant(self):
+        reference = self._solve(rtol=1e-10, atol=1e-12)
+        ratios = []
+        for factor in [1.0, 2.0]:
+            solver = pybamm.IDAKLUSolver(options={"upsampling_factor": factor})
+            parameter_values = pybamm.ParameterValues("Chen2020")
+            parameter_values["Current function [A]"] = 10
+            sim = pybamm.Simulation(
+                pybamm.lithium_ion.SPM(),
+                parameter_values=parameter_values,
+                solver=solver,
+            )
+            solution = sim.solve([0, 3600])
+            t = solution.t
+            voltage = solution["Voltage [V]"].entries
+            t_mid = 0.5 * (t[:-1] + t[1:])
+            linear = 0.5 * (voltage[:-1] + voltage[1:])
+            truth = reference["Voltage [V]"](t_mid)
+            ratios.append(
+                np.max(np.abs(linear - truth) / (solver.atol + solver.rtol * truth))
+            )
+        assert ratios[0] > 10
+        assert ratios[1] < 3
+
+    def test_sensitivities(self):
+        parameter_values = pybamm.ParameterValues("Chen2020")
+        parameter_values["Current function [A]"] = "[input]"
+        inputs = {"Current function [A]": 5}
+
+        def solve(factor, t_interp=None, output_variables=None):
+            sim = pybamm.Simulation(
+                pybamm.lithium_ion.SPM(),
+                parameter_values=parameter_values,
+                solver=pybamm.IDAKLUSolver(
+                    options={"upsampling_factor": factor},
+                    output_variables=output_variables,
+                ),
+            )
+            return sim.solve(
+                [0, 3600],
+                inputs=inputs,
+                calculate_sensitivities=True,
+                t_interp=t_interp,
+            )
+
+        baseline = solve(1.0)
+        upsampled = solve(2.0)
+        assert upsampled.t.size > baseline.t.size
+
+        # The added points carry the integrator's own interpolated sensitivities
+        reference = solve(1.0, t_interp=upsampled.t)
+        outputs_only = solve(2.0, output_variables=["Voltage [V]"])
+        np.testing.assert_array_equal(outputs_only.t, upsampled.t)
+        for solution in [upsampled, outputs_only]:
+            np.testing.assert_allclose(
+                solution["Voltage [V]"].sensitivities["Current function [A]"],
+                reference["Voltage [V]"].sensitivities["Current function [A]"],
+                rtol=1e-10,
+                atol=1e-12,
+            )
+
+    def test_option_errors(self):
+        with pytest.raises(pybamm.SolverError, match=r"must be >= 1.0"):
+            pybamm.IDAKLUSolver(options={"upsampling_factor": 0.5})
+        with pytest.raises(pybamm.SolverError, match="mutually exclusive"):
+            pybamm.IDAKLUSolver(
+                options={
+                    "upsampling_factor": 2.0,
+                    "hermite_reduction_factor": 2.0,
+                }
+            )
+        model = pybamm.lithium_ion.SPM()
+        model.upsampling_variable = "Negative particle concentration [mol.m-3]"
+        solver = pybamm.IDAKLUSolver(options={"upsampling_factor": 2.0})
+        sim = pybamm.Simulation(model, solver=solver)
+        with pytest.raises(pybamm.SolverError, match="needs a scalar variable"):
+            sim.solve([0, 600])
+
+    def test_skips_algebraic_and_missing_variables(self):
+        experiment = pybamm.Experiment(
+            ["Charge at 1C until 4.1 V", "Hold at 4.1 V for 10 minutes"]
+        )
+        sim = pybamm.Simulation(
+            pybamm.lithium_ion.SPM(),
+            experiment=experiment,
+            solver=pybamm.IDAKLUSolver(options={"upsampling_factor": 2.0}),
+        )
+        sim.solve()
+        charge, hold = experiment.steps
+        solver = sim._get_built_experiment_solver(charge)
+        model = sim._get_built_experiment_model(charge)
+        assert solver._upsampling_fcn(model, {}).n_out() == 1
+        model.upsampling_variable = "Not a variable"
+        assert solver._upsampling_fcn(model, {}).n_out() == 0
+
+        # In a voltage hold the current is an algebraic state
+        solver = sim._get_built_experiment_solver(hold)
+        model = sim._get_built_experiment_model(hold)
+        assert model.upsampling_variable == "Current [A]"
+        assert solver._upsampling_fcn(model, {}).n_out() == 0
+
+    @pytest.mark.parametrize("mode", ["legacy", "unified"])
+    def test_experiment_sets_variable_at_build(self, mode):
+        experiment = pybamm.Experiment(
+            ["Discharge at 1C until 3.0 V", "Hold at 3.0 V for 10 minutes"]
+        )
+        solutions = []
+        for factor in [1.0, 2.0]:
+            sim = pybamm.Simulation(
+                pybamm.lithium_ion.SPM(),
+                experiment=experiment,
+                solver=pybamm.IDAKLUSolver(options={"upsampling_factor": factor}),
+                experiment_model_mode=mode,
+            )
+            solutions.append(sim.solve())
+        models = [sim._get_built_experiment_model(step) for step in experiment.steps]
+        if mode == "legacy":
+            assert [model.upsampling_variable for model in models] == [
+                "Voltage [V]",
+                "Current [A]",
+            ]
+        else:
+            # One model serves every step, so it keeps the default voltage
+            assert not hasattr(models[0], "upsampling_variable")
+        assert solutions[1].cycles[0].t.size > solutions[0].cycles[0].t.size

packages/pybammsolvers/src/pybammsolvers/idaklu.cpp

@@ -51,7 +51,8 @@ IDAKLUSolverGroup *create_casadi_solver_group(
   const std::vector<CasadiFunctions::BaseFunctionType*>& dvar_dp_fcns,
   py::dict py_opts,
   const CasadiFunctions::BaseFunctionType &alg_res,
-  const CasadiFunctions::BaseFunctionType &alg_jac)
+  const CasadiFunctions::BaseFunctionType &alg_jac,
+  const CasadiFunctions::BaseFunctionType &upsampling_fcn)
 {
   auto setup_opts = SetupOptions(py_opts);
   auto solver_opts = SolverOptions(py_opts);
@@ -95,7 +96,8 @@ IDAKLUSolverGroup *create_casadi_solver_group(
       setup_opts,
       serialized_fcns ? &*serialized_fcns : nullptr,
       alg_res,
-      alg_jac
+      alg_jac,
+      upsampling_fcn
     );
     solvers.emplace_back(
       std::unique_ptr<IDAKLUSolver>(
@@ -174,6 +176,7 @@ PYBIND11_MODULE(idaklu, m)
     py::arg("options"),
     py::arg("alg_res"),
     py::arg("alg_jac"),
+    py::arg("upsampling_fcn"),
     py::return_value_policy::take_ownership);
 
   m.def("observe", &observe,

packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp

@@ -66,6 +66,8 @@ class ExpressionSet
   Expression *events = nullptr;
   Expression *alg_res = nullptr;
   Expression *alg_jac = nullptr;
+  // g(t, y, p) -> the scalar upsampled output
+  Expression *upsampling_fcn = nullptr;
 
   std::vector<Expression*> var_fcns;
   std::vector<Expression*> dvar_dy_fcns;