feat(idaklu): solver statistics, self-scheduled input-set sweeps and polymorphic Newton

CHANGELOG.md

@@ -8,6 +8,18 @@
 
 - Setting `Array.entries_string` or `Interpolant.entries_string`, and calling `StateVector.set_evaluation_array`, are in-place symbol updates and now emit `SymbolMutationDeprecationWarning`. Construct a new symbol instead. ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))
 
+## Features
+
+- `Solution.solver_statistics` reports the integrator's step, linear-solver-setup, nonlinear-iteration, nonlinear-failure and error-test-failure counts as a `pybamm.SolverStatistics` dataclass. `IDAKLUSolver` fills it for every solution, counting across `t_eval` breakpoints and separately for each input set, and the counts are summed when solutions are combined, so a stepped or experiment solution reports its total, including the cycles that `save_at_cycles` does not keep. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+- A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, whether it failed from the start or part-way through, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. Each `on_failure="warn"` warning names its set the same way. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+
+## Bug fixes
+
+- A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans, even one among numbers such as `[True, 1e-6]`, and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+- `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+- Discretising a model whose differential variables are all scalars or on zero-dimensional domains (such as a lumped current collector) no longer raises SciPy 1.18's `block_diag` `DeprecationWarning`, and its mass matrix stays a sparse matrix rather than becoming a sparse array once SciPy changes `block_diag`'s return type. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+- The CI test matrix now installs every dependency from `uv.lock`, as local nox sessions do, instead of resolving the latest releases, so only the prebuilt `pybammsolvers` wheel is installed outside the lock. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+
 ## Optimizations
 
 - Symbol mutation is no longer guarded by intercepting every attribute write, and `children` and `domain` return the stored immutable sequences instead of building a view on each access. Writes to a symbol's slots after construction are instead rejected by a static check of PyBaMM's code, tests, examples and documentation. Under the test suite's mutation guard, constructing an expression is twice as fast and parameterising it is 25% faster, which recovers most of the `test_parameterise` benchmark regression from [#5779](https://github.com/pybamm-team/PyBaMM/pull/5779). ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))

docs/source/api/solvers/solution.rst

@@ -4,4 +4,6 @@ Solutions
 .. autoclass:: pybamm.Solution
   :members:
 
+.. autoclass:: pybamm.SolverStatistics
+
 .. footbibliography::

noxfile.py

@@ -46,70 +46,58 @@ def is_macos_intel():
     return sys.platform == "darwin" and platform.machine() in ("x86_64", "i386")
 
 
-def install_locked(session, *, extras=None, groups=None, zoo=False, zoo_extras=None):
-    """Install pybamm and its dependencies into the session environment.
-
-    Two modes, selected by the ``PYBAMM_SOLVER_WHEELS`` environment variable:
-
-    * **Unset:** ``uv sync --frozen``, so ``pybammsolvers`` builds from source.
-    * **Set to a wheel directory (the CI matrix):** install this interpreter's
-      wheel by path, then pybamm with ``--no-sources``. The in-repo solver
-      version collides with the PyPI release, so nothing weaker picks the right
-      artifact, and Windows has no from-source build.
-
-    ``zoo=True`` also installs the zoo with ``zoo_extras``, which only the
-    prebuilt-wheel path needs — a workspace sync installs every member anyway.
+def install_locked(session, *, extras=None, groups=None, zoo_extras=None):
+    """Install every workspace member and its ``uv.lock`` dependencies.
+
+    ``uv sync --frozen`` installs the locked environment, building
+    ``pybammsolvers`` from source. When ``PYBAMM_SOLVER_WHEELS`` names a wheel
+    or a directory of wheels (the CI matrix), the sync skips ``pybammsolvers``
+    and this interpreter's prebuilt wheel is installed in its place, since
+    Windows has no from-source build.
     """
     env = {"UV_PROJECT_ENVIRONMENT": session.virtualenv.location}
-
     wheels = os.getenv("PYBAMM_SOLVER_WHEELS")
-    if wheels:
-        wheels_path = Path(wheels)
-        if wheels_path.is_dir():
-            # The per-OS artifact holds wheels for every Python version; pick the
-            # one whose interpreter tag matches this cell. The trailing dash in
-            # "*-cpXY-*" avoids matching free-threaded "cpXYt-" builds.
-            tag = f"cp{sys.version_info.major}{sys.version_info.minor}"
-            matches = sorted(wheels_path.glob(f"*-{tag}-*.whl"))
-            if not matches:
-                session.error(
-                    f"PYBAMM_SOLVER_WHEELS={wheels} contains no pybammsolvers "
-                    f"wheel for {tag}"
-                )
-            wheel = matches[0].as_posix()
-        else:
-            wheel = wheels_path.as_posix()
-        python_bin = os.path.join(
-            session.bin, "python.exe" if sys.platform == "win32" else "python"
-        )
-        extras_str = f"[{','.join(extras)}]" if extras else ""
-        zoo_extras_str = f"[{','.join(zoo_extras)}]" if zoo_extras else ""
-        cmd = [
-            "uv",
-            "pip",
-            "install",
-            "--python",
-            python_bin,
-            "--no-sources",
-            wheel,
-            "-e",
-            f"./packages/pybamm{extras_str}",
-        ]
-        if zoo:
-            cmd.extend(["-e", f"./packages/pybamm-model-zoo{zoo_extras_str}"])
-        for group in groups or []:
-            # Groups (dev, docs) are defined in the pybamm package, not the root.
-            cmd.extend(["--group", f"packages/pybamm/pyproject.toml:{group}"])
-        session.run(*cmd, env=env, external=True)
-        return
 
     cmd = ["uv", "sync", "--frozen"]
     for extra in [*(extras or []), *(zoo_extras or [])]:
         cmd.extend(["--extra", extra])
     for group in groups or []:
         cmd.extend(["--group", group])
+    if wheels:
+        cmd.extend(["--no-install-package", "pybammsolvers"])
     session.run(*cmd, env=env, external=True)
 
+    if not wheels:
+        return
+    wheels_path = Path(wheels)
+    if wheels_path.is_dir():
+        # The trailing dash keeps free-threaded "cpXYt-" wheels out of the match
+        tag = f"cp{sys.version_info.major}{sys.version_info.minor}"
+        matches = sorted(wheels_path.glob(f"*-{tag}-*.whl"))
+        if not matches:
+            session.error(
+                f"PYBAMM_SOLVER_WHEELS={wheels} contains no pybammsolvers "
+                f"wheel for {tag}"
+            )
+        wheel = matches[0].as_posix()
+    else:
+        wheel = wheels_path.as_posix()
+    python_bin = os.path.join(
+        session.bin, "python.exe" if sys.platform == "win32" else "python"
+    )
+    # The sync already installed the solver's locked dependencies
+    session.run(
+        "uv",
+        "pip",
+        "install",
+        "--python",
+        python_bin,
+        "--no-deps",
+        wheel,
+        env=env,
+        external=True,
+    )
+
 
 @nox.session(name="coverage", default=False)
 def run_coverage(session):
@@ -333,9 +321,7 @@ def build_docs(session):
 def install_zoo(session):
     """Install pybamm plus the zoo and every model's declared dependencies."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    install_locked(
-        session, extras=["all"], groups=["dev"], zoo=True, zoo_extras=["zoo-all"]
-    )
+    install_locked(session, extras=["all"], groups=["dev"], zoo_extras=["zoo-all"])
 
 
 def zoo_pytest(session, marker, *args, allow_empty=False):
@@ -400,7 +386,7 @@ def run_zoo_new(session):
     """Create a new model zoo entry from the template."""
     # Only the zoo itself and pybamm's version metadata are needed, so this skips
     # the extras and the dev group that install_zoo pulls in.
-    install_locked(session, zoo=True)
+    install_locked(session)
     session.run("python", f"{ZOO_TESTS}/scripts/new_model.py", *session.posargs)
 
 

packages/pybamm/pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "pybammsolvers>=0.10.0,<0.11.0",
+    "pybammsolvers>=0.11.0,<0.12.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/__init__.py

@@ -204,6 +204,7 @@
 
 # Solver classes
 from .solvers.solution import (
+    SolverStatistics,
     SolutionBase,
     Solution,
     EISSolution,

packages/pybamm/src/pybamm/discretisations/discretisation.py

@@ -769,7 +769,10 @@ def create_mass_matrix(self, model):
         N_alg = len(model.algebraic)
 
         if N_rhs > 0 or N_alg > 0:
-            mass_matrix = pybamm.Matrix(block_diag(mass_list, format="csr"))
+            # If every block were dense, block_diag would return a sparse array,
+            # not a matrix, and a 0D domain's mass matrix is dense
+            blocks = [csr_matrix(mass) for mass in mass_list]
+            mass_matrix = pybamm.Matrix(block_diag(blocks, format="csr"))
         else:
             mass_matrix = None
 

packages/pybamm/src/pybamm/simulation/simulation.py

@@ -987,6 +987,9 @@ def solve(
         # Track running termination (real cycles + skipped steps), applied to the
         # folded solution below to match the old left-fold's final termination.
         last_termination = None
+        # Solver work of the cycles left out of the fold, totalled so their
+        # solutions can be freed; None once one of them has no statistics
+        discarded_cycle_statistics = pybamm.SolverStatistics()
 
         for cycle_num, cycle_length in enumerate(
             cycle_lengths,
@@ -1172,6 +1175,16 @@ def solve(
                 cross_cycle_segments.append(cycle_solution)
                 if not isinstance(cycle_solution, pybamm.EmptySolution):
                     last_termination = cycle_solution.termination
+            elif (
+                isinstance(cycle_solution, pybamm.Solution)
+                and discarded_cycle_statistics is not None
+            ):
+                statistics = cycle_solution.solver_statistics
+                discarded_cycle_statistics = (
+                    None
+                    if statistics is None
+                    else discarded_cycle_statistics + statistics
+                )
 
             if steps:
                 if all(
@@ -1242,6 +1255,13 @@ def solve(
             self._solution = pybamm.Solution.from_sub_solutions(cross_cycle_segments)
             if last_termination is not None:
                 self._solution.termination = last_termination
+            if isinstance(self._solution, pybamm.Solution):
+                statistics = self._solution.solver_statistics
+                self._solution.solver_statistics = (
+                    None
+                    if statistics is None or discarded_cycle_statistics is None
+                    else statistics + discarded_cycle_statistics
+                )
 
         if self._solution is not None and len(all_cycle_solutions) > 0:
             self._solution.cycles = all_cycle_solutions