Add pybammsolvers nonlinear solver

CHANGELOG.md

@@ -1,5 +1,9 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Features
+
+- Added `NonlinearSolver` as the default nonlinear solver, which replaces `CasadiAlgebraicSolver`. `IDAKLUSolver` now computes the initial conditions in C++ by default. ([#5459](https://github.com/pybamm-team/PyBaMM/pull/5459))
+
 # [v26.4.1](https://github.com/pybamm-team/PyBaMM/tree/v26.4.1) - 2026-04-24
 
 ## Bug fixes

docs/source/user_guide/installation/index.rst

@@ -73,7 +73,7 @@ PyBaMM requires the following dependencies.
 =================================================================== ==========================
 Package                                                             Supported version(s)
 =================================================================== ==========================
-`PyBaMM solvers <https://github.com/pybamm-team/pybammsolvers>`__     >= 0.6.0, <0.7.0
+`PyBaMM solvers <https://github.com/pybamm-team/pybammsolvers>`__     >= 0.8.0, <0.9.0
 `NumPy <https://numpy.org>`__                                         Whatever recent versions work
 `SciPy <https://docs.scipy.org/doc/scipy/>`__                         Whatever recent versions work. >= 1.9.3
 `CasADi <https://web.casadi.org/docs/>`__                             Whatever recent versions work. >= 3.7.2

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "pybammsolvers>=0.6.0,<0.7.0",
+    "pybammsolvers>=0.8.0,<0.9.0",
     "black",
     "numpy",
     "scipy>=1.11.4",

src/pybamm/__init__.py

@@ -200,6 +200,7 @@
 
 from .solvers.idaklu_jax import IDAKLUJax
 from .solvers.idaklu_solver import IDAKLUSolver
+from .solvers.nonlinear_solver import NonlinearSolver
 
 # Experiments
 from .experiment.experiment import Experiment

src/pybamm/models/base_model.py

@@ -608,7 +608,7 @@ def default_spatial_methods(self):
     def default_solver(self):
         """Returns the default solver for the model, based on whether it is an ODE/DAE or algebraic model."""
         if len(self.rhs) == 0 and len(self.algebraic) != 0:
-            return pybamm.CasadiAlgebraicSolver()
+            return pybamm.NonlinearSolver()
         else:
             return pybamm.IDAKLUSolver()
 

src/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py

@@ -1353,7 +1353,7 @@ def _get_msmr_potential_model(parameter_values, param):
 def get_esoh_default_solver(tol: float = 1e-6) -> pybamm.CompositeSolver:
     return pybamm.CompositeSolver(
         [
-            pybamm.CasadiAlgebraicSolver(tol=tol, step_tol=0),
+            pybamm.NonlinearSolver(atol=tol, rtol=0, step_tol=0, max_backtracks=100),
             pybamm.AlgebraicSolver(tol=tol),
             pybamm.AlgebraicSolver(method="lsq", tol=tol),
             pybamm.AlgebraicSolver(method="minimize", tol=tol),

src/pybamm/solvers/__init__.py

@@ -8,6 +8,7 @@
     "dummy_solver",
     "idaklu_jax",
     "idaklu_solver",
+    "nonlinear_solver",
     "jax_bdf_solver",
     "jax_solver",
     "lrudict",

src/pybamm/solvers/algebraic_solver.py

@@ -56,8 +56,8 @@ def __init__(self, method="lm", tol=1e-6, extra_options=None):
         self._algebraic_solver = True
         pybamm.citations.register("Virtanen2020")
 
-    def set_up_root_solver(self, model, inputs_dict, t_eval):
-        """Create and return a rootfinder object. Not used for `pybamm.AlgebraicSolver`.
+    def get_root_solver(self, model, inputs_dict, t_eval):
+        """Return a rootfinder object. Not used for `pybamm.AlgebraicSolver`.
 
         Parameters
         ----------
@@ -67,8 +67,13 @@ def set_up_root_solver(self, model, inputs_dict, t_eval):
             Dictionary of inputs.
         t_eval : :class:`numpy.array`, size (k,)
             The times at which to compute the solution.
+
+        Returns
+        -------
+        None
+            Not used for `pybamm.AlgebraicSolver`.
         """
-        pass
+        return None
 
     @property
     def tol(self):

src/pybamm/solvers/base_solver.py

@@ -29,10 +29,10 @@ class BaseSolver:
         The absolute tolerance for the solver (default is 1e-6).
     root_method : str or pybamm algebraic solver class, optional
         The method to use to find initial conditions (for DAE solvers).
-        If a solver class, must be an algebraic solver class.
-        If "casadi",
-        the solver uses casadi's Newton rootfinding algorithm to find initial
-        conditions. Otherwise, the solver uses 'scipy.optimize.root' with method
+        Default is "nonlinear_solver", which uses a custom Newton solver for consistent
+        initial conditions (recommended). If "casadi", the solver uses
+        casadi's Newton rootfinding algorithm as a fallback.
+        Otherwise, the solver uses 'scipy.optimize.root' with method
         specified by 'root_method' (e.g. "lm", "hybr", ...)
     root_tol : float, optional
         The tolerance for the initial-condition solver (default is 1e-6).
@@ -65,11 +65,13 @@ def __init__(
         self.rtol = rtol
         self.atol = atol
         self.root_tol = root_tol
+        # Set ``on_failure`` before ``root_method`` so the setter can propagate
+        # the parent's failure policy down to the constructed root solver.
+        self.on_failure = on_failure or "error"
         self.root_method = root_method
         self.extrap_tol = extrap_tol or -1e-10
         self.output_variables = [] if output_variables is None else output_variables
         self.on_extrapolation = on_extrapolation or "warn"
-        self.on_failure = on_failure or "error"
         self._model_set_up = {}
 
         # Defaults, can be overwritten by specific solver
@@ -151,13 +153,24 @@ def on_failure(self, value):
             raise ValueError("on_failure must be 'warn', 'error', or 'ignore'")
         self._on_failure = value
 
+        # Keep root solver's failure policy in sync with the parent's
+        root_method = getattr(self, "_root_method", None)
+        if isinstance(root_method, pybamm.BaseSolver):
+            root_method.on_failure = value
+
     @property
     def root_method(self):
         return self._root_method
 
     @root_method.setter
     def root_method(self, method):
-        if method == "casadi":
+        if method == "nonlinear_solver":
+            # use the tighter of the two tolerances
+            atol = min(self.root_tol, self.atol)
+            method = pybamm.NonlinearSolver(
+                atol=atol, rtol=self.rtol, on_failure=self.on_failure
+            )
+        elif method == "casadi":
             method = pybamm.CasadiAlgebraicSolver(self.root_tol)
         elif isinstance(method, str):
             method = pybamm.AlgebraicSolver(method, self.root_tol)
@@ -294,7 +307,8 @@ def set_up(
         # Save CasADi functions for solvers that use CasADi
         # Note: when we pass to casadi the ode part of the problem must be in
         if isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) or isinstance(
-            self, pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver
+            self,
+            pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver,
         ):
             # can use DAE solver to solve model with algebraic equations only
             if len(model.rhs) > 0:
@@ -318,9 +332,6 @@ def set_up(
             model.casadi_sensitivities_rhs = jacp_rhs
             model.casadi_sensitivities_algebraic = jacp_algebraic
 
-        if getattr(self.root_method, "algebraic_solver", False):
-            self.root_method.set_up_root_solver(model, inputs[0], t_eval)
-
         # if output_variables specified then convert functions to casadi
         # expressions for evaluation within the respective solver
         self.computed_var_fcns = {}
@@ -968,7 +979,7 @@ def solve(
 
         # Check initial conditions don't violate events
         for y0, inpts in zip(model.y0_list, model_inputs_list, strict=True):
-            self._check_events_with_initialization(t_eval, model, y0, inpts)
+            self._check_event_violation_on_initialisation(t_eval, model, y0, inpts)
 
         # Process discontinuities
         t_eval_info = [
@@ -1066,9 +1077,11 @@ def solve(
             )
 
         # Raise error if solutions[0] only contains one timestep (except for algebraic
-        # solvers, where we may only expect one time in the solution)
+        # solvers, where we may only expect one time in the solution, or failure
+        # solutions where we may only have the initial state)
         if (
             self._algebraic_solver is False
+            and solutions[0].termination != "failure"
             and len(solutions[0].all_ts) == 1
             and len(solutions[0].all_ts[0]) == 1
         ):
@@ -1108,6 +1121,16 @@ def filter_discontinuities(t_discon: list, t_eval: list) -> np.ndarray:
         idx_end = np.searchsorted(t_discon_unique, t_eval[-1], side="left")
         return t_discon_unique[idx_start:idx_end]
 
+    def get_root_solver(self, model, inputs_dict, t_eval):
+        if not model.algebraic_root_solver:
+            model.algebraic_root_solver = self._set_up_root_solver(
+                model, inputs_dict, t_eval
+            )
+        return model.algebraic_root_solver
+
+    def _set_up_root_solver(self, model, inputs_dict, t_eval):
+        raise NotImplementedError
+
     def _get_discontinuity_start_end_indices(self, model, y0, inputs, t_eval):
         if self.supports_t_eval_discontinuities and model.t_discon_constant_symbols:
             pybamm.logger.verbose("Discontinuity events found for constant symbols")
@@ -1153,8 +1176,13 @@ def _get_discontinuity_start_end_indices(self, model, y0, inputs, t_eval):
 
         return start_indices, end_indices, t_eval
 
-    @staticmethod
-    def _check_events_with_initialization(t_eval, model, y0, inputs_dict):
+    def _check_event_violation_on_initialisation(self, *args, **kwargs):
+        self._check_event_violation(*args, **kwargs)
+
+    def _check_event_violation_post_solve(self, *args, **kwargs):
+        self._check_event_violation(*args, **kwargs)
+
+    def _check_event_violation(self, t_eval, model, y0, inputs_dict):
         num_terminate_events = len(model.terminate_events_eval)
         if num_terminate_events == 0:
             return
@@ -1459,7 +1487,9 @@ def step(
         self._set_consistent_initialization(model, t_start_shifted, [model_inputs])
 
         # Check consistent initialization doesn't violate events
-        self._check_events_with_initialization(t_eval, model, model.y0, model_inputs)
+        self._check_event_violation_on_initialisation(
+            t_eval, model, model.y0, model_inputs
+        )
 
         # Step
         pybamm.logger.verbose(f"Stepping for {t_start_shifted:.0f} < t < {t_end:.0f}")

src/pybamm/solvers/casadi_algebraic_solver.py

@@ -68,7 +68,7 @@ def step_tol(self):
     def step_tol(self, value):
         self._step_tol = value
 
-    def set_up_root_solver(self, model, inputs_dict, t_eval):
+    def _set_up_root_solver(self, model, inputs_dict, t_eval):
         """Create and return a CasADi rootfinder object. The parameter argument to the
         rootfinder is the concatenated time, differential states, and flattened inputs.
 
@@ -83,7 +83,7 @@ def set_up_root_solver(self, model, inputs_dict, t_eval):
 
         Returns
         -------
-        None
+        :class:`casadi.rootfinder`
             The rootfinder function is stored in the model as `algebraic_root_solver`.
         """
         pybamm.logger.info(f"Start building {self.name}")
@@ -137,7 +137,7 @@ def set_up_root_solver(self, model, inputs_dict, t_eval):
         constraints[model_alg_ub <= 0] = -1
 
         # Set up rootfinder
-        model.algebraic_root_solver = casadi.rootfinder(
+        out = casadi.rootfinder(
             "roots",
             "newton",
             dict(x=y_alg_sym, p=p_sym, g=alg),
@@ -150,6 +150,7 @@ def set_up_root_solver(self, model, inputs_dict, t_eval):
         )
 
         pybamm.logger.info(f"Finish building {self.name}")
+        return out
 
     def _integrate_single(self, model, t_eval, inputs_dict, y0):
         """
@@ -202,9 +203,7 @@ def _integrate_single(self, model, t_eval, inputs_dict, y0):
         # Set the inputs portion of the parameter vector
         p[1 + len_rhs :] = np.asarray(inputs).ravel()
 
-        if getattr(model, "algebraic_root_solver", None) is None:
-            self.set_up_root_solver(model, inputs_dict, t_eval)
-        roots = model.algebraic_root_solver
+        roots = self.get_root_solver(model, inputs_dict, t_eval)
 
         timer = pybamm.Timer()
         integration_time = 0

src/pybamm/solvers/casadi_solver.py

@@ -34,10 +34,10 @@ class CasadiSolver(pybamm.BaseSolver):
         The absolute tolerance for the solver (default is 1e-6).
     root_method : str or pybamm algebraic solver class, optional
         The method to use to find initial conditions (for DAE solvers).
-        If a solver class, must be an algebraic solver class.
-        If "casadi",
-        the solver uses casadi's Newton rootfinding algorithm to find initial
-        conditions. Otherwise, the solver uses 'scipy.optimize.root' with method
+        Default is "nonlinear_solver", which uses a custom Newton solver for consistent
+        initial conditions (recommended). If "casadi", the solver uses
+        casadi's Newton rootfinding algorithm as a fallback.
+        Otherwise, the solver uses 'scipy.optimize.root' with method
         specified by 'root_method' (e.g. "lm", "hybr", ...)
     root_tol : float, optional
         The tolerance for root-finding. Default is 1e-6.
@@ -82,7 +82,7 @@ def __init__(
         mode="safe",
         rtol=1e-6,
         atol=1e-6,
-        root_method="casadi",
+        root_method="nonlinear_solver",
         root_tol=1e-6,
         max_step_decrease_count=5,
         dt_max=None,
@@ -206,10 +206,14 @@ def _integrate_single(self, model, t_eval, inputs_dict, y0):
                 # create integrator once, without grid,
                 # to avoid having to create several times
                 self.create_integrator(model, y0, inputs)
-                # Initialize solution
+                # Initialize solution. Coerce y0 to a CasADi DM so that when
+                # subsequent CasADi integration segments are appended, the
+                # concatenated ``solution.y`` stays a CasADi type. When the
+                # root method is NonlinearSolver, y0 arrives as an ndarray.
+                y0_init = y0 if isinstance(y0, casadi.DM | casadi.MX) else casadi.DM(y0)
                 solution = pybamm.Solution(
                     np.array([t]),
-                    y0,
+                    y0_init,
                     model,
                     inputs_dict,
                 )