diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index 8370ce6f58..740fca3e26 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..722904c501 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..2090730379 --- /dev/null +++ b/packages/pybamm/src/pybamm/expression_tree/brent.py @@ -0,0 +1,425 @@ +# +# 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 + + +# 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. + + 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. + + 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, 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 + 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:`_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 _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): + # 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) + ) + 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( + "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", + 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 _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: + 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'." + ) + + 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..c163a8508a 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..a1e296cd96 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..d964653a8b --- /dev/null +++ b/packages/pybamm/tests/unit/test_expression_tree/test_brent.py @@ -0,0 +1,266 @@ +# +# 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) + + @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 + 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=r"rootfinder process failed|no sign change" + ): + 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_the_conversion_carries_no_process_identity(self): + # 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) + ) + # 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_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)) + 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 _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..649c7c8312 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,211 @@ 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 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") + +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 () + +# 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\" " + "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 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 () + +# 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) +set(_casadi_internal_work "${CMAKE_CURRENT_BINARY_DIR}/casadi-internal") +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) +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") + 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}" + ${_casadi_expected_hash} + 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..ac94b4931f --- /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 + m->key.clear(); + for (casadi_int i = 0; i < n_in_; ++i) { + if (i == iin_ || !m->iarg[i]) continue; + m->key.insert(m->key.end(), m->iarg[i], m->iarg[i] + nnz_in(i)); + } + 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)"; + 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 = m->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..8666d8c844 --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent.hpp @@ -0,0 +1,127 @@ +// +// 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. +// +// 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 + +#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; + // 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; +}; + +/** + * @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; + 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..b842b1df48 --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/brent_impl.hpp @@ -0,0 +1,94 @@ +// +// 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 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 +// 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; + 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. 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; + e = d; + for (k = 0; k < max_iter; ++k) { + T1 tol, xm, step; + 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 + 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)); + 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..00fdac60e8 --- /dev/null +++ b/packages/pybammsolvers/tests/test_brent_rootfinder.py @@ -0,0 +1,342 @@ +"""Unit tests for the "brent" CasADi rootfinder plugin.""" + +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 + + +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_converges_across_a_sweep(self, expr_fn): + rf, f = _solver(expr_fn) + 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 solve + got = float(rf(p)) + assert LO <= got <= HI + worst_residual = max(worst_residual, abs(float(f(got, p)))) + 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_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]) + 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)