fix(idaklu): hand IDAS each parameter's magnitude as pbar

CHANGELOG.md

@@ -25,9 +25,11 @@
 - The `integration` nox session no longer installs the `pydiffsol` extra on macOS Intel CI runners, where it has no working build. ([#5726](https://github.com/pybamm-team/PyBaMM/pull/5726))
 - The Read the Docs build uses Read the Docs' native `uv` support (`python.install` with `method: uv`) instead of installing `uv` from a GitHub release tarball through `asdf` in `build.jobs`. The tarball fetch failed the build whenever GitHub's release CDN returned a 5xx; `uv` now ships in the build image. The generated command, `uv sync --group docs --extra all`, is unchanged. ([#5727](https://github.com/pybamm-team/PyBaMM/pull/5727))
 - The memray memory benchmarks no longer fail at random. memray saw Python's small-object allocations only as whole 1 MiB pymalloc arenas, so a test's result jumped by 1 MiB whenever the heap left by earlier tests ran out of space: `test_discretise_memory` intermittently reported 1.0 MiB against its 50 KiB limit. The `benchmark-memory` session now traces every Python allocation, each memory benchmark module warms up under memray before its tests are measured, and the limits are recalibrated to the traced figures. ([#5777](https://github.com/pybamm-team/PyBaMM/pull/5777))
+- `IDAKLUSolver` now hands IDAS each differentiated parameter's magnitude as its sensitivity scale `pbar` (`IDASetSensParams`), where it left every scale at 1. IDAS holds `pbar * dy/dp` to the state tolerances, so the unscaled column of a parameter as small as `Chen2020`'s positive particle diffusivity (4e-15, `dy/dp` around 1e14) broke the corrector: a DFN solve for that sensitivity failed with `IDA_CONV_FAIL` at every `rtol = atol` from 1e-8 to 1e-12 and now completes, matching finite differences to 3e-7, and at the default tolerances it runs 4.7x faster. The sensitivity to a large parameter, such as a temperature, is now held to the state tolerances too, which costs up to 1.6x at the default tolerances. ([#5783](https://github.com/pybamm-team/PyBaMM/pull/5783))
 
 ## Breaking changes
 
+- `pybamm` now requires `pybammsolvers>=0.10.0`, because `IDAKLUSolver` uses IDAKLU binding features the 0.9 series does not have. Upgrade `pybammsolvers` together with `pybamm`. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782), [#5783](https://github.com/pybamm-team/PyBaMM/pull/5783))
 - Bumped the pinned CasADi version from 3.7.2 to 3.8.1. CasADi 3.8 ships `abi3` wheels, which cover every current and future CPython version, and changes how numpy functions dispatch on CasADi values (casadi#2959). PyBaMM pins the legacy dispatch behaviour via `casadi.GlobalOptions.setNumpyMode(-1)` at import, so expression-tree result types are unchanged. ([#5761](https://github.com/pybamm-team/PyBaMM/pull/5761))
 
 # [v26.8.0.0](https://github.com/pybamm-team/PyBaMM/tree/pybamm-v26.8.0.0) - 2026-08-13

packages/pybamm/pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "pybammsolvers>=0.9.0,<0.10.0",
+    "pybammsolvers>=0.10.0,<0.11.0",
     "black",
     "numpy>=2.0.0",
     "scipy>=1.11.4,!=1.18.0", # 1.18.0 breaks pickling/deepcopy of interpolators, fixed in 1.18.1 (scipy#25489)

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

@@ -17,6 +17,32 @@
 _UNSET = object()
 
 
+def _sensitivity_scales(inputs_dict: dict, sensitivity_names: list[str]) -> np.ndarray:
+    """IDAS ``pbar``: the magnitude of each differentiated parameter.
+
+    IDAS weights the scaled sensitivity ``pbar_i * dy/dp_i`` like a state, so
+    ``pbar_i = |p_i|`` weights each column independently of the units of
+    ``p_i``.
+
+    Parameters
+    ----------
+    inputs_dict : dict
+        Input values for one input set.
+    sensitivity_names : list of str
+        Differentiated parameter names, in the solver's column order.
+
+    Returns
+    -------
+    numpy.ndarray
+        ``|p|`` for each name in ``sensitivity_names``, or 0 for a value with no
+        entries.
+    """
+    return np.array(
+        [np.max(np.abs(inputs_dict[name]), initial=0.0) for name in sensitivity_names],
+        dtype=np.float64,
+    )
+
+
 def _flatten_inputs(inputs_dict):
     """Flatten ``{name: value}`` into a 1-D float array in dict-key order."""
     if not inputs_dict:
@@ -685,6 +711,14 @@ def _integrate(
         else:
             inputs = np.array([[]] * len(inputs_list))
 
+        sensitivity_names = self._setup["sensitivity_names"]
+        if sensitivity_names:
+            pbar = np.vstack(
+                [_sensitivity_scales(d, sensitivity_names) for d in inputs_list]
+            )
+        else:
+            pbar = np.empty((0, 0))
+
         # y0full is now a list with length = number of input sets
         y0full = np.vstack(model.y0full)
         ydot0full = np.vstack(model.ydot0full)
@@ -705,6 +739,7 @@ def _integrate(
                 ydot0full,
                 inputs,
                 logger=logger,
+                pbar=pbar,
             )
         except ValueError as e:
             # Return from None to replace the C++ runtime error

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

@@ -1989,3 +1989,148 @@ def test_solution_user_options_survive_pickle(self, tmp_path):
         np.testing.assert_allclose(
             loaded["2u"].entries, sol["2u"].entries, rtol=1e-12, atol=1e-12
         )
+
+
+class TestIDAKLUSensitivityScales:
+    """IDAS ``pbar``, the scale of each differentiated parameter."""
+
+    @staticmethod
+    def _rescalable_model():
+        """``du/dt = -(a / s) u``: at ``a = s``, ``u = exp(-t)`` whatever the
+        magnitude of ``s``, and ``a du/da = -t exp(-t)``."""
+        model = pybamm.BaseModel()
+        u = pybamm.Variable("u")
+        a = pybamm.InputParameter("a")
+        s = pybamm.InputParameter("s")
+        model.rhs = {u: -(a / s) * u}
+        model.initial_conditions = {u: 1}
+        model.variables = {"u": u}
+        pybamm.Discretisation().process_model(model)
+        return model
+
+    def test_scales_are_the_parameter_magnitudes_in_column_order(self):
+        # Dict order would have the right length too, so a mix-up would
+        # degrade the weighting silently instead of raising.
+        scales = pybamm.solvers.idaklu_solver._sensitivity_scales(
+            {"a": -3.0, "b": 4e-15, "c": 1.0}, ["b", "a"]
+        )
+        np.testing.assert_allclose(scales, [4e-15, 3.0])
+
+    @pytest.mark.parametrize(
+        ("value", "expected"),
+        [
+            (np.array([-2.0]), 2.0),
+            (np.array([[7.0]]), 7.0),
+            # The solver owns the zero clamp, so a raw 0 reaches it.
+            (0.0, 0.0),
+            (np.array([]), 0.0),
+        ],
+    )
+    def test_scale_of_each_input_form(self, value, expected):
+        scales = pybamm.solvers.idaklu_solver._sensitivity_scales({"a": value}, ["a"])
+        np.testing.assert_allclose(scales, [expected])
+
+    def test_a_zero_parameter_still_solves(self):
+        # IDAS rejects pbar = 0, so the solver has to clamp it to the unit scale.
+        model = pybamm.BaseModel()
+        u = pybamm.Variable("u")
+        a = pybamm.InputParameter("a")
+        model.rhs = {u: -u + a}
+        model.initial_conditions = {u: 1}
+        model.variables = {"u": u}
+        pybamm.Discretisation().process_model(model)
+
+        solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8)
+        sol = solver.solve(
+            model,
+            np.linspace(0, 1, 10),
+            inputs={"a": 0.0},
+            calculate_sensitivities=True,
+        )
+        # du/da = 1 - exp(-t) regardless of a, so the clamp must not skew it.
+        np.testing.assert_allclose(
+            np.asarray(sol["u"].sensitivities["a"]).ravel(),
+            1.0 - np.exp(-sol.t),
+            rtol=1e-5,
+            atol=1e-7,
+        )
+
+    def test_each_input_set_hands_over_its_own_row(self):
+        model = self._rescalable_model()
+        solver = pybamm.IDAKLUSolver()
+        solve_kwargs = {
+            "inputs": [{"a": 1e-14, "s": 1e-14}, {"a": 1e3, "s": 1e3}],
+            "calculate_sensitivities": ["a"],
+        }
+        # Solve once so the second solve reuses this group rather than rebuilding
+        # it, which would discard the spy.
+        solver.solve(model, [0, 1], **solve_kwargs)
+
+        seen = {}
+        original = solver._setup["solver"].solve
+
+        def spy(*args, **kwargs):
+            seen["pbar"] = np.asarray(kwargs["pbar"])
+            return original(*args, **kwargs)
+
+        solver._setup["solver"] = type("Spy", (), {"solve": staticmethod(spy)})()
+        solver.solve(model, [0, 1], **solve_kwargs)
+        np.testing.assert_allclose(seen["pbar"], [[1e-14], [1e3]])
+
+    def test_each_input_set_is_solved_at_its_own_scale(self):
+        # a = s is one problem in different units, so pbar = |a| hands IDAS the
+        # same scaled system; powers of two make that scaling exact.
+        scales = [2.0**-47, 2.0**10]
+        solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8)
+        solutions = solver.solve(
+            self._rescalable_model(),
+            [0, 3],
+            inputs=[{"a": scale, "s": scale} for scale in scales],
+            calculate_sensitivities=["a"],
+        )
+        tiny, large = (
+            scale * np.asarray(sol["u"].sensitivities["a"]).ravel()
+            for scale, sol in zip(scales, solutions, strict=True)
+        )
+        np.testing.assert_array_equal(solutions[0].t, solutions[1].t)
+        np.testing.assert_array_equal(tiny, large)
+        np.testing.assert_allclose(
+            tiny, -solutions[0].t * np.exp(-solutions[0].t), rtol=1e-5, atol=1e-7
+        )
+
+    def test_a_tight_dfn_diffusivity_gradient_matches_finite_differences(self):
+        # D_p is 4e-15, so dy/dD_p reaches ~1e14: at the unit scale no absolute
+        # tolerance can hold that column and the corrector fails to converge.
+        name = "Positive particle diffusivity [m2.s-1]"
+        parameter_values = pybamm.ParameterValues("Chen2020")
+        nominal = parameter_values[name]
+        parameter_values[name] = pybamm.InputParameter("D_p")
+        simulation = pybamm.Simulation(
+            pybamm.lithium_ion.DFN(),
+            parameter_values=parameter_values,
+            solver=pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10),
+        )
+        t_eval = [0, 3000]
+        t_interp = np.linspace(0, 3000, 31)
+        sol = simulation.solve(
+            t_eval,
+            inputs={"D_p": nominal},
+            calculate_sensitivities=["D_p"],
+            t_interp=t_interp,
+        )
+        gradient = np.asarray(sol["Voltage [V]"].sensitivities["D_p"]).ravel()
+
+        step = 1e-3 * nominal
+        plus, minus = (
+            simulation.solve(t_eval, inputs={"D_p": value}, t_interp=t_interp)[
+                "Voltage [V]"
+            ](sol.t)
+            for value in (nominal + step, nominal - step)
+        )
+        finite_difference = (plus - minus) / (2 * step)
+        np.testing.assert_allclose(
+            gradient,
+            finite_difference,
+            rtol=1e-4,
+            atol=1e-5 * np.abs(finite_difference).max(),
+        )

packages/pybammsolvers/pyproject.toml

@@ -18,7 +18,7 @@ requires-python = ">=3.10,<3.15"
 license-files = ["LICENSE"]
 readme = "README.md"
 # version.py stays the source of truth; a pre-commit hook mirrors it here.
-version = "0.9.1"
+version = "0.10.0"
 dependencies = [
     "casadi==3.8.1",
     "numpy>=2.0.0",

packages/pybammsolvers/src/pybammsolvers/idaklu.cpp

@@ -144,6 +144,8 @@ PYBIND11_MODULE(idaklu, m)
     py::arg("yp0"),
     py::arg("inputs"),
     py::arg("logger") = py::none(),
+    py::kw_only(),
+    py::arg("pbar") = np_array(),
     py::return_value_policy::take_ownership);
 
   m.def("create_casadi_solver_group", &create_casadi_solver_group,

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp

@@ -34,6 +34,7 @@ class IDAKLUSolver
     const sunrealtype *y0,
     const sunrealtype *yp0,
     const sunrealtype *inputs,
+    const sunrealtype *pbar,
     bool save_adaptive_steps,
     bool save_interp_steps
   ) = 0;

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp

@@ -9,7 +9,8 @@ std::vector<Solution> IDAKLUSolverGroup::solve(
     np_array y0_np,
     np_array yp0_np,
     np_array inputs,
-    py::object logger) {
+    py::object logger,
+    np_array pbar) {
   DEBUG("IDAKLUSolverGroup::solve");
 
   // If t_interp is empty, save all adaptive steps
@@ -99,9 +100,25 @@ std::vector<Solution> IDAKLUSolverGroup::solve(
   const std::size_t solves_per_thread = number_of_groups / m_solvers.size();
   const std::size_t remainder_solves = number_of_groups % m_solvers.size();
 
+  // pbar is optional: an empty array leaves IDAS at its unit default.
+  const bool has_pbar = pbar.size() > 0;
+  if (has_pbar) {
+    if (pbar.ndim() != 2)
+      throw std::domain_error("pbar has wrong number of dimensions. Expected 2 but got " + std::to_string(pbar.ndim()));
+    if (pbar.shape()[0] != number_of_groups)
+      throw std::domain_error(
+        "pbar has wrong number of rows. Expected " + std::to_string(number_of_groups) +
+        " but got " + std::to_string(pbar.shape()[0]));
+    if (pbar.shape()[1] != number_of_parameters)
+      throw std::domain_error(
+        "pbar has wrong number of cols. Expected " + std::to_string(number_of_parameters) +
+        " but got " + std::to_string(pbar.shape()[1]));
+  }
+
   const sunrealtype *y0 = y0_np.data();
   const sunrealtype *yp0 = yp0_np.data();
   const sunrealtype *inputs_data = inputs.data();
+  const sunrealtype *pbar_data = has_pbar ? pbar.data() : nullptr;
 
   std::vector<SolutionData> results(number_of_groups);
 
@@ -122,7 +139,8 @@ std::vector<Solution> IDAKLUSolverGroup::solve(
         const sunrealtype *y = y0 + index * y0_np.shape(1);
         const sunrealtype *yp = yp0 + index * yp0_np.shape(1);
         const sunrealtype *input = inputs_data + index * inputs.shape(1);
-        results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, save_adaptive_steps, save_interp_steps);
+        const sunrealtype *scales = pbar_data ? pbar_data + index * number_of_parameters : nullptr;
+        results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps);
       }
     } catch (py::error_already_set &) {
       #pragma omp critical
@@ -156,7 +174,8 @@ std::vector<Solution> IDAKLUSolverGroup::solve(
     const sunrealtype *y = y0 + index * y0_np.shape(1);
     const sunrealtype *yp = yp0 + index * yp0_np.shape(1);
     const sunrealtype *input = inputs_data + index * inputs.shape(1);
-    results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, save_adaptive_steps, save_interp_steps);
+    const sunrealtype *scales = pbar_data ? pbar_data + index * number_of_parameters : nullptr;
+    results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps);
   }
 
   flush_logs();

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp

@@ -37,7 +37,8 @@ class IDAKLUSolverGroup
     np_array y0_np,
     np_array yp0_np,
     np_array inputs,
-    py::object logger = py::none());
+    py::object logger = py::none(),
+    np_array pbar = np_array());
 
 
   private:

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp

@@ -3,6 +3,7 @@
 
 #include "IDAKLUSolver.hpp"
 #include "common.hpp"
+#include <cmath>
 #include <vector>
 #include <memory>  // For std::make_unique
 using std::vector;
@@ -188,6 +189,9 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver
   sunrealtype *yp_val_ = nullptr;
   vector<sunrealtype *> yS_val_;
   vector<sunrealtype *> ypS_val_;
+  // |p| per sensitivity parameter, so IDAS weights the scaled sensitivity
+  // pbar*yS like a state. Empty leaves IDAS at its pbar = 1 default.
+  std::vector<sunrealtype> sens_scales_;
 
   SUNContext sunctx;
 
@@ -223,6 +227,7 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver
     const sunrealtype *y0,
     const sunrealtype *yp0,
     const sunrealtype *inputs,
+    const sunrealtype *pbar,
     bool save_adaptive_steps,
     bool save_interp_steps
   ) override;
@@ -296,6 +301,14 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver
    */
   void ReinitializeIntegrator(const sunrealtype& t_val);
 
+  /**
+   * @brief Hand IDAS the per-parameter scales held in sens_scales_.
+   *
+   * Must run after every IDASensInit/IDASensReInit, both of which reset pbar
+   * to 1.0.
+   */
+  void ApplySensitivityScales();
+
   /**
    * @brief Set a consistent initialization for the system of equations.
    */
@@ -365,7 +378,8 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver
     const std::vector<sunrealtype> &t_eval,
     const sunrealtype *y0,
     const sunrealtype *yp0,
-    const sunrealtype *inputs
+    const sunrealtype *inputs,
+    const sunrealtype *pbar
   );
 
   /**