From 3a72f044b726c52a38e6456f35b35de9558cb763 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:36:58 -0400 Subject: [PATCH 01/12] feat: add pybamm.Brent, a bracketed scalar rootfind A `Brent` node solves `residual == 0` for its unknown over a bracket inside the CasADi graph rather than in Python: it converts to a `rootfinder` using a native "brent" plugin registered by pybammsolvers, so an expression containing one costs no extra Python frames per evaluation and survives codegen to C. Brent needs only a sign change over the bracket, so it converges where a Newton iteration stalls or leaves the domain, and the answer cannot leave the bracket. Derivatives come from CasADi's implicit function theorem. The plugin is built out of tree, which needs four things from CMake that an in-tree CasADi plugin gets for free: the internal headers staged from a pinned sdist rather than vendored, since CasADi does not install them and they are LGPL; an ABI check reading the linked version from CasADi's own config.h; the compile flags CasADi itself used, one of which adds a static mutex to Rootfinder; and the iteration stringified so codegen emits the text it compiles. brent.hpp says so at the top. Also fixes an unrelated test bug found on the way: test_pybamm_import unloaded every pybamm module from sys.modules and never restored them, leaving a second set of classes behind so `isinstance` failed across the boundary. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 1 + docs/source/api/expression_tree/brent.rst | 8 + docs/source/api/expression_tree/index.rst | 1 + packages/pybamm/src/pybamm/__init__.py | 1 + .../pybamm/discretisations/discretisation.py | 6 +- .../src/pybamm/expression_tree/brent.py | 291 ++++++++++++++++ .../parameters/parameter_substitutor.py | 1 + packages/pybamm/tests/strategies/symbols.py | 33 ++ .../unit/test_expression_tree/test_brent.py | 220 ++++++++++++ packages/pybamm/tests/unit/test_util.py | 53 +-- packages/pybammsolvers/CMakeLists.txt | 193 +++++++++++ .../src/pybammsolvers/idaklu.cpp | 5 + .../src/pybammsolvers/idaklu_source/brent.cpp | 254 ++++++++++++++ .../src/pybammsolvers/idaklu_source/brent.hpp | 122 +++++++ .../idaklu_source/brent_impl.hpp | 89 +++++ .../tests/test_brent_rootfinder.py | 320 ++++++++++++++++++ 16 files changed, 1574 insertions(+), 24 deletions(-) create mode 100644 docs/source/api/expression_tree/brent.rst create mode 100644 packages/pybamm/src/pybamm/expression_tree/brent.py create mode 100644 packages/pybamm/tests/unit/test_expression_tree/test_brent.py create mode 100644 packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp create mode 100644 packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp create mode 100644 packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp create mode 100644 packages/pybammsolvers/tests/test_brent_rootfinder.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a17810776c..b4d96be756 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Features +- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. ([#TBD](https://github.com/pybamm-team/PyBaMM/pull/TBD)) - Added unstructured mesh support (`UnstructuredSubMesh`, generators, and interface coupling) for arbitrary 2D/3D domains. Hexahedra must have planar faces (warped hexes raise a `GeometryError`), and `UserSuppliedUnstructuredMesh` accepts tetrahedral, triangular, and quadrilateral cells only. ([#5687](https://github.com/pybamm-team/PyBaMM/pull/5687)) - Generalised `VectorField` to N components and added `Component`/`Norm` operators for multi-dimensional vector fields. ([#5686](https://github.com/pybamm-team/PyBaMM/pull/5686)) - Removed the left sidebar from the documentation home page for a cleaner landing experience. ([#5699](https://github.com/pybamm-team/PyBaMM/pull/5699)) diff --git a/docs/source/api/expression_tree/brent.rst b/docs/source/api/expression_tree/brent.rst new file mode 100644 index 0000000000..05cf147a23 --- /dev/null +++ b/docs/source/api/expression_tree/brent.rst @@ -0,0 +1,8 @@ +Brent +===== + +.. autoclass:: pybamm.Brent + :members: + +.. autoclass:: pybamm.BrentUnknown + :members: diff --git a/docs/source/api/expression_tree/index.rst b/docs/source/api/expression_tree/index.rst index 0a6f3d757c..de5e5abfbb 100644 --- a/docs/source/api/expression_tree/index.rst +++ b/docs/source/api/expression_tree/index.rst @@ -19,4 +19,5 @@ Expression Tree functions input_parameter interpolant + brent operations/index diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index 8370ce6f58..02cf7665d5 100644 --- a/packages/pybamm/src/pybamm/__init__.py +++ b/packages/pybamm/src/pybamm/__init__.py @@ -32,6 +32,7 @@ from .expression_tree.broadcasts import * from .expression_tree.functions import * from .expression_tree.conditional import Conditional +from .expression_tree.brent import Brent, BrentUnknown from .expression_tree.interpolant import Interpolant from .expression_tree.discrete_time_sum import * from .expression_tree.input_parameter import InputParameter diff --git a/packages/pybamm/src/pybamm/discretisations/discretisation.py b/packages/pybamm/src/pybamm/discretisations/discretisation.py index cdb404505d..5fb30ab815 100644 --- a/packages/pybamm/src/pybamm/discretisations/discretisation.py +++ b/packages/pybamm/src/pybamm/discretisations/discretisation.py @@ -1148,7 +1148,11 @@ def _process_symbol(self, symbol): else: return symbol.create_copy(new_children=[disc_child]) - elif isinstance(symbol, (pybamm.Function, pybamm.Conditional)): + elif isinstance(symbol, pybamm.BrentUnknown): + # bound by its Brent, so it has no state-vector slice to resolve + return symbol.create_copy() + + elif isinstance(symbol, (pybamm.Function, pybamm.Conditional, pybamm.Brent)): disc_children = [self.process_symbol(child) for child in symbol.children] return symbol.create_copy(disc_children) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py new file mode 100644 index 0000000000..9d2ace2f5e --- /dev/null +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -0,0 +1,291 @@ +# +# A bracketed scalar rootfind as an expression tree node +# +from __future__ import annotations + +import casadi +import numpy as np + +import pybamm + + +class BrentUnknown(pybamm.Symbol): + """ + The scalar a :class:`Brent` solves for. + + Bound by the rootfinder, not by the model, so it is deliberately not a + :class:`pybamm.Variable`: the checks that enumerate model states must not count it + as one, or a ``Brent`` inside ``model.rhs`` or ``model.algebraic`` looks like an + extra unknown with no equation. + + Two unknowns of the same name are the same unknown, as for any other symbol. Give + each one its own name: a ``Brent`` nested inside another whose unknown shares its + name shadows the outer binding, and CasADi rejects the oracle it builds. + + Parameters + ---------- + name : str + Name of the node. Must be unique among the unknowns it is nested with. + """ + + def _evaluate_for_shape(self): + # a scalar, but shaped like every other column-vector node so that the + # broadcasting helpers can read shape[1] + return np.nan * np.ones((1, 1)) + + def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): + """See :meth:`pybamm.Symbol._to_casadi()`.""" + try: + return inputs[self.name] + except KeyError: + raise pybamm.ModelError( + f"'{self}' only has a value inside the Brent that solves for it" + ) from None + + def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None): + # set by the enclosing Brent while it iterates + return self._value + + _value = np.nan * np.ones((1, 1)) + + +def _nodes_reading(root: pybamm.Symbol, unknown: pybamm.Symbol) -> set: + """The symbols in ``root``'s graph whose subtree contains ``unknown``. + + Visits each node once; ``pre_order`` re-yields shared nodes once per path. + """ + reads: set = set() + seen: set = set() + stack = [(root, False)] + while stack: + node, expanded = stack.pop() + if expanded: + if node == unknown or any(child in reads for child in node.children): + reads.add(node) + elif node not in seen: + seen.add(node) + stack.append((node, True)) + stack.extend((child, False) for child in node.children) + return reads + + +class _OracleCache(dict): + """Conversion cache that keeps a :class:`Brent`'s own binding out of the shared one. + + Keys in ``local`` are held here and discarded with the oracle; the rest is written + through to ``shared``, so the graph a rootfind shares with its surroundings is + converted once. + """ + + def __init__(self, shared: dict, local: set): + super().__init__() + self._shared = shared + self._local = local + + def get(self, key, default=None): + value = super().get(key) + return self._shared.get(key, default) if value is None else value + + def __setitem__(self, key, value): + if key in self._local: + super().__setitem__(key, value) + else: + self._shared[key] = value + + +class Brent(pybamm.Symbol): + """ + Solve ``residual == 0`` for ``unknown`` within ``bounds``, by Brent's method. + + Every argument is an expression, so the bounds and anything inside ``residual`` + may be a :class:`pybamm.InputParameter` or any other symbol. Nothing is solved in + Python: the node converts to a CasADi ``rootfinder`` using the native ``brent`` + plugin registered by ``pybammsolvers``, so the whole solve runs inside the CasADi + graph. + + Brent needs only a sign change over the bounds, so it converges where a Newton + iteration stalls, and the answer cannot leave them. Derivatives come from CasADi's + implicit function theorem, exactly. + + Parameters + ---------- + residual : :class:`pybamm.Symbol` + The expression to drive to zero. Must contain ``unknown``. To invert ``f`` at + a target, pass ``f - target``. + unknown : :class:`pybamm.BrentUnknown` + The value being solved for. Must appear in ``residual`` and nowhere else in + the surrounding expression. + bounds : tuple + ``(lo, hi)``, the bracket to search. Either may be an expression. + abstol : float, optional + Absolute tolerance on the unknown. Unlike ``bounds``, a plain number, fixed + when the node is built. + max_iter : int, optional + Iteration cap. A plain number, as ``abstol``. + name : str, optional + Name of the node. + + Examples + -------- + .. code-block:: python + + # invert an open-circuit potential at a given voltage + sto = pybamm.BrentUnknown("stoichiometry") + node = pybamm.Brent(param.n.prim.U(sto, T) - voltage, sto, (0, 1)) + """ + + def __init__( + self, + residual: pybamm.Symbol, + unknown: BrentUnknown, + bounds: tuple, + *, + abstol: float = 1e-14, + max_iter: int = 100, + name: str = "brent", + ): + if not isinstance(residual, pybamm.Symbol): + raise TypeError( + f"residual must be a pybamm.Symbol, got {type(residual).__name__}" + ) + if not isinstance(unknown, BrentUnknown): + raise TypeError( + f"unknown must be a pybamm.BrentUnknown, got {type(unknown).__name__}" + ) + if not any(node == unknown for node in residual.pre_order()): + raise pybamm.ModelError(f"'{unknown}' does not appear in '{residual}'") + if len(bounds) != 2: + raise pybamm.ModelError(f"bounds must be a (lo, hi) pair, got {bounds}") + self.abstol = abstol + self.max_iter = max_iter + # the unknown is a child so that it survives copying and serialisation with + # the rest of the tree; it also appears inside the residual + children = [ + residual, + unknown, + *(pybamm.convert_to_symbol(bound) for bound in bounds), + ] + super().__init__(name, children=children) + + def set_id(self): + # the conversion cache is keyed on the id, so two Brents differing only in + # tolerance must not be served each other's rootfinder + super().set_id() + self._id = hash((self._id, self.abstol, self.max_iter)) + + @property + def residual(self) -> pybamm.Symbol: + return self.children[0] + + @property + def unknown(self) -> BrentUnknown: + return self.children[1] + + @property + def bounds(self) -> tuple[pybamm.Symbol, pybamm.Symbol]: + return tuple(self.children[2:]) + + def create_copy(self, new_children=None, perform_simplifications=True): + residual, unknown, lo, hi = self._children_for_copying(new_children) + return Brent( + residual, + unknown, + (lo, hi), + abstol=self.abstol, + max_iter=self.max_iter, + name=self.name, + ) + + def _evaluate_for_shape(self): + return pybamm.evaluate_for_shape_using_domain(self.domains) + + def to_json(self): + json_dict = super().to_json() + json_dict.update({"abstol": self.abstol, "max_iter": self.max_iter}) + return json_dict + + @classmethod + def _from_json(cls, snippet: dict): + residual, unknown, lo, hi = snippet["children"] + return cls( + residual, + unknown, + (lo, hi), + abstol=snippet["abstol"], + max_iter=snippet["max_iter"], + name=snippet["name"], + ) + + def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): + unknown = casadi.MX.sym(f"brent_unknown_{abs(self.id)}") + cache = _OracleCache( + casadi_symbols, _nodes_reading(self.residual, self.unknown) + ) + equation = self.children[0]._to_casadi_inner( + t, y, y_dot, {**inputs, self.unknown.name: unknown}, cache + ) + lo, hi = ( + bound._to_casadi_inner(t, y, y_dot, inputs, casadi_symbols) + for bound in self.bounds + ) + + # Pass only the symbols the residual actually reads. Handing the oracle the + # whole state vector would copy it into the solve on every evaluation. + free = [s for s in casadi.symvar(equation) if not casadi.is_equal(s, unknown)] + # the plugin takes the bracket from inputs 1 and 2, which the residual itself + # does not read, so they are declared here and left unused + lo_sym, hi_sym = casadi.MX.sym("lo"), casadi.MX.sym("hi") + oracle = casadi.Function( + f"brent_oracle_{abs(self.id)}", [unknown, lo_sym, hi_sym, *free], [equation] + ) + solver = casadi.rootfinder( + f"brent_{abs(self.id)}", + "brent", + oracle, + {"abstol": self.abstol, "max_iter": self.max_iter}, + ) + # the guess is required by the rootfinder interface and ignored by a bracketed + # method, so either end of the bracket does + return solver(lo, lo, hi, *free) + + def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None): + """Solve with :func:`scipy.optimize.brentq`, for the non-CasADi path.""" + from scipy.optimize import brentq + + def scalar(value): + return np.asarray(value, dtype=float).reshape(-1)[0] + + lo, hi = (scalar(b.evaluate(t, y, y_dot, inputs)) for b in self.bounds) + unknown = self.unknown + if not isinstance(unknown, BrentUnknown): + raise TypeError( + "NumPy evaluation needs a pybamm.BrentUnknown as the unknown, got " + f"{type(unknown).__name__}" + ) + + def g(value): + unknown._value = np.array([[value]]) + return scalar(self.residual.evaluate(t, y, y_dot, inputs)) + + # a shape probe carries NaN through the bounds, so there is nothing to solve + if not (np.isfinite(lo) and np.isfinite(hi)): + return np.nan * np.ones((1, 1)) + + try: + root = brentq(g, lo, hi, xtol=self.abstol, maxiter=self.max_iter) + except (ValueError, RuntimeError): + # a probe walks the tree before the parameters are in place + return np.nan * np.ones((1, 1)) + finally: + unknown._value = np.nan * np.ones((1, 1)) + return np.array([[root]]) + + def diff(self, variable): + raise NotImplementedError( + "Brent has no symbolic derivative; use convert_to_format='casadi'." + ) + + def _jac(self, variable): + raise NotImplementedError( + "Brent has no symbolic jacobian; use convert_to_format='casadi'." + ) diff --git a/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py b/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py index ab88e133f7..f9df3b70d0 100644 --- a/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py +++ b/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py @@ -256,6 +256,7 @@ def _process_symbol(self, symbol: pybamm.Symbol) -> pybamm.Symbol: pybamm.Concatenation, pybamm.BinaryOperator, pybamm.Conditional, + pybamm.Brent, ), ): new_children = [self.process_symbol(child) for child in symbol.children] diff --git a/packages/pybamm/tests/strategies/symbols.py b/packages/pybamm/tests/strategies/symbols.py index 4ad7a89627..b57ee6037f 100644 --- a/packages/pybamm/tests/strategies/symbols.py +++ b/packages/pybamm/tests/strategies/symbols.py @@ -443,6 +443,34 @@ def _reg_power_branch( ) +def _brent_branch( + child_strategy: st.SearchStrategy[pybamm.Symbol], +) -> st.SearchStrategy[pybamm.Brent]: + """Brent over ``residual(unknown) == 0``, bracketed by two further expressions.""" + + def make_brent(parts): + residual, lo, hi, abstol, max_iter, tag = parts + # nested Brents must not share an unknown, and a drawn tag replays identically + unknown = pybamm.BrentUnknown(f"brent unknown {tag}") + return pybamm.Brent( + # subtraction so the unknown cannot be simplified away, as `0 * x` would + unknown - residual, + unknown, + (lo, hi), + abstol=abstol, + max_iter=max_iter, + ) + + return st.tuples( + child_strategy, + child_strategy, + child_strategy, + st.floats(min_value=1e-16, max_value=1e-6, allow_nan=False), + st.integers(min_value=1, max_value=200), + st.integers(min_value=0, max_value=2**40), + ).map(make_brent) + + def _interpolant_branch( _child_strategy: st.SearchStrategy[pybamm.Symbol], ) -> st.SearchStrategy[pybamm.Interpolant]: @@ -870,6 +898,9 @@ def _vector_branch( pybamm.Parameter: lambda _children: parameter_strategy(), pybamm.Time: lambda _children: time_strategy(), pybamm.InputParameter: lambda _children: input_parameter_strategy(), + pybamm.BrentUnknown: lambda _children: st.integers( + min_value=0, max_value=2**40 + ).map(lambda tag: pybamm.BrentUnknown(f"brent unknown {tag}")), pybamm.Negate: lambda children: _unary_branch(children, pybamm.Negate), pybamm.AbsoluteValue: lambda children: _unary_branch( children, pybamm.AbsoluteValue @@ -938,6 +969,7 @@ def _vector_branch( pybamm.RegPower: _reg_power_branch, # data-bearing leaves pybamm.Interpolant: _interpolant_branch, + pybamm.Brent: _brent_branch, ExpressionFunctionParameter: _expression_function_parameter_branch, # n-ary / complex branch strategies pybamm.Conditional: _conditional_branch, @@ -1065,6 +1097,7 @@ def _vector_branch( pybamm.Parameter, pybamm.Time, pybamm.InputParameter, + pybamm.BrentUnknown, } ) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py new file mode 100644 index 0000000000..c1bdc50f41 --- /dev/null +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -0,0 +1,220 @@ +# +# Tests for the Brent expression tree node +# + +import sys + +import casadi +import numpy as np +import pytest +from scipy.optimize import brentq + +import pybamm +from pybamm.expression_tree.operations.serialise import ( + convert_symbol_from_json, + convert_symbol_to_json, +) + + +def _evaluate(symbol, **inputs): + """Evaluate a Brent expression over casadi symbols for each named input.""" + symbols = {name: casadi.MX.sym(name) for name in inputs} + expression = symbol.to_casadi(inputs=symbols) + function = casadi.Function("f", list(symbols.values()), [expression]) + return float(function(*[inputs[name] for name in symbols])) + + +class TestBrent: + def test_solves_a_scalar_equation(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(pybamm.exp(x) + x - 2.0, x, (-5.0, 5.0)) + got = float(casadi.evalf(node.to_casadi(inputs={}))) + want = brentq(lambda v: np.exp(v) + v - 2.0, -5.0, 5.0, xtol=2e-12) + assert got == pytest.approx(want, abs=1e-12) + + def test_target_may_be_an_input_parameter(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) + assert _evaluate(node, target=9.0) == pytest.approx(3.0, abs=1e-12) + assert _evaluate(node, target=4.0) == pytest.approx(2.0, abs=1e-12) + + def test_bracket_may_be_input_parameters(self): + # x^2 = 6 has roots at +-sqrt(6); the bracket selects one, at solve time + x = pybamm.BrentUnknown("x") + node = pybamm.Brent( + x * x - 6.0, x, (pybamm.InputParameter("lo"), pybamm.InputParameter("hi")) + ) + assert _evaluate(node, lo=0.0, hi=10.0) == pytest.approx(np.sqrt(6), abs=1e-12) + assert _evaluate(node, lo=-10.0, hi=0.0) == pytest.approx( + -np.sqrt(6), abs=1e-12 + ) + + def test_the_expression_may_contain_input_parameters(self): + # a x^2 = 9 has positive root 3 / sqrt(a) + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(pybamm.InputParameter("a") * x * x - 9.0, x, (0.0, 10.0)) + for a in (1.0, 4.0, 9.0): + assert _evaluate(node, a=a) == pytest.approx(3.0 / np.sqrt(a), abs=1e-12) + + def test_every_argument_may_be_an_input_parameter_at_once(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent( + pybamm.InputParameter("a") * x * x - pybamm.InputParameter("target"), + x, + (pybamm.InputParameter("lo"), pybamm.InputParameter("hi")), + ) + got = _evaluate(node, a=4.0, target=9.0, lo=0.0, hi=10.0) + assert got == pytest.approx(1.5, abs=1e-12) + got = _evaluate(node, a=4.0, target=9.0, lo=-10.0, hi=0.0) + assert got == pytest.approx(-1.5, abs=1e-12) + + def test_solves_over_the_state_vector(self): + state = pybamm.StateVector(slice(0, 1)) + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * state - 6.0, x, (0.0, 10.0)) + y = casadi.MX.sym("y", 1) + expression = casadi.Function("f", [y], [node.to_casadi(y=y, inputs={})]) + assert float(expression(2.0)) == pytest.approx(3.0, abs=1e-12) + assert float(expression(3.0)) == pytest.approx(2.0, abs=1e-12) + + def test_derivative_is_exact(self): + # x = sqrt(target), so dx/d(target) = 1 / (2 sqrt(target)) + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) + symbol = casadi.MX.sym("target") + root = node.to_casadi(inputs={"target": symbol}) + derivative = casadi.Function("J", [symbol], [casadi.jacobian(root, symbol)]) + assert float(derivative(9.0)) == pytest.approx(1 / 6, rel=1e-12) + + def test_composes_into_a_larger_expression(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + got = float(casadi.evalf((3 * node + pybamm.Scalar(1)).to_casadi(inputs={}))) + assert got == pytest.approx(10.0, abs=1e-12) + + def test_nests(self): + # the inner solve gives sqrt(16) = 4, so the outer gives sqrt(4) = 2 + inner_x = pybamm.BrentUnknown("inner") + inner = pybamm.Brent(inner_x * inner_x - 16.0, inner_x, (0.0, 10.0)) + outer_x = pybamm.BrentUnknown("outer") + outer = pybamm.Brent(outer_x * outer_x - inner, outer_x, (0.0, 10.0)) + assert float(casadi.evalf(outer.to_casadi(inputs={}))) == pytest.approx(2.0) + + def test_evaluating_does_not_re_enter_python(self): + # the whole solve runs in the CasADi graph, so a Brent node must cost no more + # python frames per evaluation than the same expression without one + state = pybamm.StateVector(slice(3, 4)) + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(pybamm.exp(x) + x * state - 2.0, x, (-5.0, 5.0)) + y = casadi.MX.sym("y", 500) + with_brent = casadi.Function("a", [y], [3 * node.to_casadi(y=y, inputs={}) + 1]) + without = casadi.Function("b", [y], [3 * casadi.exp(y[3]) + 1]) + values = np.zeros(500) + values[3] = 1.5 + + def count_frames(function): + calls = 0 + + def profile(frame, event, arg): + nonlocal calls + if event == "call": + calls += 1 + + sys.setprofile(profile) + try: + for _ in range(50): + function(values) + finally: + sys.setprofile(None) + return calls + + assert count_frames(with_brent) == count_frames(without) + + def test_the_oracle_only_reads_what_the_residual_needs(self): + # a residual that ignores time must not drag time into the solve + state = pybamm.StateVector(slice(0, 1)) + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * state - 6.0, x, (0.0, 10.0)) + t, y = casadi.MX.sym("t"), casadi.MX.sym("y", 1) + names = [s.name() for s in casadi.symvar(node.to_casadi(t=t, y=y, inputs={}))] + assert names == ["y"] + + def test_no_sign_change_fails_rather_than_guessing(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x + 1 - 0.0, x, (0.0, 1.0)) + with pytest.raises(RuntimeError, match="rootfinder process failed"): + casadi.evalf(node.to_casadi(inputs={})) + + def test_children_and_copy(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + assert len(node.children) == 4 + copy = node.create_copy() + assert copy.name == node.name + assert float(casadi.evalf(copy.to_casadi(inputs={}))) == pytest.approx(3.0) + + def test_errors(self): + x = pybamm.BrentUnknown("x") + with pytest.raises(TypeError, match=r"unknown must be a pybamm\.BrentUnknown"): + pybamm.Brent(x * x - 9.0, 1.0, (0, 1)) + with pytest.raises(TypeError, match=r"residual must be a pybamm\.Symbol"): + pybamm.Brent(1.0, x, (0, 1)) + with pytest.raises(pybamm.ModelError, match="does not appear in"): + pybamm.Brent(pybamm.Scalar(2) * pybamm.t - 9.0, x, (0, 1)) + with pytest.raises(pybamm.ModelError, match="bounds must be a"): + pybamm.Brent(x * x - 9.0, x, (0, 1, 2)) + + node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + with pytest.raises(NotImplementedError, match="no symbolic derivative"): + node.diff(pybamm.t) + with pytest.raises(NotImplementedError, match="no symbolic jacobian"): + node._jac(pybamm.t) + + def test_round_trips_through_json(self): + x = pybamm.BrentUnknown("x") + node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0), abstol=1e-12, max_iter=42) + rebuilt = convert_symbol_from_json(convert_symbol_to_json(node)) + assert rebuilt.abstol == 1e-12 + assert rebuilt.max_iter == 42 + assert rebuilt.unknown == node.unknown + # the two references in `x * x` and the explicit child are one unknown + assert ( + len({s.name for s in rebuilt.pre_order() if type(s) is pybamm.BrentUnknown}) + == 1 + ) + assert float(casadi.evalf(rebuilt.to_casadi(inputs={}))) == pytest.approx(3.0) + + def test_an_unknown_outside_its_brent_is_an_error(self): + # the Brent binds it in `inputs` only while converting its own residual + with pytest.raises( + pybamm.ModelError, match="only has a value inside the Brent" + ): + pybamm.BrentUnknown("x").to_casadi(inputs={}) + + def test_the_tolerances_are_part_of_the_identity(self): + # a Brent converted inside another's oracle is shared with the enclosing + # conversion, so two differing only in tolerance must not be served each other + unknown = pybamm.BrentUnknown("s") + target = pybamm.InputParameter("p") + residual = unknown * unknown - target + coarse = pybamm.Brent(residual, unknown, (0.01, 10), abstol=1e-1, max_iter=100) + fine = pybamm.Brent(residual, unknown, (0.01, 10), abstol=1e-14, max_iter=100) + assert coarse.id != fine.id + + outer_unknown = pybamm.BrentUnknown("outer") + outer = pybamm.Brent(outer_unknown - coarse, outer_unknown, (0.01, 50)) + inputs = {"p": casadi.MX.sym("p")} + shared: dict = {} + converted = [ + node.to_casadi( + casadi.MX.sym("t"), + casadi.MX.sym("y", 1), + inputs=inputs, + casadi_symbols=shared, + ) + for node in (outer, fine) + ] + function = casadi.Function("f", [inputs["p"]], [casadi.vertcat(*converted)]) + np.testing.assert_allclose( + float(np.asarray(function(2.0)).reshape(-1)[1]), np.sqrt(2), rtol=1e-13 + ) diff --git a/packages/pybamm/tests/unit/test_util.py b/packages/pybamm/tests/unit/test_util.py index 8f82706617..cdcc77eb6e 100644 --- a/packages/pybamm/tests/unit/test_util.py +++ b/packages/pybamm/tests/unit/test_util.py @@ -13,20 +13,6 @@ ) -@pytest.fixture -def restore_sys_modules(): - """ - Put ``sys.modules`` back exactly as it was, so a test that hides or reloads - modules cannot leak a stale entry or a second copy of pybamm into later tests. - """ - saved = sys.modules.copy() - try: - yield - finally: - sys.modules.clear() - sys.modules.update(saved) - - class TestUtil: """ Test the functionality in util.py @@ -104,15 +90,16 @@ def test_get_parameters_filepath(self, tmp_path): os.path.join(pybamm.root_dir(), "src", "pybamm", temppath) ) - @pytest.mark.usefixtures("restore_sys_modules") def test_import_optional_dependency(self): optional_distribution_deps = get_optional_distribution_deps("pybamm") present_optional_import_deps = get_present_optional_import_deps( "pybamm", optional_distribution_deps=optional_distribution_deps ) - # Make the optional dependencies not importable + # Save optional dependencies, then make them not importable + modules = {} for import_pkg in present_optional_import_deps: + modules[import_pkg] = sys.modules.get(import_pkg) sys.modules[import_pkg] = None # Test import optional dependency @@ -123,22 +110,32 @@ def test_import_optional_dependency(self): ): pybamm.util.import_optional_dependency(import_pkg) - @pytest.mark.usefixtures("restore_sys_modules") + # Restore optional dependencies + for import_pkg in present_optional_import_deps: + sys.modules[import_pkg] = modules[import_pkg] + def test_pybamm_import(self): + original_symbol = pybamm.Symbol + optional_distribution_deps = get_optional_distribution_deps("pybamm") present_optional_import_deps = get_present_optional_import_deps( "pybamm", optional_distribution_deps=optional_distribution_deps ) - # Make the optional dependencies and their sub-modules not importable - for module_name in list(sys.modules): - if module_name.split(".")[0] in present_optional_import_deps: + # Save optional dependencies and their sub-modules, then make them not importable + modules = {} + for module_name, module in sys.modules.items(): + base_module_name = module_name.split(".")[0] + if base_module_name in present_optional_import_deps: + modules[module_name] = module sys.modules[module_name] = None # Unload pybamm and its sub-modules - for module_name in list(sys.modules): - if module_name.split(".")[0] == "pybamm": - del sys.modules[module_name] + unloaded = {} + for module_name in list(sys.modules.keys()): + base_module_name = module_name.split(".")[0] + if base_module_name == "pybamm": + unloaded[module_name] = sys.modules.pop(module_name) # Test pybamm is still importable try: @@ -147,6 +144,16 @@ def test_pybamm_import(self): pytest.fail( f"Import of 'pybamm' shouldn't require optional dependencies. Error: {error}" ) + finally: + # Restore optional dependencies and their sub-modules + for module_name, module in modules.items(): + sys.modules[module_name] = module + # Restore the original classes: the re-import built a second set, and code + # imported earlier still holds the first, so `isinstance` fails across them. + sys.modules.update(unloaded) + + assert pybamm.Symbol is original_symbol + assert importlib.import_module("pybamm").Symbol is original_symbol def test_optional_dependencies(self): optional_distribution_deps = get_optional_distribution_deps("pybamm") diff --git a/packages/pybammsolvers/CMakeLists.txt b/packages/pybammsolvers/CMakeLists.txt index 579692239d..ea7a620111 100644 --- a/packages/pybammsolvers/CMakeLists.txt +++ b/packages/pybammsolvers/CMakeLists.txt @@ -119,6 +119,9 @@ pybind11_add_module(idaklu src/pybammsolvers/idaklu_source/Options.cpp src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp + src/pybammsolvers/idaklu_source/brent.hpp + src/pybammsolvers/idaklu_source/brent.cpp + src/pybammsolvers/idaklu_source/brent_impl.hpp # IDAKLU expressions / function evaluation [abstract] src/pybammsolvers/idaklu_source/Expressions/Expressions.hpp src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp @@ -170,10 +173,18 @@ if (${USE_PYTHON_CASADI}) file(TO_CMAKE_PATH ${CASADI_INCLUDE_DIR} CASADI_INCLUDE_DIR) message("Found Python CasADi include directory: ${CASADI_INCLUDE_DIR}") target_include_directories(idaklu PRIVATE ${CASADI_INCLUDE_DIR}) + # CasADi's internal headers include their public siblings unqualified + target_include_directories(idaklu PRIVATE ${CASADI_INCLUDE_DIR}/casadi/core) else () message(FATAL_ERROR "Could not find CasADi include directory") endif () + # fallback version, used only if casadi/config.h cannot be found below + execute_process( + COMMAND "${PYTHON_EXECUTABLE}" -c "import casadi; print(casadi.__version__)" + OUTPUT_VARIABLE CASADI_LINKED_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE) + execute_process( COMMAND "${PYTHON_EXECUTABLE}" -c "import casadi; from pathlib import Path; import glob; lib_dir = Path(casadi.__file__).parent; lib_files = list(lib_dir.glob('*casadi*')); print(str(lib_dir) if lib_files else '')" @@ -228,7 +239,189 @@ else () else() target_link_libraries(idaklu PRIVATE casadi) endif() + # CasADi's internal headers include their public siblings unqualified + if(TARGET casadi::casadi) + get_target_property(_casadi_incs casadi::casadi INTERFACE_INCLUDE_DIRECTORIES) + foreach(_casadi_inc IN LISTS _casadi_incs) + target_include_directories(idaklu PRIVATE ${_casadi_inc}/casadi/core) + endforeach() + endif() +endif () + +# casadi::Rootfinder's headers are not installed and are LGPL-3.0-or-later, so stage +# them from the pinned sdist rather than vendoring them into this BSD-3-Clause tree. +set(CASADI_INTERNAL_VERSION "3.7.2" + CACHE STRING "CasADi version to source the internal headers from") +set(CASADI_INTERNAL_SHA256 "b4d7bd8acdc4180306903ae1c9eddaf41be2a3ae2fa7154c57174ae64acdc60d" + CACHE STRING "SHA-256 of the CasADi source distribution") +set(CASADI_SOURCE_DIR "" + CACHE PATH "Local CasADi source tree; set this for offline builds to skip the download") + +option(CASADI_ALLOW_UNVERIFIED_INTERNALS + "Build the \"brent\" plugin even if the linked CasADi cannot be identified" OFF) + +# Identify the linked CasADi from its own config.h: the system path (Windows/vcpkg, +# conda) has no interpreter to ask, and that is where a mismatch would go unnoticed. +set(_casadi_include_dirs "") +if (CASADI_INCLUDE_DIR) + list(APPEND _casadi_include_dirs "${CASADI_INCLUDE_DIR}") +endif () +if (TARGET casadi::casadi) + get_target_property(_casadi_target_incs casadi::casadi INTERFACE_INCLUDE_DIRECTORIES) + if (_casadi_target_incs) + list(APPEND _casadi_include_dirs ${_casadi_target_incs}) + endif () +endif () + +unset(CASADI_CONFIG_HEADER CACHE) +find_file(CASADI_CONFIG_HEADER casadi/config.h PATHS ${_casadi_include_dirs} NO_DEFAULT_PATH) + +set(_casadi_config_flags "") +if (CASADI_CONFIG_HEADER) + file(READ "${CASADI_CONFIG_HEADER}" _casadi_config_contents) + string(REGEX MATCH "CASADI_VERSION_STRING[ \t]+\"([0-9]+\\.[0-9]+\\.[0-9]+)\"" + _casadi_version_match "${_casadi_config_contents}") + set(CASADI_LINKED_VERSION "${CMAKE_MATCH_1}") + string(REGEX MATCH "CASADI_COMPILER_FLAGS[ \t]+\"([^\"]*)\"" + _casadi_flags_match "${_casadi_config_contents}") + set(_casadi_config_flags "${CMAKE_MATCH_1}") +endif () + +if (NOT CASADI_LINKED_VERSION) + if (CASADI_ALLOW_UNVERIFIED_INTERNALS) + message(WARNING + "Could not read the linked CasADi version from casadi/config.h. The \"brent\" " + "plugin is being built against CasADi ${CASADI_INTERNAL_VERSION} internals " + "without checking that they match; an ABI mismatch will not be diagnosed.") + else () + message(FATAL_ERROR + "Could not find casadi/config.h in any of: ${_casadi_include_dirs}\n" + "The \"brent\" rootfinder plugin subclasses CasADi internals that are not " + "ABI-stable, so the linked version must be known before it can be built. " + "Configure with -DCASADI_ALLOW_UNVERIFIED_INTERNALS=ON to build anyway.") + endif () +elseif (NOT CASADI_LINKED_VERSION STREQUAL CASADI_INTERNAL_VERSION) + message(FATAL_ERROR + "CasADi version mismatch: linking CasADi ${CASADI_LINKED_VERSION} but sourcing the " + "internal headers for the \"brent\" plugin from ${CASADI_INTERNAL_VERSION}. These " + "headers are not ABI-stable, so they must match. Update CASADI_INTERNAL_VERSION and " + "CASADI_INTERNAL_SHA256 in this file to match the casadi pin in pyproject.toml.") +else () + message(STATUS "idaklu: \"brent\" plugin verified against CasADi ${CASADI_LINKED_VERSION}") +endif () + +set(_casadi_internal_headers + casadi_os.hpp + function_internal.hpp + oracle_function.hpp + plugin_interface.hpp + rootfinder_impl.hpp) +set(_casadi_internal_work "${CMAKE_CURRENT_BINARY_DIR}/casadi-internal") +set(_casadi_internal_include "${_casadi_internal_work}/include") + +# editable.rebuild reconfigures on every import; skip the work once staged. +set(_casadi_internal_staged TRUE) +foreach (_header IN LISTS _casadi_internal_headers) + if (NOT EXISTS "${_casadi_internal_include}/casadi/core/${_header}") + set(_casadi_internal_staged FALSE) + endif () +endforeach () + +if (NOT _casadi_internal_staged) + if (NOT CASADI_SOURCE_DIR AND DEFINED ENV{CASADI_SOURCE_DIR}) + set(CASADI_SOURCE_DIR "$ENV{CASADI_SOURCE_DIR}") + endif () + + if (CASADI_SOURCE_DIR) + file(TO_CMAKE_PATH "${CASADI_SOURCE_DIR}" _casadi_internal_src) + message(STATUS "idaklu: taking CasADi internal headers from ${_casadi_internal_src}") + else () + set(_casadi_sdist "${_casadi_internal_work}/casadi-${CASADI_INTERNAL_VERSION}.tar.gz") + set(_casadi_sdist_url + "https://pypi.org/packages/source/c/casadi/casadi-${CASADI_INTERNAL_VERSION}.tar.gz") + message(STATUS "idaklu: downloading ${_casadi_sdist_url} for the CasADi internal headers") + file(DOWNLOAD "${_casadi_sdist_url}" "${_casadi_sdist}" + EXPECTED_HASH SHA256=${CASADI_INTERNAL_SHA256} + STATUS _casadi_download_status) + list(GET _casadi_download_status 0 _casadi_download_rc) + if (_casadi_download_rc) + list(GET _casadi_download_status 1 _casadi_download_message) + file(REMOVE "${_casadi_sdist}") + message(FATAL_ERROR + "Could not download the CasADi source distribution, which supplies the " + "internal headers the \"brent\" rootfinder plugin is built against:\n" + " ${_casadi_sdist_url}\n" + " ${_casadi_download_message}\n" + "For an offline build, unpack that archive yourself and configure with " + "-DCASADI_SOURCE_DIR= (or set the CASADI_SOURCE_DIR environment " + "variable), pointing at a CasADi ${CASADI_INTERNAL_VERSION} source tree.") + endif () + set(_casadi_internal_members "") + foreach (_header IN LISTS _casadi_internal_headers) + list(APPEND _casadi_internal_members + "casadi-${CASADI_INTERNAL_VERSION}/casadi/core/${_header}") + endforeach () + execute_process( + COMMAND "${CMAKE_COMMAND}" -E tar xzf "${_casadi_sdist}" ${_casadi_internal_members} + WORKING_DIRECTORY "${_casadi_internal_work}" + RESULT_VARIABLE _casadi_extract_rc) + if (_casadi_extract_rc) + message(FATAL_ERROR + "Could not extract the CasADi internal headers from ${_casadi_sdist} " + "(exit ${_casadi_extract_rc}).") + endif () + set(_casadi_internal_src "${_casadi_internal_work}/casadi-${CASADI_INTERNAL_VERSION}") + endif () + + # Stage only the listed headers, so the rest of a CasADi source tree can never + # shadow the generated public headers of the CasADi actually being linked. + foreach (_header IN LISTS _casadi_internal_headers) + if (NOT EXISTS "${_casadi_internal_src}/casadi/core/${_header}") + message(FATAL_ERROR + "${_casadi_internal_src} is not a CasADi source tree: " + "casadi/core/${_header} is missing.") + endif () + configure_file( + "${_casadi_internal_src}/casadi/core/${_header}" + "${_casadi_internal_include}/casadi/core/${_header}" + COPYONLY) + endforeach () +endif () + +target_include_directories(idaklu PRIVATE ${_casadi_internal_include}) + +# brent_impl.hpp is emitted verbatim into generated C, so stringify it rather than +# keeping a second copy. +set(BRENT_IMPL_HEADER + "${CMAKE_CURRENT_SOURCE_DIR}/src/pybammsolvers/idaklu_source/brent_impl.hpp") +file(READ "${BRENT_IMPL_HEADER}" BRENT_IMPL_SOURCE) +file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/generated/brent_impl_source.hpp" + "// Generated by CMake from brent_impl.hpp. Do not edit.\n" + "namespace casadi {\n" + "static const char brent_impl_str[] = R\"PYBAMM_BRENT(\n" + "${BRENT_IMPL_SOURCE}" + ")PYBAMM_BRENT\";\n" + "} // namespace casadi\n") +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${BRENT_IMPL_HEADER}") +target_include_directories(idaklu PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/generated") +# Reuse CASADI_COMPILER_FLAGS from config.h: CASADI_WITH_THREADSAFE_SYMBOLICS adds a +# static mutex to Rootfinder, so guessing it wrong is an ABI break. +set(_casadi_internal_defs "") +if (_casadi_config_flags) + separate_arguments(_casadi_config_flag_list UNIX_COMMAND "${_casadi_config_flags}") + foreach (_flag IN LISTS _casadi_config_flag_list) + if (_flag MATCHES "^-D(.+)$") + list(APPEND _casadi_internal_defs "${CMAKE_MATCH_1}") + endif () + endforeach () +endif () +if (NOT _casadi_internal_defs) + # No config.h to read: fall back to the flags CasADi ${CASADI_INTERNAL_VERSION} uses. + set(_casadi_internal_defs CASADI_VERSION=31 WITH_DL + CASADI_WITH_THREADSAFE_SYMBOLICS CASADI_WITH_THREAD) endif () +message(STATUS "idaklu: CasADi internal header flags: ${_casadi_internal_defs}") +target_compile_definitions(idaklu PRIVATE ${_casadi_internal_defs}) # Match idaklu's libstdc++ std::string ABI to the linked libcasadi (GNU only; # no-op for MSVC and macOS/libc++). CasADi's PyPI casadi-config.cmake forces diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp index 1440d4d2d6..db994629e8 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp @@ -18,6 +18,7 @@ #include "idaklu_source/sundials_error_handler.hpp" #include "idaklu_source/reduce.hpp" #include "idaklu_source/StandaloneNewtonSolver.hpp" +#include "idaklu_source/brent.hpp" casadi::Function generate_casadi_function(const std::string &data) @@ -270,6 +271,10 @@ PYBIND11_MODULE(idaklu, m) py::arg("t_eval"), py::arg("y0_alg"), py::arg("inputs"), py::return_value_policy::move); + // Register the "brent" rootfinder with CasADi. Rootfinder::solvers_ is a + // process-global, so registering here covers every CasADi user in the process. + casadi::casadi_load_rootfinder_brent(); + py::class_(m, "Function"); py::class_(m, "solution") diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp new file mode 100644 index 0000000000..c8bb5db3be --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp @@ -0,0 +1,254 @@ +#include "brent.hpp" + +#include + +#include + +// the iteration, compiled here and stringified by CMake into brent_impl_str +#include "brent_impl.hpp" +#include "brent_impl_source.hpp" + +namespace casadi { + +constexpr casadi_int BRACKET_LO = 1; +constexpr casadi_int BRACKET_HI = 2; + +extern "C" int CASADI_ROOTFINDER_BRENT_EXPORT +casadi_register_rootfinder_brent(Rootfinder::Plugin* plugin) { + plugin->creator = Brent::creator; + plugin->name = "brent"; + plugin->doc = Brent::meta_doc.c_str(); + plugin->version = CASADI_VERSION; + plugin->options = &Brent::options_; + plugin->deserialize = &Brent::deserialize; + return 0; +} + +extern "C" void CASADI_ROOTFINDER_BRENT_EXPORT casadi_load_rootfinder_brent() { + Rootfinder::registerPlugin(casadi_register_rootfinder_brent); +} + +const std::string Brent::meta_doc = + "Brent's method for a scalar residual on a bracket. Derivative free, and the " + "iterate never leaves the bracket."; + +const Options Brent::options_ = { + {&Rootfinder::options_}, + {{"abstol", {OT_DOUBLE, "Absolute tolerance on the unknown"}}, + {"max_iter", {OT_INT, "Maximum number of iterations"}}}}; + +void Brent::init(const Dict& opts) { + Rootfinder::init(opts); + + for (auto&& op : opts) { + if (op.first == "abstol") { + abstol_ = op.second; + } else if (op.first == "max_iter") { + max_iter_ = op.second; + } + } + + casadi_assert(n_ == 1, "Brent solves a scalar residual, got n=" + str(n_)); + casadi_assert(n_in_ >= 3, + "Brent reads its bracket from inputs 1 and 2, so the oracle must be " + "g(x, lo, hi, ...); got " + str(n_in_) + " input(s)"); + casadi_assert(max_iter_ > 0, "max_iter must be positive, got " + str(max_iter_)); + casadi_assert(abstol_ > 0, "abstol must be positive, got " + str(abstol_)); + + set_function(oracle_, "g"); +} + +int Brent::init_mem(void* mem) const { + if (Rootfinder::init_mem(mem)) return 1; + auto m = static_cast(mem); + m->iter = 0; + m->return_status = "unset"; + return 0; +} + +int Brent::residual(void* user_data, double x, double* fx) { + auto ctx = static_cast(user_data); + const Brent* self = ctx->solver; + BrentMemory* m = ctx->mem; + std::copy_n(m->iarg, self->n_in_, m->arg); + m->arg[self->iin_] = &x; + std::copy_n(m->ires, self->n_out_, m->res); + m->res[self->iout_] = fx; + return self->calc_function(m, "g"); +} + +int Brent::solve(void* mem) const { + auto m = static_cast(mem); + Context ctx{this, m}; + + const double a = m->iarg[BRACKET_LO][0]; + const double b = m->iarg[BRACKET_HI][0]; + + // the root never depends on the guess, so the guess is left out of the key + key_.clear(); + for (casadi_int i = 0; i < n_in_; ++i) { + if (i == iin_ || !m->iarg[i]) continue; + key_.insert(key_.end(), m->iarg[i], m->iarg[i] + nnz_in(i)); + } + if (m->cache_valid && m->cache_key == key_) { + ++m->cache_hits; + if (m->ires[iout_]) m->ires[iout_][0] = m->cache_root; + m->return_status = "success (cached)"; + m->success = true; + return 0; + } + + double root = 0; + const int flag = casadi_brent(&Brent::residual, &ctx, a, b, abstol_, max_iter_, + &root, &m->iter); + if (flag) { + m->return_status = + flag == 2 ? "no sign change over the bracket" : + flag == 3 ? "iteration limit reached without converging" : "residual failed"; + m->unified_return_status = SOLVER_RET_UNKNOWN; + m->success = false; + return 0; + } + + m->cache_key = key_; + m->cache_root = root; + m->cache_valid = true; + + if (m->ires[iout_]) m->ires[iout_][0] = root; + m->return_status = "success"; + m->success = true; + return 0; +} + +void Brent::codegen_declarations(CodeGenerator& g) const { + // Adding the oracle first keeps its definition out of the middle of the wrapper + // emitted below, which is written straight to the buffer. + g.add_dependency(get_function("g")); + + // one definition per file, not per node; add_shorthand stays off so the name is not + // CASADI_PREFIX-renamed, which would defeat the guard + g << "#ifndef CASADI_BRENT_IMPL\n" + << "#define CASADI_BRENT_IMPL\n" + // The cache below is per instance and per thread. Generated code is expected to + // be reentrant, so plain statics will not do. + << "#if defined(__cplusplus) && __cplusplus >= 201103L\n" + << "#define CASADI_BRENT_TLS thread_local\n" + << "#elif defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L" + " && !defined(__STDC_NO_THREADS__)\n" + << "#define CASADI_BRENT_TLS _Thread_local\n" + << "#elif defined(__GNUC__) || defined(__clang__)\n" + << "#define CASADI_BRENT_TLS __thread\n" + << "#else\n" + << "#define CASADI_BRENT_TLS\n" + << "#endif\n" + << "struct casadi_brent_data {\n" + << " const casadi_real** arg;\n" + << " casadi_real** res;\n" + << " casadi_int* iw;\n" + << " casadi_real* w;\n" + << "};\n" + << g.sanitize_source(brent_impl_str, {"casadi_real"}, false) + << "#endif\n\n"; + + // The residual callback is per instance: it hard-codes this oracle and this + // implicit input/output pair. + g << "static int " << g.shorthand("brent_res_" + codegen_name(g, false)) + << "(void* user_data, casadi_real x, casadi_real* fx) {\n" + << " struct casadi_brent_data* d = (struct casadi_brent_data*) user_data;\n" + << " const casadi_real** arg1 = d->arg + " << n_in_ << ";\n" + << " casadi_real** res1 = d->res + " << n_out_ << ";\n"; + for (casadi_int i = 0; i < n_in_; ++i) { + g << " arg1[" << i << "] = " << (i == iin_ ? "&x" : "d->" + g.arg(i)) << ";\n"; + } + for (casadi_int i = 0; i < n_out_; ++i) { + g << " res1[" << i << "] = " << (i == iout_ ? "fx" : "d->" + g.res(i)) << ";\n"; + } + g << " return " << g(get_function("g"), "arg1", "res1", "d->iw", "d->w") << ";\n" + << "}\n\n"; + + // the same cache as BrentMemory, per instance and per thread + const std::string c = cache_name(g); + g << "static CASADI_BRENT_TLS casadi_real " << c << "_key[" << cache_size() << "];\n" + << "static CASADI_BRENT_TLS casadi_real " << c << "_root;\n" + << "static CASADI_BRENT_TLS int " << c << "_valid = 0;\n\n"; +} + +std::string Brent::cache_name(CodeGenerator& g) const { + return g.shorthand("brent_cache_" + codegen_name(g, false)); +} + +casadi_int Brent::cache_size() const { + casadi_int n = 0; + for (casadi_int i = 0; i < n_in_; ++i) if (i != iin_) n += nnz_in(i); + return n; +} + +void Brent::codegen_body(CodeGenerator& g) const { + g.local("brent_data", "struct casadi_brent_data"); + g.local("brent_iter", "casadi_int"); + g.local("brent_root", "casadi_real"); + g.local("brent_flag", "int"); + + // sz_w_per_ is zero, so the oracle's scratch starts at w -- the same slice + // calc_function hands the oracle on the interpreted path. + g << "brent_data.arg = arg;\n" + << "brent_data.res = res;\n" + << "brent_data.iw = iw;\n" + << "brent_data.w = w;\n"; + + const std::string lo = g.arg(BRACKET_LO) + "[0]"; + const std::string hi = g.arg(BRACKET_HI) + "[0]"; + const std::string c = cache_name(g); + + std::vector key; + for (casadi_int i = 0; i < n_in_; ++i) { + if (i == iin_) continue; + for (casadi_int e = 0; e < nnz_in(i); ++e) { + const std::string a = g.arg(i); + key.push_back("(" + a + " ? " + a + "[" + str(e) + "] : 0)"); + } + } + + g.local("brent_hit", "int"); + g << "brent_hit = " << c << "_valid;\n"; + for (casadi_int j = 0; j < static_cast(key.size()); ++j) { + g << "if (brent_hit && " << c << "_key[" << j << "] != " << key[j] + << ") brent_hit = 0;\n"; + } + g << "if (!brent_hit) {\n" + << " brent_flag = casadi_brent(" + << g.shorthand("brent_res_" + codegen_name(g, false)) + << ", &brent_data, " << lo << ", " << hi << ", " << g.constant(abstol_) << ", " + << max_iter_ << ", &brent_root, &brent_iter);\n" + << " if (brent_flag) return 1;\n"; + for (casadi_int j = 0; j < static_cast(key.size()); ++j) { + g << " " << c << "_key[" << j << "] = " << key[j] << ";\n"; + } + g << " " << c << "_root = brent_root;\n" + << " " << c << "_valid = 1;\n" + << "}\n" + << "if (" << g.res(iout_) << ") " << g.res(iout_) << "[0] = " << c << "_root;\n"; +} + +void Brent::serialize_body(SerializingStream& s) const { + Rootfinder::serialize_body(s); + s.version("Brent", 1); + s.pack("Brent::abstol", abstol_); + s.pack("Brent::max_iter", max_iter_); +} + +Brent::Brent(DeserializingStream& s) : Rootfinder(s) { + s.version("Brent", 1); + s.unpack("Brent::abstol", abstol_); + s.unpack("Brent::max_iter", max_iter_); +} + +Dict Brent::get_stats(void* mem) const { + Dict stats = Rootfinder::get_stats(mem); + auto m = static_cast(mem); + stats["iter_count"] = m->iter; + stats["return_status"] = m->return_status; + return stats; +} + +} // namespace casadi diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp new file mode 100644 index 0000000000..d9ba19b172 --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp @@ -0,0 +1,122 @@ +// +// A CasADi rootfinder plugin built out of tree, which takes four things from +// CMakeLists.txt that an in-tree CasADi plugin would get for free: +// +// 1. `rootfinder_impl.hpp` and four siblings. CasADi does not install these +// (INSTALL_INTERNAL_HEADERS is off by default) and they are LGPL, so CMake stages +// them from the pinned sdist at configure time instead of vendoring them here. +// Use -DCASADI_SOURCE_DIR= to build offline. +// 2. A version check. These headers are not ABI-stable, so CMake reads the linked +// CasADi's version from its own installed casadi/config.h -- the one source that +// exists on every discovery path -- and refuses to build on a mismatch. +// 3. The flags CasADi itself was compiled with, also from config.h. +// CASADI_WITH_THREADSAFE_SYMBOLICS adds a static mutex to Rootfinder, so guessing +// it would be an ABI break rather than a warning. +// 4. `brent_impl_str`, the iteration stringified from brent_impl.hpp, so codegen can +// emit the same text this file compiles. +// +// The export macro below is the fifth: in tree it comes from a generated header. +// +#ifndef PYBAMM_BRENT_HPP +#define PYBAMM_BRENT_HPP + +#include "casadi/core/rootfinder_impl.hpp" +#include + +// whatever the toolchain spells "visible in this shared object" +#if defined(_WIN32) || defined(__CYGWIN__) +#define CASADI_ROOTFINDER_BRENT_EXPORT __declspec(dllexport) +#elif defined(__GNUC__) || defined(__clang__) +#define CASADI_ROOTFINDER_BRENT_EXPORT __attribute__((visibility("default"))) +#else +#define CASADI_ROOTFINDER_BRENT_EXPORT +#endif + +namespace casadi { + +struct CASADI_ROOTFINDER_BRENT_EXPORT BrentMemory : public RootfinderMemory { + casadi_int iter; + const char* return_status; + // Last solve, keyed on every input but the guess, so a Brent nested inside another + // is not re-solved on every iteration of the enclosing one. + std::vector cache_key; + double cache_root = 0; + bool cache_valid = false; + casadi_int cache_hits = 0; +}; + +/** + * @brief Brent's method, registered as the CasADi rootfinder plugin "brent". + * + * Solves a scalar g(x, p) = 0 on a bracket. Brent needs only a sign change over that + * bracket, so it converges on residuals where a Newton iteration stalls or leaves the + * domain, and the iterate is confined to the bracket by construction. + * + * The oracle must be ``g(x, lo, hi, ...)``: the bracket is read from inputs 1 and 2 at + * solve time, so it can be a live value in the surrounding graph rather than a constant. + * + * Derivatives come from :class:`Rootfinder`, which applies the implicit function theorem; + * nothing here is differentiated. + */ +class CASADI_ROOTFINDER_BRENT_EXPORT Brent : public Rootfinder { +public: + explicit Brent(const std::string& name, const Function& f) : Rootfinder(name, f) {} + ~Brent() override { clear_mem(); } + + const char* plugin_name() const override { return "brent"; } + std::string class_name() const override { return "Brent"; } + + static Rootfinder* creator(const std::string& name, const Function& f) { + return new Brent(name, f); + } + + static const Options options_; + const Options& get_options() const override { return options_; } + static const std::string meta_doc; + + void init(const Dict& opts) override; + mutable std::vector key_; // cache scratch + int solve(void* mem) const override; + + void* alloc_mem() const override { return new BrentMemory(); } + int init_mem(void* mem) const override; + void free_mem(void* mem) const override { delete static_cast(mem); } + Dict get_stats(void* mem) const override; + + // Emit the iteration as C so an expression containing a Brent node survives + // Function::generate() and JIT, which is how PyBaMM AOT-compiles its functions. + bool has_codegen() const override { return true; } + void codegen_declarations(CodeGenerator& g) const override; + void codegen_body(CodeGenerator& g) const override; + std::string cache_name(CodeGenerator& g) const; + casadi_int cache_size() const; + + // PyBaMM round-trips its functions through serialize(), so the tolerances have to + // survive it or a deserialised Brent silently reverts to the defaults. + void serialize_body(SerializingStream& s) const override; + static ProtoFunction* deserialize(DeserializingStream& s) { return new Brent(s); } + +protected: + explicit Brent(DeserializingStream& s); + + /// What the residual needs to reach the oracle, passed through casadi_brent's void*. + struct Context { + const Brent* solver; + BrentMemory* mem; + }; + + /// casadi_brent's residual callback for the interpreted path. + static int residual(void* user_data, double x, double* fx); + + double abstol_{1e-14}; + casadi_int max_iter_{100}; +}; + +extern "C" int CASADI_ROOTFINDER_BRENT_EXPORT +casadi_register_rootfinder_brent(Rootfinder::Plugin* plugin); + +extern "C" void CASADI_ROOTFINDER_BRENT_EXPORT casadi_load_rootfinder_brent(); + +} // namespace casadi + +#endif // PYBAMM_BRENT_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp new file mode 100644 index 0000000000..cd7dc02f2a --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp @@ -0,0 +1,89 @@ +// +// Brent's method on a bracket, written in CasADi's runtime style (a `T1` template +// that `CodeGenerator::sanitize_source` turns into plain C). +// +// This text is both compiled and stringified into the generated C, so keep it valid C +// once the template header and comments are stripped: declarations first, no C++-only +// constructs, no preprocessor directives (`sanitize_source` drops `#define`/`#undef`). +// +// `res_fn` evaluates the residual at `x` into `*fx` and returns nonzero on failure. +// Returns 0 on success (`*out` is the root, `*iter` the iteration count), 1 if the +// residual failed, 2 if the bracket shows no sign change, 3 if `max_iter` was reached +// without the bracket shrinking to `abstol`. +// + +// `static` so that two generated files, each carrying its own copy, still link. +template +static int casadi_brent(int (*res_fn)(void*, T1, T1*), void* user_data, + T1 a, T1 b, T1 abstol, casadi_int max_iter, + T1* out, casadi_int* iter) { + T1 fa, fb, c, fc, d, e, tol, xm, p, q, s, qq, r, step; + casadi_int k; + *iter = 0; + if (res_fn(user_data, a, &fa)) return 1; + if (res_fn(user_data, b, &fb)) return 1; + // No sign change, so the bracket contains no root. Report it rather than + // returning whichever end happens to be closer. + if (!(fa * fb <= 0)) return 2; + c = a; + fc = fa; + d = b - a; + e = d; + for (k = 0; k < max_iter; ++k) { + if (fb * fc > 0) { + c = a; + fc = fa; + d = b - a; + e = d; + } + if (fabs(fc) < fabs(fb)) { + a = b; + b = c; + c = a; + fa = fb; + fb = fc; + fc = fa; + } + // DBL_EPSILON, spelled out: the generated C includes no headers of its own + tol = 2 * 2.2204460492503131e-16 * fabs(b) + 0.5 * abstol; + xm = 0.5 * (c - b); + if (fabs(xm) <= tol || fb == 0) break; + if (fabs(e) >= tol && fabs(fa) > fabs(fb)) { + // Inverse quadratic interpolation, or secant when only two points are distinct + s = fb / fa; + if (a == c) { + p = 2 * xm * s; + q = 1 - s; + } else { + qq = fa / fc; + r = fb / fc; + p = s * (2 * xm * qq * (qq - r) - (b - a) * (r - 1)); + q = (qq - 1) * (r - 1) * (s - 1); + } + if (p > 0) q = -q; + p = fabs(p); + // Take the interpolated step only while it keeps bisecting; else bisect + if (2 * p < fmin(3 * xm * q - fabs(tol * q), fabs(e * q))) { + e = d; + d = p / q; + } else { + d = xm; + e = d; + } + } else { + d = xm; + e = d; + } + a = b; + fa = fb; + step = fabs(d) > tol ? d : (xm > 0 ? tol : -tol); + b += step; + if (res_fn(user_data, b, &fb)) return 1; + } + *iter = k; + // Falling out of the loop means the bracket never shrank to `abstol`, so `b` is + // not a root. Reporting it as one would hand back an arbitrary point silently. + if (k >= max_iter) return 3; + *out = b; + return 0; +} diff --git a/packages/pybammsolvers/tests/test_brent_rootfinder.py b/packages/pybammsolvers/tests/test_brent_rootfinder.py new file mode 100644 index 0000000000..26dc808669 --- /dev/null +++ b/packages/pybammsolvers/tests/test_brent_rootfinder.py @@ -0,0 +1,320 @@ +"""Unit tests for the "brent" CasADi rootfinder plugin.""" + +from __future__ import annotations + +import casadi +import numpy as np +import pytest +from scipy.optimize import brentq + +import pybammsolvers.idaklu # noqa: F401 registers the plugin on import + +LO, HI = 1e-9, 1 - 1e-9 + + +def non_monotone(x, p): + """Several turning points inside the bracket, where a Newton iteration stalls.""" + return casadi.sin(12 * x) * 0.08 + (1.6 - x) ** 3 - p + + +def _solver(expr_fn, lo=LO, hi=HI, **opts): + """A rootfinder over ``expr_fn(x, p) == 0``, and a callable for the residual. + + The bracket is bound to constants here, so the returned solver takes ``p``. + """ + x, lo_s, hi_s, p = (casadi.MX.sym(n) for n in ("x", "lo", "hi", "p")) + g = casadi.Function("g", [x, lo_s, hi_s, p], [expr_fn(x, p)]) + rf = casadi.rootfinder("rf", "brent", g, opts) + bound = casadi.Function("bound", [p], [rf(0.0, lo, hi, p)]) + return bound, casadi.Function("f", [x, p], [expr_fn(x, p)]) + + +class TestBrentPlugin: + def test_importing_pybammsolvers_registers_the_plugin(self): + assert casadi.has_rootfinder("brent") + + def test_solves_a_scalar_equation(self): + rf, _ = _solver(lambda x, p: x * x - p * x - 6.0, lo=0.0, hi=10.0) + assert float(rf(1.0)) == pytest.approx(3.0, abs=1e-12) + + @pytest.mark.parametrize( + "expr_fn", [lambda x, p: casadi.exp(x) + p * x - 2.0, non_monotone] + ) + def test_matches_scipy_over_a_sweep(self, expr_fn): + rf, f = _solver(expr_fn) + worst_difference = worst_residual = 0.0 + for p in np.linspace(0.2, 1.8, 40): + if float(f(LO, p)) * float(f(HI, p)) > 0: + continue # no bracket for this p, nothing to compare + got = float(rf(p)) + want = brentq(lambda x, p=p: float(f(x, p)), LO, HI, xtol=2e-12) + assert LO <= got <= HI + worst_difference = max(worst_difference, abs(got - want)) + worst_residual = max(worst_residual, abs(float(f(got, p)))) + assert worst_difference < 1e-9 + assert worst_residual < 1e-12 + + def test_bracket_can_be_a_graph_input(self): + # x^2 - 6 has roots at +-sqrt(6); which one is found is set by the bracket, + # passed as a live value rather than an option + x, lo, hi = casadi.MX.sym("x"), casadi.MX.sym("lo"), casadi.MX.sym("hi") + f = casadi.Function("f", [x, lo, hi], [x * x - 6.0]) + rf = casadi.rootfinder("rf", "brent", f, {}) + assert float(rf(0.0, 0.0, 10.0)) == pytest.approx(np.sqrt(6), abs=1e-12) + assert float(rf(0.0, -10.0, 0.0)) == pytest.approx(-np.sqrt(6), abs=1e-12) + + def test_derivatives_come_from_the_implicit_function_theorem(self): + # x^2 - p0 x - p1 = 0 at (1, 6) has root 3; dx/dp = [x, 1] / (2x - p0) + x, lo, hi = casadi.MX.sym("x"), casadi.MX.sym("lo"), casadi.MX.sym("hi") + p = casadi.MX.sym("p", 2) + g = casadi.Function("g", [x, lo, hi, p], [x * x - p[0] * x - p[1]]) + rf = casadi.rootfinder("rf", "brent", g, {}) + root = rf(0.0, 0.0, 10.0, p) + jacobian = casadi.Function("J", [p], [casadi.jacobian(root, p)]) + np.testing.assert_allclose( + np.asarray(jacobian(casadi.DM([1.0, 6.0]))).ravel(), [0.6, 0.2], rtol=1e-12 + ) + + def test_composes_inside_a_graph(self): + rf, _ = _solver(lambda x, p: x * x - p * x - 6.0, lo=0.0, hi=10.0) + p = casadi.MX.sym("p") + composed = casadi.Function("composed", [p], [3 * rf(p) + 1]) + assert float(composed(1.0)) == pytest.approx(10.0, abs=1e-12) + + def test_survives_a_serialize_round_trip(self): + # pybamm hands functions to IDAKLU as generate_function(fn.serialize()) + rf, _ = _solver(lambda x, p: x * x - p * x - 6.0, lo=0.0, hi=10.0) + p = casadi.MX.sym("p") + composed = casadi.Function("composed", [p], [3 * rf(p) + 1]) + rebuilt = casadi.Function.deserialize(composed.serialize()) + assert float(rebuilt(1.0)) == float(composed(1.0)) == pytest.approx(10.0) + + def test_no_sign_change_fails_rather_than_guessing(self): + rf, _ = _solver(lambda x, p: x * x + 1.0, lo=0.0, hi=1.0) + with pytest.raises(RuntimeError, match="rootfinder process failed"): + rf(0.0) + + def test_reports_iteration_count(self): + x, lo, hi, p = (casadi.MX.sym(n) for n in ("x", "lo", "hi", "p")) + g = casadi.Function("g", [x, lo, hi, p], [casadi.exp(x) + p * x - 2.0]) + rf = casadi.rootfinder("rf", "brent", g, {}) + rf(0.0, LO, HI, 1.0) + assert rf.stats()["iter_count"] > 0 + assert rf.stats()["return_status"] == "success" + + def test_rejects_a_non_scalar_system(self): + x = casadi.MX.sym("x", 2) + lo, hi = casadi.MX.sym("lo"), casadi.MX.sym("hi") + g = casadi.Function("g", [x, lo, hi], [x - 1.0]) + with pytest.raises(RuntimeError, match="Brent solves a scalar residual"): + casadi.rootfinder("rf", "brent", g, {}) + + def test_rejects_an_oracle_without_bracket_inputs(self): + x, p = casadi.MX.sym("x"), casadi.MX.sym("p") + g = casadi.Function("g", [x, p], [x - p]) + with pytest.raises(RuntimeError, match="g\\(x, lo, hi"): + casadi.rootfinder("rf", "brent", g, {}) + + +class TestBrentCodegen: + """PyBaMM AOT-compiles its CasADi functions, so a Brent node has to survive + Function.generate() -> C -> compile -> casadi.external unchanged.""" + + @staticmethod + def _compile(function, tmp_path, name): + """Generate, compile and load ``function``; returns the C and the external. + + CasADi's own Importer drives the compiler, so this works wherever CasADi + does rather than only where a POSIX ``cc`` is on the path. + """ + function.generate(f"{name}.c", {"with_header": False}) + source = tmp_path / f"{name}.c" + external = casadi.external(name, casadi.Importer(str(source), "shell")) + return source.read_text(), external + + @pytest.fixture + def _in_tmp_path(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + return tmp_path + + def test_generates_compilable_c(self, _in_tmp_path, recwarn): + rf, _ = _solver(non_monotone) + p = casadi.MX.sym("p") + composed = casadi.Function("composed", [p], [rf(p)]) + source, _ = self._compile(composed, _in_tmp_path, "composed") + assert "#error" not in source + assert not [w for w in recwarn if "code generated" in str(w.message)] + + def test_generated_c_matches_the_interpreted_plugin(self, _in_tmp_path): + rf, f = _solver(non_monotone) + p = casadi.MX.sym("p") + composed = casadi.Function("composed", [p], [rf(p)]) + _, external = self._compile(composed, _in_tmp_path, "composed") + + compared = 0 + for target in np.linspace(0.2, 1.8, 61): + if float(f(LO, target)) * float(f(HI, target)) > 0: + continue # no bracket for this target + np.testing.assert_allclose( + float(external(target)), float(composed(target)), rtol=0, atol=1e-14 + ) + compared += 1 + assert compared >= 40 + + def test_a_bracket_read_from_an_input_survives_codegen(self, _in_tmp_path): + x, lo, hi = casadi.MX.sym("x"), casadi.MX.sym("lo"), casadi.MX.sym("hi") + f = casadi.Function("f", [x, lo, hi], [x * x - 6.0]) + rf = casadi.rootfinder("rf", "brent", f, {}) + a, b = casadi.MX.sym("a"), casadi.MX.sym("b") + composed = casadi.Function("bracketed", [a, b], [rf(0.0, a, b)]) + _, external = self._compile(composed, _in_tmp_path, "bracketed") + for bracket in ((0.0, 10.0), (-10.0, 0.0)): + np.testing.assert_allclose( + float(external(*bracket)), float(composed(*bracket)), rtol=0, atol=1e-14 + ) + + def test_two_brent_nodes_share_one_iteration(self, _in_tmp_path): + first, _ = _solver(non_monotone) + second, _ = _solver(lambda x, p: casadi.exp(x) + p * x - 2.0) + p = casadi.MX.sym("p") + composed = casadi.Function("both", [p], [first(p) + second(p)]) + source, external = self._compile(composed, _in_tmp_path, "both") + # one iteration behind an include guard, one residual wrapper per node + assert source.count("#ifndef CASADI_BRENT_IMPL") == 2 + assert source.count("static int casadi_brent_res_") == 2 + np.testing.assert_allclose( + float(external(0.9)), float(composed(0.9)), rtol=0, atol=1e-14 + ) + + def test_the_derivative_survives_codegen(self, _in_tmp_path): + rf, f = _solver(non_monotone) + p = casadi.MX.sym("p") + composed = casadi.Function("composed", [p], [rf(p)]) + jacobian = casadi.Function("djac", [p], [casadi.jacobian(composed(p), p)]) + _, external = self._compile(jacobian, _in_tmp_path, "djac") + for target in (0.4, 0.9, 1.4): + np.testing.assert_allclose( + float(external(target)), float(jacobian(target)), rtol=0, atol=1e-14 + ) + + +class TestBrentCache: + """A Brent nested inside another must not re-solve on every enclosing iteration. + + Covered on the interpreted and the generated path. + """ + + @staticmethod + def _nested(): + """``x`` such that ``g(x) = 0`` where ``g`` reads an inner solve of its own.""" + target = casadi.MX.sym("target") + inner_rf, _ = _solver(non_monotone) + inner = inner_rf(target) + outer_x = casadi.MX.sym("x") + outer = casadi.rootfinder( + "outer", + "brent", + casadi.Function( + "outer_g", + [outer_x, casadi.MX.sym("lo"), casadi.MX.sym("hi"), target], + [outer_x - inner], + ), + {"abstol": 1e-13, "max_iter": 200}, + ) + return casadi.Function("nested", [target], [outer(LO, LO, HI, target)]) + + def test_repeating_the_inputs_reuses_the_last_solve(self): + x, lo, hi, p = (casadi.MX.sym(n) for n in ("x", "lo", "hi", "p")) + rf = casadi.rootfinder( + "rf", + "brent", + casadi.Function("g", [x, lo, hi, p], [non_monotone(x, p)]), + ) + first = float(rf(0.0, LO, HI, 1.0)) + assert rf.stats()["return_status"] == "success" + iterations = rf.stats()["iter_count"] + + # Same inputs: the root comes back from the cache, so the iteration count + # does not move and the status says so. + assert float(rf(0.0, LO, HI, 1.0)) == first + assert rf.stats()["return_status"] == "success (cached)" + assert rf.stats()["iter_count"] == iterations + + # A different target has to be solved afresh. + rf(0.0, LO, HI, 1.4) + assert rf.stats()["return_status"] == "success" + + def test_generated_c_carries_the_cache(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + nested = self._nested() + nested.generate("nested.c", {"with_header": False}) + source = (tmp_path / "nested.c").read_text() + # The generated iteration keeps its own copy of the cache, so an + # ahead-of-time compiled function does not regress to re-solving. + assert "brent_cache" in source + assert "CASADI_BRENT_TLS" in source + + external = casadi.external( + "nested", casadi.Importer(str(tmp_path / "nested.c"), "shell") + ) + for target in (0.7, 1.0, 1.4): + np.testing.assert_allclose( + float(external(target)), float(nested(target)), rtol=0, atol=1e-14 + ) + + +class TestIterationLimit: + """Running out of iterations is a failure, not a root. + + Exhausting ``max_iter`` leaves an arbitrary point inside the bracket. + """ + + @staticmethod + def _steep(): + """``exp(x) - 100``, whose root at ``log(100)`` takes many iterations to reach.""" + x, lo, hi, p = (casadi.MX.sym(n) for n in ("x", "lo", "hi", "p")) + return casadi.Function("g", [x, lo, hi, p], [casadi.exp(x) - 100]) + + def test_exhausting_the_iterations_is_not_a_success(self): + rootfinder = casadi.rootfinder( + "rf", "brent", self._steep(), {"max_iter": 1, "error_on_fail": False} + ) + rootfinder(0.0, 0.0, 10.0, 0.0) + assert rootfinder.stats()["return_status"] == ( + "iteration limit reached without converging" + ) + + def test_the_same_problem_converges_when_given_the_iterations(self): + rootfinder = casadi.rootfinder("rf", "brent", self._steep(), {"max_iter": 200}) + root = float(rootfinder(0.0, 0.0, 10.0, 0.0)) + assert rootfinder.stats()["return_status"] == "success" + np.testing.assert_allclose(root, np.log(100), rtol=1e-12) + + def test_an_exhausted_solve_is_not_cached(self): + rootfinder = casadi.rootfinder( + "rf", "brent", self._steep(), {"max_iter": 1, "error_on_fail": False} + ) + rootfinder(0.0, 0.0, 10.0, 0.0) + rootfinder(0.0, 0.0, 10.0, 0.0) + # A cached hit would report "success (cached)" and hand back the non-root. + assert rootfinder.stats()["return_status"] == ( + "iteration limit reached without converging" + ) + + def test_the_generated_c_also_refuses_to_converge(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + rootfinder = casadi.rootfinder( + "rf", "brent", self._steep(), {"max_iter": 1, "error_on_fail": False} + ) + # `p` stays symbolic: an all-constant call is folded away at build time and + # never reaches the generated C. + target = casadi.MX.sym("p") + wrapped = casadi.Function( + "wrapped", [target], [rootfinder(0.0, 0.0, 10.0, target)] + ) + wrapped.generate("wrapped.c", {"with_header": False}) + external = casadi.external( + "wrapped", casadi.Importer(str(tmp_path / "wrapped.c"), "shell") + ) + with pytest.raises(RuntimeError): + external(0.0) From 9d08775eddd0307a16fc5fb45d1d5e85f8177d01 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:59:47 -0400 Subject: [PATCH 02/12] docs: point the changelog at the pull request --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b4d96be756..2ef2936e21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Features -- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. ([#TBD](https://github.com/pybamm-team/PyBaMM/pull/TBD)) +- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. ([#5729](https://github.com/pybamm-team/PyBaMM/pull/5729)) - Added unstructured mesh support (`UnstructuredSubMesh`, generators, and interface coupling) for arbitrary 2D/3D domains. Hexahedra must have planar faces (warped hexes raise a `GeometryError`), and `UserSuppliedUnstructuredMesh` accepts tetrahedral, triangular, and quadrilateral cells only. ([#5687](https://github.com/pybamm-team/PyBaMM/pull/5687)) - Generalised `VectorField` to N components and added `Component`/`Norm` operators for multi-dimensional vector fields. ([#5686](https://github.com/pybamm-team/PyBaMM/pull/5686)) - Removed the left sidebar from the documentation home page for a cleaner landing experience. ([#5699](https://github.com/pybamm-team/PyBaMM/pull/5699)) From d925888df548adbfb9e15c2a2fae8662a6419ee7 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:14:23 -0400 Subject: [PATCH 03/12] test: drop the scipy cross-check from the brent sweep test cibuildwheel's test env installs only pybammsolvers[dev], which has no scipy, so the module-level `from scipy.optimize import brentq` failed collection and took every wheel job down with exit code 2. The residual assertion already pins the root, and which root a bracketed solve lands on for the non-monotone case was never a contract, so the brentq comparison bought nothing that survived being unrunnable. Co-Authored-By: Claude Opus 5 --- packages/pybammsolvers/tests/test_brent_rootfinder.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/packages/pybammsolvers/tests/test_brent_rootfinder.py b/packages/pybammsolvers/tests/test_brent_rootfinder.py index 26dc808669..a823d2406a 100644 --- a/packages/pybammsolvers/tests/test_brent_rootfinder.py +++ b/packages/pybammsolvers/tests/test_brent_rootfinder.py @@ -5,7 +5,6 @@ import casadi import numpy as np import pytest -from scipy.optimize import brentq import pybammsolvers.idaklu # noqa: F401 registers the plugin on import @@ -40,18 +39,15 @@ def test_solves_a_scalar_equation(self): @pytest.mark.parametrize( "expr_fn", [lambda x, p: casadi.exp(x) + p * x - 2.0, non_monotone] ) - def test_matches_scipy_over_a_sweep(self, expr_fn): + def test_converges_across_a_sweep(self, expr_fn): rf, f = _solver(expr_fn) - worst_difference = worst_residual = 0.0 + worst_residual = 0.0 for p in np.linspace(0.2, 1.8, 40): if float(f(LO, p)) * float(f(HI, p)) > 0: - continue # no bracket for this p, nothing to compare + continue # no bracket for this p, nothing to solve got = float(rf(p)) - want = brentq(lambda x, p=p: float(f(x, p)), LO, HI, xtol=2e-12) assert LO <= got <= HI - worst_difference = max(worst_difference, abs(got - want)) worst_residual = max(worst_residual, abs(float(f(got, p)))) - assert worst_difference < 1e-9 assert worst_residual < 1e-12 def test_bracket_can_be_a_graph_input(self): From 2cefac90d603bb368f7276ccc09bd42bd255240c Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:43:13 -0400 Subject: [PATCH 04/12] fix: clear the cppcheck findings on the brent plugin Codacy runs cppcheck over each file in the diff on its own, so brent.hpp is analysed without brent.cpp and all seven of its members come back unused. Suppress the check for the header rather than annotating members one by one. The variableScope finding is real: `tol`, `xm`, `step` and the interpolation scratch each live for a single iteration, so declare them in the block that uses them. One declaration block per scope still satisfies the C the text is stringified into. Co-Authored-By: Claude Opus 5 --- .../src/pybammsolvers/idaklu_source/brent.hpp | 3 +++ .../src/pybammsolvers/idaklu_source/brent_impl.hpp | 10 +++++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp index d9ba19b172..a2926d6fd8 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp @@ -17,6 +17,9 @@ // // The export macro below is the fifth: in tree it comes from a generated header. // +// The members below are all read from brent.cpp, but cppcheck also analyses this +// header on its own, where nothing uses them, so it calls every one of them unused. +// cppcheck-suppress-file unusedStructMember #ifndef PYBAMM_BRENT_HPP #define PYBAMM_BRENT_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp index cd7dc02f2a..6e035ad1fd 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp @@ -3,8 +3,9 @@ // that `CodeGenerator::sanitize_source` turns into plain C). // // This text is both compiled and stringified into the generated C, so keep it valid C -// once the template header and comments are stripped: declarations first, no C++-only -// constructs, no preprocessor directives (`sanitize_source` drops `#define`/`#undef`). +// once the template header and comments are stripped: declarations at the top of each +// block, no C++-only constructs, no preprocessor directives (`sanitize_source` drops +// `#define`/`#undef`). // // `res_fn` evaluates the residual at `x` into `*fx` and returns nonzero on failure. // Returns 0 on success (`*out` is the root, `*iter` the iteration count), 1 if the @@ -17,7 +18,7 @@ template static int casadi_brent(int (*res_fn)(void*, T1, T1*), void* user_data, T1 a, T1 b, T1 abstol, casadi_int max_iter, T1* out, casadi_int* iter) { - T1 fa, fb, c, fc, d, e, tol, xm, p, q, s, qq, r, step; + T1 fa, fb, c, fc, d, e; casadi_int k; *iter = 0; if (res_fn(user_data, a, &fa)) return 1; @@ -30,6 +31,7 @@ static int casadi_brent(int (*res_fn)(void*, T1, T1*), void* user_data, d = b - a; e = d; for (k = 0; k < max_iter; ++k) { + T1 tol, xm, step; if (fb * fc > 0) { c = a; fc = fa; @@ -50,11 +52,13 @@ static int casadi_brent(int (*res_fn)(void*, T1, T1*), void* user_data, if (fabs(xm) <= tol || fb == 0) break; if (fabs(e) >= tol && fabs(fa) > fabs(fb)) { // Inverse quadratic interpolation, or secant when only two points are distinct + T1 p, q, s; s = fb / fa; if (a == c) { p = 2 * xm * s; q = 1 - s; } else { + T1 qq, r; qq = fa / fc; r = fb / fc; p = s * (2 * xm * qq * (qq - r) - (b - a) * (r - 1)); From b2b403b189c5cf821d0111f046cb696fa2bd901d Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:29:21 -0400 Subject: [PATCH 05/12] fix: stage CasADi's options.hpp so the plugin builds on Windows The staged internal headers include `options.hpp` unqualified. When that does not resolve next to the file asking for it, MSVC widens the search to the include stack, reaches `idaklu_source/`, and answers with our own Options.hpp -- equal to `options.hpp` on a case-insensitive filesystem, and no declaration of `casadi::Options` in sight. Every internal header that names `Options` then fails to parse, starting at function_internal.hpp:117. Staging CasADi's copy puts it in the directory the including file is in, which is the first place either compiler looks, so the collision never comes up. The guard is shared with the installed copy, so a translation unit that pulls in both still sees one definition, and the config.h version check already keeps the two identical. Co-Authored-By: Claude Opus 5 --- packages/pybammsolvers/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/pybammsolvers/CMakeLists.txt b/packages/pybammsolvers/CMakeLists.txt index ea7a620111..61ef03d59f 100644 --- a/packages/pybammsolvers/CMakeLists.txt +++ b/packages/pybammsolvers/CMakeLists.txt @@ -310,9 +310,13 @@ else () message(STATUS "idaklu: \"brent\" plugin verified against CasADi ${CASADI_LINKED_VERSION}") endif () +# options.hpp is public, but stage it too: the headers below include it unqualified, and +# MSVC then searches the include stack, where our own case-insensitively equal Options.hpp +# would answer for it. set(_casadi_internal_headers casadi_os.hpp function_internal.hpp + options.hpp oracle_function.hpp plugin_interface.hpp rootfinder_impl.hpp) From 2138f27c725a33c76b9dfd8454781b45f2f1c88f Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Thu, 20 Aug 2026 13:27:37 -0400 Subject: [PATCH 06/12] fix: treat the brent plugin as unavailable on Windows The plugin registers into whichever CasADi is linked into the pybammsolvers extension. On Linux and macOS that is the casadi wheel's own libcasadi, so the registration is visible to the CasADi Python calls. On Windows the extension is built with MSVC against vcpkg's CasADi while the casadi wheel is built with MinGW -- it ships libstdc++-6.dll and .dll.a import libraries -- so the two hold separate copies of CasADi, and separate plugin registries. Registering ours cannot reach Python's, which falls back to loading libcasadi_rootfinder_brent.dll and fails with WIN32 error 126. Closing that gap needs the plugin built with MinGW against the wheel's libcasadi.dll.a, a second toolchain in the Windows job. Until then it is unavailable there, so skip both test modules on win32 -- 23 in pybammsolvers, 17 in pybamm -- and raise a NotImplementedError naming the cause when a Brent node is converted to CasADi without the plugin, in place of CasADi's "Plugin 'brent' is not found". evaluate() is unaffected: it solves with SciPy. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 2 +- packages/pybamm/src/pybamm/expression_tree/brent.py | 12 ++++++++++++ .../tests/unit/test_expression_tree/test_brent.py | 8 ++++++++ .../pybammsolvers/tests/test_brent_rootfinder.py | 9 +++++++++ 4 files changed, 30 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ef2936e21..a58841a9cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Features -- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. ([#5729](https://github.com/pybamm-team/PyBaMM/pull/5729)) +- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. Not available on Windows, where the CasADi wheel is built with MinGW and `pybammsolvers` with MSVC. ([#5729](https://github.com/pybamm-team/PyBaMM/pull/5729)) - Added unstructured mesh support (`UnstructuredSubMesh`, generators, and interface coupling) for arbitrary 2D/3D domains. Hexahedra must have planar faces (warped hexes raise a `GeometryError`), and `UserSuppliedUnstructuredMesh` accepts tetrahedral, triangular, and quadrilateral cells only. ([#5687](https://github.com/pybamm-team/PyBaMM/pull/5687)) - Generalised `VectorField` to N components and added `Component`/`Norm` operators for multi-dimensional vector fields. ([#5686](https://github.com/pybamm-team/PyBaMM/pull/5686)) - Removed the left sidebar from the documentation home page for a cleaner landing experience. ([#5699](https://github.com/pybamm-team/PyBaMM/pull/5699)) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py index 9d2ace2f5e..3e7f04f024 100644 --- a/packages/pybamm/src/pybamm/expression_tree/brent.py +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -103,6 +103,10 @@ class Brent(pybamm.Symbol): plugin registered by ``pybammsolvers``, so the whole solve runs inside the CasADi graph. + Not available on Windows: the CasADi wheel there is built with MinGW and + ``pybammsolvers`` with MSVC, so the two hold separate copies of CasADi and the + plugin never reaches the one Python calls. ``evaluate()`` still works, via SciPy. + Brent needs only a sign change over the bounds, so it converges where a Newton iteration stalls, and the answer cannot leave them. Derivatives come from CasADi's implicit function theorem, exactly. @@ -217,6 +221,14 @@ def _from_json(cls, snippet: dict): ) def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): + if not casadi.has_rootfinder("brent"): + raise NotImplementedError( + "the 'brent' rootfinder plugin is not registered with the CasADi that " + "Python is using, so this node cannot be converted to CasADi. On " + "Windows the casadi wheel is built with MinGW and pybammsolvers with " + "MSVC, which leaves the two with separate copies of CasADi. Call " + "evaluate() instead, which solves with SciPy." + ) unknown = casadi.MX.sym(f"brent_unknown_{abs(self.id)}") cache = _OracleCache( casadi_symbols, _nodes_reading(self.residual, self.unknown) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py index c1bdc50f41..e0c2cb2b83 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -24,6 +24,14 @@ def _evaluate(symbol, **inputs): return float(function(*[inputs[name] for name in symbols])) +# On Windows our MSVC-built extension and the MinGW-built casadi wheel hold separate +# copies of CasADi, so a plugin registered in ours is invisible to the one Python calls. +pytestmark = pytest.mark.skipif( + sys.platform == "win32", + reason="the brent plugin cannot reach the casadi wheel's CasADi on Windows", +) + + class TestBrent: def test_solves_a_scalar_equation(self): x = pybamm.BrentUnknown("x") diff --git a/packages/pybammsolvers/tests/test_brent_rootfinder.py b/packages/pybammsolvers/tests/test_brent_rootfinder.py index a823d2406a..866e9f4f7b 100644 --- a/packages/pybammsolvers/tests/test_brent_rootfinder.py +++ b/packages/pybammsolvers/tests/test_brent_rootfinder.py @@ -2,12 +2,21 @@ from __future__ import annotations +import sys + import casadi import numpy as np import pytest import pybammsolvers.idaklu # noqa: F401 registers the plugin on import +# On Windows this wheel is MSVC-built and the casadi wheel MinGW-built, so the two hold +# separate copies of CasADi and the plugin registered here never reaches Python's. +pytestmark = pytest.mark.skipif( + sys.platform == "win32", + reason="the brent plugin cannot reach the casadi wheel's CasADi on Windows", +) + LO, HI = 1e-9, 1 - 1e-9 From 55f1d8a525d4ef796625aa3e38e73194998c0495 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:11:19 -0400 Subject: [PATCH 07/12] fix: correct three brent defects and withdraw its public name Three defects, none of which depend on the API shape: The cache scratch vector lived on the solver object, so two threads evaluating one Function shared it. Moved onto BrentMemory, which CasADi hands out per concurrent evaluation. The bracket test read `fa * fb <= 0`, and two residuals near the underflow limit multiply to zero, which reads as a sign change. Now tested by sign, with NaN handled explicitly, so a bracket holding no root is reported as one. The NumPy path returned NaN where the plugin raises: a blanket `except (ValueError, RuntimeError)` around brentq swallowed both an empty bracket and a failure to converge along with the tree probes it was meant to absorb. Those two now raise SolverError, and only a probe returns NaN. `Brent` and `BrentUnknown` also lose their public names, becoming `_Brent` and `_BrentUnknown` alongside `_BaseAverage`, and the docs page goes. The node has no in-tree consumer yet, and lifting the residual into a closed sub-expression -- which is where this is headed, and what the Rust tape backend in #5732 wants too -- will change the constructor. Cheaper to withdraw the commitment now than to break it later. The changelog bullet goes with it: nothing user-facing is left to announce. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 1 - docs/source/api/expression_tree/brent.rst | 8 -- docs/source/api/expression_tree/index.rst | 1 - packages/pybamm/src/pybamm/__init__.py | 2 +- .../pybamm/discretisations/discretisation.py | 4 +- .../src/pybamm/expression_tree/brent.py | 61 ++++++++---- .../parameters/parameter_substitutor.py | 2 +- packages/pybamm/tests/strategies/symbols.py | 14 +-- .../unit/test_expression_tree/test_brent.py | 96 ++++++++++--------- .../src/pybammsolvers/idaklu_source/brent.cpp | 8 +- .../src/pybammsolvers/idaklu_source/brent.hpp | 4 +- .../idaklu_source/brent_impl.hpp | 7 +- .../tests/test_brent_rootfinder.py | 17 ++++ 13 files changed, 133 insertions(+), 92 deletions(-) delete mode 100644 docs/source/api/expression_tree/brent.rst diff --git a/CHANGELOG.md b/CHANGELOG.md index a58841a9cf..a17810776c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,6 @@ ## Features -- Added `pybamm.Brent`, a bracketed scalar rootfind backed by a native CasADi plugin. Not available on Windows, where the CasADi wheel is built with MinGW and `pybammsolvers` with MSVC. ([#5729](https://github.com/pybamm-team/PyBaMM/pull/5729)) - Added unstructured mesh support (`UnstructuredSubMesh`, generators, and interface coupling) for arbitrary 2D/3D domains. Hexahedra must have planar faces (warped hexes raise a `GeometryError`), and `UserSuppliedUnstructuredMesh` accepts tetrahedral, triangular, and quadrilateral cells only. ([#5687](https://github.com/pybamm-team/PyBaMM/pull/5687)) - Generalised `VectorField` to N components and added `Component`/`Norm` operators for multi-dimensional vector fields. ([#5686](https://github.com/pybamm-team/PyBaMM/pull/5686)) - Removed the left sidebar from the documentation home page for a cleaner landing experience. ([#5699](https://github.com/pybamm-team/PyBaMM/pull/5699)) diff --git a/docs/source/api/expression_tree/brent.rst b/docs/source/api/expression_tree/brent.rst deleted file mode 100644 index 05cf147a23..0000000000 --- a/docs/source/api/expression_tree/brent.rst +++ /dev/null @@ -1,8 +0,0 @@ -Brent -===== - -.. autoclass:: pybamm.Brent - :members: - -.. autoclass:: pybamm.BrentUnknown - :members: diff --git a/docs/source/api/expression_tree/index.rst b/docs/source/api/expression_tree/index.rst index de5e5abfbb..0a6f3d757c 100644 --- a/docs/source/api/expression_tree/index.rst +++ b/docs/source/api/expression_tree/index.rst @@ -19,5 +19,4 @@ Expression Tree functions input_parameter interpolant - brent operations/index diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index 02cf7665d5..740fca3e26 100644 --- a/packages/pybamm/src/pybamm/__init__.py +++ b/packages/pybamm/src/pybamm/__init__.py @@ -32,7 +32,7 @@ from .expression_tree.broadcasts import * from .expression_tree.functions import * from .expression_tree.conditional import Conditional -from .expression_tree.brent import Brent, BrentUnknown +from .expression_tree.brent import _Brent, _BrentUnknown from .expression_tree.interpolant import Interpolant from .expression_tree.discrete_time_sum import * from .expression_tree.input_parameter import InputParameter diff --git a/packages/pybamm/src/pybamm/discretisations/discretisation.py b/packages/pybamm/src/pybamm/discretisations/discretisation.py index 5fb30ab815..722904c501 100644 --- a/packages/pybamm/src/pybamm/discretisations/discretisation.py +++ b/packages/pybamm/src/pybamm/discretisations/discretisation.py @@ -1148,11 +1148,11 @@ def _process_symbol(self, symbol): else: return symbol.create_copy(new_children=[disc_child]) - elif isinstance(symbol, pybamm.BrentUnknown): + elif isinstance(symbol, pybamm._BrentUnknown): # bound by its Brent, so it has no state-vector slice to resolve return symbol.create_copy() - elif isinstance(symbol, (pybamm.Function, pybamm.Conditional, pybamm.Brent)): + elif isinstance(symbol, (pybamm.Function, pybamm.Conditional, pybamm._Brent)): disc_children = [self.process_symbol(child) for child in symbol.children] return symbol.create_copy(disc_children) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py index 3e7f04f024..5d786e2a4f 100644 --- a/packages/pybamm/src/pybamm/expression_tree/brent.py +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -9,17 +9,17 @@ import pybamm -class BrentUnknown(pybamm.Symbol): +class _BrentUnknown(pybamm.Symbol): """ - The scalar a :class:`Brent` solves for. + The scalar a :class:`_Brent` solves for. Bound by the rootfinder, not by the model, so it is deliberately not a :class:`pybamm.Variable`: the checks that enumerate model states must not count it - as one, or a ``Brent`` inside ``model.rhs`` or ``model.algebraic`` looks like an + as one, or a ``_Brent`` inside ``model.rhs`` or ``model.algebraic`` looks like an extra unknown with no equation. Two unknowns of the same name are the same unknown, as for any other symbol. Give - each one its own name: a ``Brent`` nested inside another whose unknown shares its + each one its own name: a ``_Brent`` nested inside another whose unknown shares its name shadows the outer binding, and CasADi rejects the oracle it builds. Parameters @@ -70,7 +70,7 @@ def _nodes_reading(root: pybamm.Symbol, unknown: pybamm.Symbol) -> set: class _OracleCache(dict): - """Conversion cache that keeps a :class:`Brent`'s own binding out of the shared one. + """Conversion cache that keeps a :class:`_Brent`'s own binding out of the shared one. Keys in ``local`` are held here and discarded with the oracle; the rest is written through to ``shared``, so the graph a rootfind shares with its surroundings is @@ -93,7 +93,7 @@ def __setitem__(self, key, value): self._shared[key] = value -class Brent(pybamm.Symbol): +class _Brent(pybamm.Symbol): """ Solve ``residual == 0`` for ``unknown`` within ``bounds``, by Brent's method. @@ -116,7 +116,7 @@ class Brent(pybamm.Symbol): residual : :class:`pybamm.Symbol` The expression to drive to zero. Must contain ``unknown``. To invert ``f`` at a target, pass ``f - target``. - unknown : :class:`pybamm.BrentUnknown` + unknown : :class:`_BrentUnknown` The value being solved for. Must appear in ``residual`` and nowhere else in the surrounding expression. bounds : tuple @@ -134,14 +134,14 @@ class Brent(pybamm.Symbol): .. code-block:: python # invert an open-circuit potential at a given voltage - sto = pybamm.BrentUnknown("stoichiometry") - node = pybamm.Brent(param.n.prim.U(sto, T) - voltage, sto, (0, 1)) + sto = pybamm._BrentUnknown("stoichiometry") + node = pybamm._Brent(param.n.prim.U(sto, T) - voltage, sto, (0, 1)) """ def __init__( self, residual: pybamm.Symbol, - unknown: BrentUnknown, + unknown: _BrentUnknown, bounds: tuple, *, abstol: float = 1e-14, @@ -152,9 +152,9 @@ def __init__( raise TypeError( f"residual must be a pybamm.Symbol, got {type(residual).__name__}" ) - if not isinstance(unknown, BrentUnknown): + if not isinstance(unknown, _BrentUnknown): raise TypeError( - f"unknown must be a pybamm.BrentUnknown, got {type(unknown).__name__}" + f"unknown must be a _BrentUnknown, got {type(unknown).__name__}" ) if not any(node == unknown for node in residual.pre_order()): raise pybamm.ModelError(f"'{unknown}' does not appear in '{residual}'") @@ -182,7 +182,7 @@ def residual(self) -> pybamm.Symbol: return self.children[0] @property - def unknown(self) -> BrentUnknown: + def unknown(self) -> _BrentUnknown: return self.children[1] @property @@ -191,7 +191,7 @@ def bounds(self) -> tuple[pybamm.Symbol, pybamm.Symbol]: def create_copy(self, new_children=None, perform_simplifications=True): residual, unknown, lo, hi = self._children_for_copying(new_children) - return Brent( + return _Brent( residual, unknown, (lo, hi), @@ -269,9 +269,9 @@ def scalar(value): lo, hi = (scalar(b.evaluate(t, y, y_dot, inputs)) for b in self.bounds) unknown = self.unknown - if not isinstance(unknown, BrentUnknown): + if not isinstance(unknown, _BrentUnknown): raise TypeError( - "NumPy evaluation needs a pybamm.BrentUnknown as the unknown, got " + "NumPy evaluation needs a _BrentUnknown as the unknown, got " f"{type(unknown).__name__}" ) @@ -284,20 +284,43 @@ def g(value): return np.nan * np.ones((1, 1)) try: - root = brentq(g, lo, hi, xtol=self.abstol, maxiter=self.max_iter) + residual_lo, residual_hi = g(lo), g(hi) except (ValueError, RuntimeError): # a probe walks the tree before the parameters are in place return np.nan * np.ones((1, 1)) finally: unknown._value = np.nan * np.ones((1, 1)) + + if not (np.isfinite(residual_lo) and np.isfinite(residual_hi)): + return np.nan * np.ones((1, 1)) + + # the same two failures the plugin reports, rather than a quiet NaN + if ( + residual_lo != 0 + and residual_hi != 0 + and (residual_lo > 0) == (residual_hi > 0) + ): + raise pybamm.SolverError( + f"no sign change over the bracket ({lo}, {hi}), where the residual is " + f"{residual_lo} and {residual_hi}, so it holds no root" + ) + + try: + root = brentq(g, lo, hi, xtol=self.abstol, maxiter=self.max_iter) + except RuntimeError as error: + raise pybamm.SolverError( + f"the rootfind did not converge in {self.max_iter} iterations" + ) from error + finally: + unknown._value = np.nan * np.ones((1, 1)) return np.array([[root]]) def diff(self, variable): raise NotImplementedError( - "Brent has no symbolic derivative; use convert_to_format='casadi'." + "_Brent has no symbolic derivative; use convert_to_format='casadi'." ) def _jac(self, variable): raise NotImplementedError( - "Brent has no symbolic jacobian; use convert_to_format='casadi'." + "_Brent has no symbolic jacobian; use convert_to_format='casadi'." ) diff --git a/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py b/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py index f9df3b70d0..c163a8508a 100644 --- a/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py +++ b/packages/pybamm/src/pybamm/parameters/parameter_substitutor.py @@ -256,7 +256,7 @@ def _process_symbol(self, symbol: pybamm.Symbol) -> pybamm.Symbol: pybamm.Concatenation, pybamm.BinaryOperator, pybamm.Conditional, - pybamm.Brent, + pybamm._Brent, ), ): new_children = [self.process_symbol(child) for child in symbol.children] diff --git a/packages/pybamm/tests/strategies/symbols.py b/packages/pybamm/tests/strategies/symbols.py index b57ee6037f..a1e296cd96 100644 --- a/packages/pybamm/tests/strategies/symbols.py +++ b/packages/pybamm/tests/strategies/symbols.py @@ -445,14 +445,14 @@ def _reg_power_branch( def _brent_branch( child_strategy: st.SearchStrategy[pybamm.Symbol], -) -> st.SearchStrategy[pybamm.Brent]: +) -> st.SearchStrategy[pybamm._Brent]: """Brent over ``residual(unknown) == 0``, bracketed by two further expressions.""" def make_brent(parts): residual, lo, hi, abstol, max_iter, tag = parts # nested Brents must not share an unknown, and a drawn tag replays identically - unknown = pybamm.BrentUnknown(f"brent unknown {tag}") - return pybamm.Brent( + unknown = pybamm._BrentUnknown(f"brent unknown {tag}") + return pybamm._Brent( # subtraction so the unknown cannot be simplified away, as `0 * x` would unknown - residual, unknown, @@ -898,9 +898,9 @@ def _vector_branch( pybamm.Parameter: lambda _children: parameter_strategy(), pybamm.Time: lambda _children: time_strategy(), pybamm.InputParameter: lambda _children: input_parameter_strategy(), - pybamm.BrentUnknown: lambda _children: st.integers( + pybamm._BrentUnknown: lambda _children: st.integers( min_value=0, max_value=2**40 - ).map(lambda tag: pybamm.BrentUnknown(f"brent unknown {tag}")), + ).map(lambda tag: pybamm._BrentUnknown(f"brent unknown {tag}")), pybamm.Negate: lambda children: _unary_branch(children, pybamm.Negate), pybamm.AbsoluteValue: lambda children: _unary_branch( children, pybamm.AbsoluteValue @@ -969,7 +969,7 @@ def _vector_branch( pybamm.RegPower: _reg_power_branch, # data-bearing leaves pybamm.Interpolant: _interpolant_branch, - pybamm.Brent: _brent_branch, + pybamm._Brent: _brent_branch, ExpressionFunctionParameter: _expression_function_parameter_branch, # n-ary / complex branch strategies pybamm.Conditional: _conditional_branch, @@ -1097,7 +1097,7 @@ def _vector_branch( pybamm.Parameter, pybamm.Time, pybamm.InputParameter, - pybamm.BrentUnknown, + pybamm._BrentUnknown, } ) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py index e0c2cb2b83..8b6e3397fd 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -34,22 +34,22 @@ def _evaluate(symbol, **inputs): class TestBrent: def test_solves_a_scalar_equation(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(pybamm.exp(x) + x - 2.0, x, (-5.0, 5.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(pybamm.exp(x) + x - 2.0, x, (-5.0, 5.0)) got = float(casadi.evalf(node.to_casadi(inputs={}))) want = brentq(lambda v: np.exp(v) + v - 2.0, -5.0, 5.0, xtol=2e-12) assert got == pytest.approx(want, abs=1e-12) def test_target_may_be_an_input_parameter(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) assert _evaluate(node, target=9.0) == pytest.approx(3.0, abs=1e-12) assert _evaluate(node, target=4.0) == pytest.approx(2.0, abs=1e-12) def test_bracket_may_be_input_parameters(self): # x^2 = 6 has roots at +-sqrt(6); the bracket selects one, at solve time - x = pybamm.BrentUnknown("x") - node = pybamm.Brent( + x = pybamm._BrentUnknown("x") + node = pybamm._Brent( x * x - 6.0, x, (pybamm.InputParameter("lo"), pybamm.InputParameter("hi")) ) assert _evaluate(node, lo=0.0, hi=10.0) == pytest.approx(np.sqrt(6), abs=1e-12) @@ -59,14 +59,14 @@ def test_bracket_may_be_input_parameters(self): def test_the_expression_may_contain_input_parameters(self): # a x^2 = 9 has positive root 3 / sqrt(a) - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(pybamm.InputParameter("a") * x * x - 9.0, x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(pybamm.InputParameter("a") * x * x - 9.0, x, (0.0, 10.0)) for a in (1.0, 4.0, 9.0): assert _evaluate(node, a=a) == pytest.approx(3.0 / np.sqrt(a), abs=1e-12) def test_every_argument_may_be_an_input_parameter_at_once(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent( + x = pybamm._BrentUnknown("x") + node = pybamm._Brent( pybamm.InputParameter("a") * x * x - pybamm.InputParameter("target"), x, (pybamm.InputParameter("lo"), pybamm.InputParameter("hi")), @@ -78,8 +78,8 @@ def test_every_argument_may_be_an_input_parameter_at_once(self): def test_solves_over_the_state_vector(self): state = pybamm.StateVector(slice(0, 1)) - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * state - 6.0, x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * state - 6.0, x, (0.0, 10.0)) y = casadi.MX.sym("y", 1) expression = casadi.Function("f", [y], [node.to_casadi(y=y, inputs={})]) assert float(expression(2.0)) == pytest.approx(3.0, abs=1e-12) @@ -87,33 +87,33 @@ def test_solves_over_the_state_vector(self): def test_derivative_is_exact(self): # x = sqrt(target), so dx/d(target) = 1 / (2 sqrt(target)) - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - pybamm.InputParameter("target"), x, (0.0, 10.0)) symbol = casadi.MX.sym("target") root = node.to_casadi(inputs={"target": symbol}) derivative = casadi.Function("J", [symbol], [casadi.jacobian(root, symbol)]) assert float(derivative(9.0)) == pytest.approx(1 / 6, rel=1e-12) def test_composes_into_a_larger_expression(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0)) got = float(casadi.evalf((3 * node + pybamm.Scalar(1)).to_casadi(inputs={}))) assert got == pytest.approx(10.0, abs=1e-12) def test_nests(self): # the inner solve gives sqrt(16) = 4, so the outer gives sqrt(4) = 2 - inner_x = pybamm.BrentUnknown("inner") - inner = pybamm.Brent(inner_x * inner_x - 16.0, inner_x, (0.0, 10.0)) - outer_x = pybamm.BrentUnknown("outer") - outer = pybamm.Brent(outer_x * outer_x - inner, outer_x, (0.0, 10.0)) + inner_x = pybamm._BrentUnknown("inner") + inner = pybamm._Brent(inner_x * inner_x - 16.0, inner_x, (0.0, 10.0)) + outer_x = pybamm._BrentUnknown("outer") + outer = pybamm._Brent(outer_x * outer_x - inner, outer_x, (0.0, 10.0)) assert float(casadi.evalf(outer.to_casadi(inputs={}))) == pytest.approx(2.0) def test_evaluating_does_not_re_enter_python(self): # the whole solve runs in the CasADi graph, so a Brent node must cost no more # python frames per evaluation than the same expression without one state = pybamm.StateVector(slice(3, 4)) - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(pybamm.exp(x) + x * state - 2.0, x, (-5.0, 5.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(pybamm.exp(x) + x * state - 2.0, x, (-5.0, 5.0)) y = casadi.MX.sym("y", 500) with_brent = casadi.Function("a", [y], [3 * node.to_casadi(y=y, inputs={}) + 1]) without = casadi.Function("b", [y], [3 * casadi.exp(y[3]) + 1]) @@ -141,53 +141,61 @@ def profile(frame, event, arg): def test_the_oracle_only_reads_what_the_residual_needs(self): # a residual that ignores time must not drag time into the solve state = pybamm.StateVector(slice(0, 1)) - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * state - 6.0, x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * state - 6.0, x, (0.0, 10.0)) t, y = casadi.MX.sym("t"), casadi.MX.sym("y", 1) names = [s.name() for s in casadi.symvar(node.to_casadi(t=t, y=y, inputs={}))] assert names == ["y"] def test_no_sign_change_fails_rather_than_guessing(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x + 1 - 0.0, x, (0.0, 1.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x + 1 - 0.0, x, (0.0, 1.0)) with pytest.raises(RuntimeError, match="rootfinder process failed"): casadi.evalf(node.to_casadi(inputs={})) + def test_evaluate_reports_an_empty_bracket_rather_than_nan(self): + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x + 1.0, x, (0.0, 1.0)) + with pytest.raises(pybamm.SolverError, match="no sign change over the bracket"): + node.evaluate() + def test_children_and_copy(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0)) assert len(node.children) == 4 copy = node.create_copy() assert copy.name == node.name assert float(casadi.evalf(copy.to_casadi(inputs={}))) == pytest.approx(3.0) def test_errors(self): - x = pybamm.BrentUnknown("x") - with pytest.raises(TypeError, match=r"unknown must be a pybamm\.BrentUnknown"): - pybamm.Brent(x * x - 9.0, 1.0, (0, 1)) + x = pybamm._BrentUnknown("x") + with pytest.raises(TypeError, match=r"unknown must be a _BrentUnknown"): + pybamm._Brent(x * x - 9.0, 1.0, (0, 1)) with pytest.raises(TypeError, match=r"residual must be a pybamm\.Symbol"): - pybamm.Brent(1.0, x, (0, 1)) + pybamm._Brent(1.0, x, (0, 1)) with pytest.raises(pybamm.ModelError, match="does not appear in"): - pybamm.Brent(pybamm.Scalar(2) * pybamm.t - 9.0, x, (0, 1)) + pybamm._Brent(pybamm.Scalar(2) * pybamm.t - 9.0, x, (0, 1)) with pytest.raises(pybamm.ModelError, match="bounds must be a"): - pybamm.Brent(x * x - 9.0, x, (0, 1, 2)) + pybamm._Brent(x * x - 9.0, x, (0, 1, 2)) - node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0)) + node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0)) with pytest.raises(NotImplementedError, match="no symbolic derivative"): node.diff(pybamm.t) with pytest.raises(NotImplementedError, match="no symbolic jacobian"): node._jac(pybamm.t) def test_round_trips_through_json(self): - x = pybamm.BrentUnknown("x") - node = pybamm.Brent(x * x - 9.0, x, (0.0, 10.0), abstol=1e-12, max_iter=42) + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0), abstol=1e-12, max_iter=42) rebuilt = convert_symbol_from_json(convert_symbol_to_json(node)) assert rebuilt.abstol == 1e-12 assert rebuilt.max_iter == 42 assert rebuilt.unknown == node.unknown # the two references in `x * x` and the explicit child are one unknown assert ( - len({s.name for s in rebuilt.pre_order() if type(s) is pybamm.BrentUnknown}) + len( + {s.name for s in rebuilt.pre_order() if type(s) is pybamm._BrentUnknown} + ) == 1 ) assert float(casadi.evalf(rebuilt.to_casadi(inputs={}))) == pytest.approx(3.0) @@ -197,20 +205,20 @@ def test_an_unknown_outside_its_brent_is_an_error(self): with pytest.raises( pybamm.ModelError, match="only has a value inside the Brent" ): - pybamm.BrentUnknown("x").to_casadi(inputs={}) + pybamm._BrentUnknown("x").to_casadi(inputs={}) def test_the_tolerances_are_part_of_the_identity(self): # a Brent converted inside another's oracle is shared with the enclosing # conversion, so two differing only in tolerance must not be served each other - unknown = pybamm.BrentUnknown("s") + unknown = pybamm._BrentUnknown("s") target = pybamm.InputParameter("p") residual = unknown * unknown - target - coarse = pybamm.Brent(residual, unknown, (0.01, 10), abstol=1e-1, max_iter=100) - fine = pybamm.Brent(residual, unknown, (0.01, 10), abstol=1e-14, max_iter=100) + coarse = pybamm._Brent(residual, unknown, (0.01, 10), abstol=1e-1, max_iter=100) + fine = pybamm._Brent(residual, unknown, (0.01, 10), abstol=1e-14, max_iter=100) assert coarse.id != fine.id - outer_unknown = pybamm.BrentUnknown("outer") - outer = pybamm.Brent(outer_unknown - coarse, outer_unknown, (0.01, 50)) + outer_unknown = pybamm._BrentUnknown("outer") + outer = pybamm._Brent(outer_unknown - coarse, outer_unknown, (0.01, 50)) inputs = {"p": casadi.MX.sym("p")} shared: dict = {} converted = [ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp index c8bb5db3be..ac94b4931f 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.cpp @@ -85,12 +85,12 @@ int Brent::solve(void* mem) const { const double b = m->iarg[BRACKET_HI][0]; // the root never depends on the guess, so the guess is left out of the key - key_.clear(); + m->key.clear(); for (casadi_int i = 0; i < n_in_; ++i) { if (i == iin_ || !m->iarg[i]) continue; - key_.insert(key_.end(), m->iarg[i], m->iarg[i] + nnz_in(i)); + m->key.insert(m->key.end(), m->iarg[i], m->iarg[i] + nnz_in(i)); } - if (m->cache_valid && m->cache_key == key_) { + if (m->cache_valid && m->cache_key == m->key) { ++m->cache_hits; if (m->ires[iout_]) m->ires[iout_][0] = m->cache_root; m->return_status = "success (cached)"; @@ -110,7 +110,7 @@ int Brent::solve(void* mem) const { return 0; } - m->cache_key = key_; + m->cache_key = m->key; m->cache_root = root; m->cache_valid = true; diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp index a2926d6fd8..8666d8c844 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp @@ -43,6 +43,9 @@ struct CASADI_ROOTFINDER_BRENT_EXPORT BrentMemory : public RootfinderMemory { // Last solve, keyed on every input but the guess, so a Brent nested inside another // is not re-solved on every iteration of the enclosing one. std::vector cache_key; + // scratch for the key of the solve in flight; per-memory, so two threads + // evaluating one Function do not share it + std::vector key; double cache_root = 0; bool cache_valid = false; casadi_int cache_hits = 0; @@ -78,7 +81,6 @@ class CASADI_ROOTFINDER_BRENT_EXPORT Brent : public Rootfinder { static const std::string meta_doc; void init(const Dict& opts) override; - mutable std::vector key_; // cache scratch int solve(void* mem) const override; void* alloc_mem() const override { return new BrentMemory(); } diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp index 6e035ad1fd..b842b1df48 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp @@ -23,9 +23,10 @@ static int casadi_brent(int (*res_fn)(void*, T1, T1*), void* user_data, *iter = 0; if (res_fn(user_data, a, &fa)) return 1; if (res_fn(user_data, b, &fb)) return 1; - // No sign change, so the bracket contains no root. Report it rather than - // returning whichever end happens to be closer. - if (!(fa * fb <= 0)) return 2; + // No sign change, so the bracket contains no root. Tested by sign rather than + // by product, which underflows to zero for two tiny residuals and fakes one. + if (fa != fa || fb != fb) return 2; + if (fa != 0 && fb != 0 && (fa > 0) == (fb > 0)) return 2; c = a; fc = fa; d = b - a; diff --git a/packages/pybammsolvers/tests/test_brent_rootfinder.py b/packages/pybammsolvers/tests/test_brent_rootfinder.py index 866e9f4f7b..00fdac60e8 100644 --- a/packages/pybammsolvers/tests/test_brent_rootfinder.py +++ b/packages/pybammsolvers/tests/test_brent_rootfinder.py @@ -99,6 +99,23 @@ def test_no_sign_change_fails_rather_than_guessing(self): with pytest.raises(RuntimeError, match="rootfinder process failed"): rf(0.0) + def test_two_tiny_residuals_are_not_a_bracket(self): + # both ends are the same sign, but their product underflows to zero, which a + # `fa * fb <= 0` test reads as a sign change + rf, _ = _solver(lambda x, p: 1e-200 * (x + 1.0), lo=0.0, hi=1.0) + with pytest.raises(RuntimeError, match="rootfinder process failed"): + rf(0.0) + + def test_solves_the_same_function_from_several_threads(self): + # the cache scratch lives on the solver's memory, so two evaluations of one + # Function must not share it + rf, f = _solver(lambda x, p: casadi.exp(x) + p * x - 2.0) + targets = np.linspace(0.2, 1.8, 64) + mapped = rf.map(len(targets), "thread", 8) + roots = np.array(mapped(targets.reshape(1, -1))).reshape(-1) + worst = max(abs(float(f(root, p))) for root, p in zip(roots, targets)) + assert worst < 1e-12 + def test_reports_iteration_count(self): x, lo, hi, p = (casadi.MX.sym(n) for n in ("x", "lo", "hi", "p")) g = casadi.Function("g", [x, lo, hi, p], [casadi.exp(x) + p * x - 2.0]) From 9a292ed4e6fe4030777578c5fd929246ca9b9f28 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:20:43 -0400 Subject: [PATCH 08/12] refactor: read the CasADi internals version from the library being linked The version and its sdist hash were pinned here as well as in pyproject.toml, so a CasADi bump meant editing both. The configure step already reads the linked CasADi's version out of its own config.h to check the two agree, so use that as the source instead: pyproject.toml stays the only pin. CASADI_INTERNAL_VERSION survives as an override, and now disagreeing with the linked library is the only thing that trips the mismatch error. The sdist hash cannot be pinned in tree once the version floats, so it is optional and checked only when passed; a truncated download still fails, at extraction. The staging directory is keyed on the version too. It was not, so a bump would have found the previous version's headers already in place and skipped restaging, compiling against internals from the wrong release. Co-Authored-By: Claude Opus 5 --- packages/pybammsolvers/CMakeLists.txt | 40 +++++++++++++++++++-------- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/packages/pybammsolvers/CMakeLists.txt b/packages/pybammsolvers/CMakeLists.txt index 61ef03d59f..649c7c8312 100644 --- a/packages/pybammsolvers/CMakeLists.txt +++ b/packages/pybammsolvers/CMakeLists.txt @@ -249,11 +249,11 @@ else () endif () # casadi::Rootfinder's headers are not installed and are LGPL-3.0-or-later, so stage -# them from the pinned sdist rather than vendoring them into this BSD-3-Clause tree. -set(CASADI_INTERNAL_VERSION "3.7.2" - CACHE STRING "CasADi version to source the internal headers from") -set(CASADI_INTERNAL_SHA256 "b4d7bd8acdc4180306903ae1c9eddaf41be2a3ae2fa7154c57174ae64acdc60d" - CACHE STRING "SHA-256 of the CasADi source distribution") +# them from the matching sdist rather than vendoring them into this BSD-3-Clause tree. +set(CASADI_INTERNAL_VERSION "" + CACHE STRING "CasADi version to source the internal headers from; read from the linked CasADi when empty") +set(CASADI_INTERNAL_SHA256 "" + CACHE STRING "SHA-256 of the CasADi source distribution, checked on download when set") set(CASADI_SOURCE_DIR "" CACHE PATH "Local CasADi source tree; set this for offline builds to skip the download") @@ -287,7 +287,21 @@ if (CASADI_CONFIG_HEADER) set(_casadi_config_flags "${CMAKE_MATCH_1}") endif () -if (NOT CASADI_LINKED_VERSION) +# Follow the CasADi being linked rather than a version pinned here, so the pin in +# pyproject.toml stays the only one. +set(_casadi_requested_version "${CASADI_INTERNAL_VERSION}") +if (NOT _casadi_requested_version) + set(CASADI_INTERNAL_VERSION "${CASADI_LINKED_VERSION}") +endif () + +if (NOT CASADI_INTERNAL_VERSION) + message(FATAL_ERROR + "Could not find casadi/config.h in any of: ${_casadi_include_dirs}\n" + "The \"brent\" rootfinder plugin subclasses CasADi internals that are not " + "ABI-stable, so the version to source them from has to be known. Pass " + "-DCASADI_INTERNAL_VERSION=, with " + "-DCASADI_ALLOW_UNVERIFIED_INTERNALS=ON to skip the check that it matches.") +elseif (NOT CASADI_LINKED_VERSION) if (CASADI_ALLOW_UNVERIFIED_INTERNALS) message(WARNING "Could not read the linked CasADi version from casadi/config.h. The \"brent\" " @@ -303,9 +317,9 @@ if (NOT CASADI_LINKED_VERSION) elseif (NOT CASADI_LINKED_VERSION STREQUAL CASADI_INTERNAL_VERSION) message(FATAL_ERROR "CasADi version mismatch: linking CasADi ${CASADI_LINKED_VERSION} but sourcing the " - "internal headers for the \"brent\" plugin from ${CASADI_INTERNAL_VERSION}. These " - "headers are not ABI-stable, so they must match. Update CASADI_INTERNAL_VERSION and " - "CASADI_INTERNAL_SHA256 in this file to match the casadi pin in pyproject.toml.") + "internal headers for the \"brent\" plugin from the requested " + "${CASADI_INTERNAL_VERSION}. These headers are not ABI-stable, so they must match; " + "drop -DCASADI_INTERNAL_VERSION to follow the linked CasADi.") else () message(STATUS "idaklu: \"brent\" plugin verified against CasADi ${CASADI_LINKED_VERSION}") endif () @@ -321,7 +335,7 @@ set(_casadi_internal_headers plugin_interface.hpp rootfinder_impl.hpp) set(_casadi_internal_work "${CMAKE_CURRENT_BINARY_DIR}/casadi-internal") -set(_casadi_internal_include "${_casadi_internal_work}/include") +set(_casadi_internal_include "${_casadi_internal_work}/${CASADI_INTERNAL_VERSION}/include") # editable.rebuild reconfigures on every import; skip the work once staged. set(_casadi_internal_staged TRUE) @@ -344,8 +358,12 @@ if (NOT _casadi_internal_staged) set(_casadi_sdist_url "https://pypi.org/packages/source/c/casadi/casadi-${CASADI_INTERNAL_VERSION}.tar.gz") message(STATUS "idaklu: downloading ${_casadi_sdist_url} for the CasADi internal headers") + set(_casadi_expected_hash "") + if (CASADI_INTERNAL_SHA256) + set(_casadi_expected_hash EXPECTED_HASH SHA256=${CASADI_INTERNAL_SHA256}) + endif () file(DOWNLOAD "${_casadi_sdist_url}" "${_casadi_sdist}" - EXPECTED_HASH SHA256=${CASADI_INTERNAL_SHA256} + ${_casadi_expected_hash} STATUS _casadi_download_status) list(GET _casadi_download_status 0 _casadi_download_rc) if (_casadi_download_rc) From 861da37f6c1e2bb440c15c45b63c59ef9a5305f0 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:43:11 -0400 Subject: [PATCH 09/12] fix: keep the process out of the names CasADi serialises The conversion named the unknown, the oracle and the rootfinder after `abs(self.id)`. `Symbol.set_id` hashes a tuple containing a string, so ids are randomised per process, and those names reach `fn.serialize()` -- the AOT compile cache key. Any model holding a rootfind therefore missed the on-disk cache and recompiled from scratch in every process. Measured over three fresh interpreters before: a Brent-free function keyed to 4da87061f2ba3c82 each time, one holding a Brent to 5229b804, 0d83e675 and 277998d7. After: 3e26c07750303134 each time. The identity was never needed. CasADi names generated code positionally -- two distinct oracles both called "brent_oracle" emit casadi_f1 and casadi_f3, and the plugin's cache variable is already keyed on that positional name -- so constants do the job. Two Brents that differ only in tolerance still get their own rootfinder, because the conversion cache is keyed on the id, which is where identity belongs. Co-Authored-By: Claude Opus 5 --- .../src/pybamm/expression_tree/brent.py | 9 ++++-- .../unit/test_expression_tree/test_brent.py | 30 +++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py index 5d786e2a4f..f8882e1f4a 100644 --- a/packages/pybamm/src/pybamm/expression_tree/brent.py +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -229,7 +229,10 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): "MSVC, which leaves the two with separate copies of CasADi. Call " "evaluate() instead, which solves with SciPy." ) - unknown = casadi.MX.sym(f"brent_unknown_{abs(self.id)}") + # These names reach fn.serialize(), which keys the AOT compile cache, so they + # must not carry `id`: it is a per-process hash. CasADi names generated code + # positionally, so duplicates across two Brent nodes are harmless. + unknown = casadi.MX.sym("brent_unknown") cache = _OracleCache( casadi_symbols, _nodes_reading(self.residual, self.unknown) ) @@ -248,10 +251,10 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): # does not read, so they are declared here and left unused lo_sym, hi_sym = casadi.MX.sym("lo"), casadi.MX.sym("hi") oracle = casadi.Function( - f"brent_oracle_{abs(self.id)}", [unknown, lo_sym, hi_sym, *free], [equation] + "brent_oracle", [unknown, lo_sym, hi_sym, *free], [equation] ) solver = casadi.rootfinder( - f"brent_{abs(self.id)}", + "brent", "brent", oracle, {"abstol": self.abstol, "max_iter": self.max_iter}, diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py index 8b6e3397fd..476e945b17 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -2,7 +2,9 @@ # Tests for the Brent expression tree node # +import subprocess import sys +import textwrap import casadi import numpy as np @@ -159,6 +161,34 @@ def test_evaluate_reports_an_empty_bracket_rather_than_nan(self): with pytest.raises(pybamm.SolverError, match="no sign change over the bracket"): node.evaluate() + def test_the_conversion_carries_no_process_identity(self): + # Symbol.id is a per-process hash, and anything the conversion puts in a name + # reaches fn.serialize(), which is the AOT compile cache key + script = textwrap.dedent(""" + import hashlib + + import casadi + import pybamm + + y = casadi.MX.sym("y") + x = pybamm._BrentUnknown("x") + node = pybamm._Brent( + x * x - pybamm.StateVector(slice(0, 1)), x, (0.0, 10.0) + ) + serialised = casadi.Function("f", [y], [node.to_casadi(y=y)]).serialize() + print(hashlib.sha256(serialised.encode()).hexdigest()) + """) + keys = { + subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + for _ in range(3) + } + assert len(keys) == 1, f"the key moved between processes: {keys}" + def test_children_and_copy(self): x = pybamm._BrentUnknown("x") node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0)) From f4a2c2c8862358261ca33ceba16c85c9dfb2369f Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:49:08 -0400 Subject: [PATCH 10/12] test: pin the cache key without shelling out Bandit flags the subprocess call the previous version used to compare the serialised function across interpreters, and Codacy gates on new issues. Forcing a different id in-process tests the same invariant more directly: two structurally identical nodes, one carrying the id another process would have hashed, must serialise to the same bytes. It catches each of the three names independently, and drops the three interpreter startups the old one paid for. A substring scan for the id would not have worked here: serialize() is encoded, so neither the id nor "brent_oracle" appears in it as text. Co-Authored-By: Claude Opus 5 --- .../unit/test_expression_tree/test_brent.py | 30 +++++-------------- 1 file changed, 8 insertions(+), 22 deletions(-) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py index 476e945b17..1f3edef08f 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -2,9 +2,7 @@ # Tests for the Brent expression tree node # -import subprocess import sys -import textwrap import casadi import numpy as np @@ -162,32 +160,20 @@ def test_evaluate_reports_an_empty_bracket_rather_than_nan(self): node.evaluate() def test_the_conversion_carries_no_process_identity(self): - # Symbol.id is a per-process hash, and anything the conversion puts in a name - # reaches fn.serialize(), which is the AOT compile cache key - script = textwrap.dedent(""" - import hashlib - - import casadi - import pybamm + # Symbol.id is a per-process hash, so anything the conversion derives from one + # moves fn.serialize() -- the AOT compile cache key -- on every run + def serialised(id_offset): y = casadi.MX.sym("y") x = pybamm._BrentUnknown("x") node = pybamm._Brent( x * x - pybamm.StateVector(slice(0, 1)), x, (0.0, 10.0) ) - serialised = casadi.Function("f", [y], [node.to_casadi(y=y)]).serialize() - print(hashlib.sha256(serialised.encode()).hexdigest()) - """) - keys = { - subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - check=True, - ).stdout.strip() - for _ in range(3) - } - assert len(keys) == 1, f"the key moved between processes: {keys}" + # the id another process would have hashed for the same node + node._id = node.id + id_offset + return casadi.Function("f", [y], [node.to_casadi(y=y)]).serialize() + + assert serialised(0) == serialised(1) def test_children_and_copy(self): x = pybamm._BrentUnknown("x") From e48f09dda123985330574079d898e31c842ea54f Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:28:32 -0400 Subject: [PATCH 11/12] feat: run the brent iteration in Python where the plugin cannot be reached Windows holds two copies of CasADi -- ours built with MSVC, the wheel's with MinGW -- so the plugin registered in one is invisible to the other, and the node refused to convert there at all. That took composite electrode SOH in #5730 with it, which works on Windows today. A CasADi Callback closes the gap. The residual is still compiled into the oracle exactly as the plugin path builds it; only the bracketing iteration runs in Python, driven by scipy.optimize.brentq. Everything around the rootfind stays in the graph, so a rootfind nested inside another costs one callback per enclosing iteration rather than a Python tree walk per residual evaluation: 200 nested solves take 0.26s this way against 0.004s through the plugin, where evaluating the tree directly could not finish one composite SOH solve in fourteen minutes. Derivatives come from the implicit function theorem, as the plugin's do: dx/dp = -(dF/dp) / (dF/dx) at the root, emitted as a forward mode over the same oracle. The tests confirm both paths agree to 1e-9 on a case with a known answer. Two things the driver cannot do. It re-enters Python, so the test that pins that out of the plugin path now skips when the plugin is absent, and it holds a callback, so an expression containing one cannot be code-generated. Co-Authored-By: Claude Opus 5 --- .../src/pybamm/expression_tree/brent.py | 117 ++++++++++++++++-- .../unit/test_expression_tree/test_brent.py | 32 +++-- 2 files changed, 129 insertions(+), 20 deletions(-) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py index f8882e1f4a..98e07d7028 100644 --- a/packages/pybamm/src/pybamm/expression_tree/brent.py +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -93,6 +93,101 @@ def __setitem__(self, key, value): self._shared[key] = value +# CasADi holds a Callback by weak reference, so a driver that fell out of scope would +# be collected while the graph still pointed at it. +_LIVE_DRIVERS: list[casadi.Callback] = [] + + +class _BrentDriver(casadi.Callback): + """Drive :func:`scipy.optimize.brentq` over a compiled residual. + + Stands in for the ``"brent"`` plugin where it is not registered with the CasADi + the process calls, which is every Windows install. Only the iteration re-enters + Python; the residual and everything around it stay in the graph, so a rootfind + nested inside another costs one call per enclosing iteration rather than a tree + walk per residual evaluation. + """ + + def __init__(self, oracle, abstol, max_iter): + super().__init__() + self._oracle = oracle + self._abstol = abstol + self._max_iter = max_iter + self.construct("brent_driver", {}) + + def get_n_in(self): + # the oracle's inputs are (unknown, lo, hi, *free); the unknown is ours to find + return self._oracle.n_in() - 1 + + def get_n_out(self): + return 1 + + def get_sparsity_in(self, i): + return self._oracle.sparsity_in(i + 1) + + def get_sparsity_out(self, i): + return casadi.Sparsity.dense(1, 1) + + def has_forward(self, nfwd): + # one direction at a time, so a seed has the same shape as its input + return nfwd == 1 + + def get_forward(self, nfwd, name, inames, onames, opts): + """The implicit function theorem, as the plugin's own derivative is.""" + oracle = self._oracle + unknown = casadi.MX.sym("x") + arguments = [ + casadi.MX.sym(f"i{i}", oracle.sparsity_in(i + 1)) + for i in range(oracle.n_in() - 1) + ] + root = casadi.MX.sym("root") + seeds = [ + casadi.MX.sym(f"d{i}", argument.sparsity()) + for i, argument in enumerate(arguments) + ] + residual = oracle(unknown, *arguments) + # dx/dp = -(dF/dp) / (dF/dx), at the root: the bracket does not move it + free = casadi.vertcat(*[casadi.vec(a) for a in arguments[2:]]) + free_seed = casadi.vertcat(*[casadi.vec(seed) for seed in seeds[2:]]) + numerator = casadi.jtimes(residual, free, free_seed) if arguments[2:] else 0 + derivative = -numerator / casadi.jacobian(residual, unknown) + return casadi.Function( + name, + [*arguments, root, *seeds], + [casadi.substitute(derivative, unknown, root)], + inames, + onames, + {"allow_free": True, **opts}, + ) + + def eval(self, arg): + from scipy.optimize import brentq + + lo, hi = float(arg[0]), float(arg[1]) + free = list(arg[2:]) + + def residual(value): + return float(self._oracle(value, lo, hi, *free)) + + residual_lo, residual_hi = residual(lo), residual(hi) + if ( + residual_lo != 0 + and residual_hi != 0 + and (residual_lo > 0) == (residual_hi > 0) + ): + raise RuntimeError( + f"no sign change over the bracket ({lo}, {hi}), where the residual is " + f"{residual_lo} and {residual_hi}, so it holds no root" + ) + try: + root = brentq(residual, lo, hi, xtol=self._abstol, maxiter=self._max_iter) + except RuntimeError as error: + raise RuntimeError( + f"the rootfind did not converge in {self._max_iter} iterations" + ) from error + return [casadi.DM(root)] + + class _Brent(pybamm.Symbol): """ Solve ``residual == 0`` for ``unknown`` within ``bounds``, by Brent's method. @@ -103,9 +198,12 @@ class _Brent(pybamm.Symbol): plugin registered by ``pybammsolvers``, so the whole solve runs inside the CasADi graph. - Not available on Windows: the CasADi wheel there is built with MinGW and - ``pybammsolvers`` with MSVC, so the two hold separate copies of CasADi and the - plugin never reaches the one Python calls. ``evaluate()`` still works, via SciPy. + Where that plugin is not registered with the CasADi the process calls -- every + Windows install, whose CasADi wheel is built with MinGW while ``pybammsolvers`` + is built with MSVC, leaving two copies of CasADi in the process -- the iteration + runs in Python over the compiled residual instead. Same roots and the same + derivatives, a few tenths of a millisecond slower per solve, but the result + holds a callback and so cannot be code-generated. Brent needs only a sign change over the bounds, so it converges where a Newton iteration stalls, and the answer cannot leave them. Derivatives come from CasADi's @@ -221,14 +319,6 @@ def _from_json(cls, snippet: dict): ) def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): - if not casadi.has_rootfinder("brent"): - raise NotImplementedError( - "the 'brent' rootfinder plugin is not registered with the CasADi that " - "Python is using, so this node cannot be converted to CasADi. On " - "Windows the casadi wheel is built with MinGW and pybammsolvers with " - "MSVC, which leaves the two with separate copies of CasADi. Call " - "evaluate() instead, which solves with SciPy." - ) # These names reach fn.serialize(), which keys the AOT compile cache, so they # must not carry `id`: it is a per-process hash. CasADi names generated code # positionally, so duplicates across two Brent nodes are harmless. @@ -253,6 +343,11 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): oracle = casadi.Function( "brent_oracle", [unknown, lo_sym, hi_sym, *free], [equation] ) + if not casadi.has_rootfinder("brent"): + driver = _BrentDriver(oracle, self.abstol, self.max_iter) + _LIVE_DRIVERS.append(driver) + return driver(lo, hi, *free) + solver = casadi.rootfinder( "brent", "brent", diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py index 1f3edef08f..d964653a8b 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_brent.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -24,14 +24,6 @@ def _evaluate(symbol, **inputs): return float(function(*[inputs[name] for name in symbols])) -# On Windows our MSVC-built extension and the MinGW-built casadi wheel hold separate -# copies of CasADi, so a plugin registered in ours is invisible to the one Python calls. -pytestmark = pytest.mark.skipif( - sys.platform == "win32", - reason="the brent plugin cannot reach the casadi wheel's CasADi on Windows", -) - - class TestBrent: def test_solves_a_scalar_equation(self): x = pybamm._BrentUnknown("x") @@ -108,6 +100,10 @@ def test_nests(self): outer = pybamm._Brent(outer_x * outer_x - inner, outer_x, (0.0, 10.0)) assert float(casadi.evalf(outer.to_casadi(inputs={}))) == pytest.approx(2.0) + @pytest.mark.skipif( + not casadi.has_rootfinder("brent"), + reason="only the plugin keeps the iteration out of Python", + ) def test_evaluating_does_not_re_enter_python(self): # the whole solve runs in the CasADi graph, so a Brent node must cost no more # python frames per evaluation than the same expression without one @@ -150,7 +146,9 @@ def test_the_oracle_only_reads_what_the_residual_needs(self): def test_no_sign_change_fails_rather_than_guessing(self): x = pybamm._BrentUnknown("x") node = pybamm._Brent(x * x + 1 - 0.0, x, (0.0, 1.0)) - with pytest.raises(RuntimeError, match="rootfinder process failed"): + with pytest.raises( + RuntimeError, match=r"rootfinder process failed|no sign change" + ): casadi.evalf(node.to_casadi(inputs={})) def test_evaluate_reports_an_empty_bracket_rather_than_nan(self): @@ -175,6 +173,22 @@ def serialised(id_offset): assert serialised(0) == serialised(1) + def test_the_driver_matches_the_plugin(self, monkeypatch): + # what Windows runs: the iteration in Python over a compiled residual + monkeypatch.setattr(casadi, "has_rootfinder", lambda name: False) + p = pybamm.InputParameter("p") + x = pybamm._BrentUnknown("x") + node = pybamm._Brent(x * x - p, x, (0.0, 10.0)) + symbol = casadi.MX.sym("p") + converted = node.to_casadi(inputs={"p": symbol}) + root = casadi.Function("root", [symbol], [converted]) + assert float(root(6.0)) == pytest.approx(np.sqrt(6), abs=1e-12) + # dx/dp = 1 / (2x), by the implicit function theorem + derivative = casadi.Function( + "d", [symbol], [casadi.jacobian(converted, symbol)] + ) + assert float(derivative(6.0)) == pytest.approx(1 / (2 * np.sqrt(6)), abs=1e-9) + def test_children_and_copy(self): x = pybamm._BrentUnknown("x") node = pybamm._Brent(x * x - 9.0, x, (0.0, 10.0)) From fa12f8d6f1d4b7cd7359bfb7b5f9db064b9cad88 Mon Sep 17 00:00:00 2001 From: Marc Berliner <34451391+MarcBerliner@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:44:48 -0400 Subject: [PATCH 12/12] docs: record that the brent driver cannot cross into the solver An expression holding one is a CasADi Callback, so IDAKLU's deserialise step rejects it -- observing such a variable raises "not found 'CallbackInternal'". Same root cause as the codegen limitation already noted, and worth naming next to it. Co-Authored-By: Claude Opus 5 --- packages/pybamm/src/pybamm/expression_tree/brent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/pybamm/src/pybamm/expression_tree/brent.py b/packages/pybamm/src/pybamm/expression_tree/brent.py index 98e07d7028..2090730379 100644 --- a/packages/pybamm/src/pybamm/expression_tree/brent.py +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -203,7 +203,8 @@ class _Brent(pybamm.Symbol): is built with MSVC, leaving two copies of CasADi in the process -- the iteration runs in Python over the compiled residual instead. Same roots and the same derivatives, a few tenths of a millisecond slower per solve, but the result - holds a callback and so cannot be code-generated. + holds a callback, so it can be neither code-generated nor serialised into the + solver: reading such a variable back through IDAKLU fails to deserialise it. Brent needs only a sign change over the bounds, so it converges where a Newton iteration stalls, and the answer cannot leave them. Derivatives come from CasADi's