From a72301be5a4ea5f8f35e60efd10969729d19e2c9 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 21 Oct 2026 14:53:41 +0100 Subject: [PATCH] fix: make CasADi's RegPower and uniform-grid interpolants match evaluate RegPower's CasADi form, x * hypot(x, delta)^(a-1) * scale^a, rounded differently from evaluate's x * (x^2 + delta^2)^((a-1)/2) * scale^a. Both now evaluate one sequence of operations, and a literal exponent of 0.5 or 2 takes its powers as square roots, so CasADi and evaluate agree bitwise for those exponents. Linear and pchip interpolants on a uniform grid used CasADi's "exact" lookup, which can pick the neighbouring segment for a point at or one ulp beside a knot, so the Jacobian took that segment's slope. They now use binary search, as on a non-uniform grid. --- CHANGELOG.md | 2 + .../src/pybamm/expression_tree/functions.py | 78 ++++++++++--------- .../src/pybamm/expression_tree/interpolant.py | 26 ++----- .../test_expression_tree/test_functions.py | 42 ++++------ .../test_expression_tree/test_interpolant.py | 29 +++++++ .../test_operations/test_convert_to_casadi.py | 49 ++++++++++++ 6 files changed, 145 insertions(+), 81 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba95d816b7..c99486239e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,8 @@ - Sensitivities from `Simulation.solve(initial_soc=...)` now include how inputs such as active material volume fractions, maximum concentrations, voltage limits and OCP parameters set the initial state. Before, they could be off by 100% or have the wrong sign. ([#5948](https://github.com/pybamm-team/PyBaMM/pull/5948)) - `from pybamm.expression_tree import *` no longer raises `AttributeError`, nor does the same import from `pybamm.parameters`, `pybamm.solvers` or `pybamm.models.submodels.interface.kinetics`; their `__all__` listed modules that were moved or removed. `CoupledVariable.create_copy` now takes the same arguments as other symbols, so `create_copy(new_children)` and `new_copy()` no longer raise `TypeError`. Evaluating a concatenation before discretisation raises `NotImplementedError`, as an undiscretised variable does, so `evaluate_ignoring_errors()` returns `None` instead of raising `TypeError`. ([#5974](https://github.com/pybamm-team/PyBaMM/pull/5974)) - Fixed copying a function without simplification, as `ParameterValues` does for a parameter set loaded with `Serialise.load_parameters`, returning a generic `pybamm.Function`: `pybamm.exp` and the other specific functions lost their class, and an `Arcsinh2` copy could not be evaluated. ([#5954](https://github.com/pybamm-team/PyBaMM/pull/5954)) +- `RegPower` now evaluates the same operations in `evaluate` and CasADi, taking a literal exponent of 0.5 or 2 as square roots, so they agree bitwise; CasADi used a `hypot` form. With `delta=0`, `evaluate` and the Jacobian use `base ** exponent`, as CasADi did. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955)) +- CasADi linear and `pchip` interpolants on a uniform grid now find a point's segment by binary search. The `"exact"` lookup could pick the neighbouring segment at or one ulp beside a knot, giving the Jacobian that segment's slope. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955)) ## Optimizations diff --git a/packages/pybamm/src/pybamm/expression_tree/functions.py b/packages/pybamm/src/pybamm/expression_tree/functions.py index e03f7679b7..911a3d4b14 100644 --- a/packages/pybamm/src/pybamm/expression_tree/functions.py +++ b/packages/pybamm/src/pybamm/expression_tree/functions.py @@ -19,17 +19,11 @@ _CASADI_EQUIVALENTS = {special.erf: casadi.erf} -def _is_constant_value(x, value: float) -> bool: - # True iff ``x`` is a scalar constant equal to ``value``. Conservative on - # symbolic or multi-element inputs. - if isinstance(x, casadi.MX) and np.prod(x.size(), initial=1) > 2: - return False - if isinstance(x, np.ndarray): - return bool(np.all(x == value)) - try: - return bool(x == value) - except (TypeError, ValueError, RuntimeError): - return False +def _array_sqrt(value): + # np.sqrt cannot trace a JAX array, which the generated Python evaluator passes + if isinstance(value, (np.ndarray, np.generic, float, int)): + return np.sqrt(value) + return value**0.5 class Function(pybamm.Symbol): @@ -973,9 +967,29 @@ def __init__( def _reg_power_evaluate(self, base, exponent, scale): """Evaluate reg_power using numpy.""" + return self._decomposition( + base, exponent, scale, _array_sqrt, lambda a, b: a**b, lambda v: v + ) + + def _decomposition(self, base, exponent, scale, sqrt, power, constant): + """``x * (x*x + delta^2)^((a-1)/2) * scale^a`` with ``x = base / scale``. + + Every backend evaluates this sequence of operations, given its own ``sqrt``, + ``power(base, exponent)`` and ``constant(float)``. A literal ``a`` of 0.5 or 2 + takes its powers as square roots, so the backends agree bitwise. With + ``delta == 0`` it is ``power(base, exponent)``. + """ + if self.delta == 0: + return power(base, exponent) x = base / scale - x2_d2 = x**2 + self.delta**2 - return x * (x2_d2 ** ((exponent - 1) / 2)) * (scale**exponent) + inner = x * x + constant(self.delta * self.delta) + literal = self.children[1] + if isinstance(literal, pybamm.Scalar) and float(literal.value) == 0.5: + return x / sqrt(sqrt(inner)) * sqrt(scale) + if isinstance(literal, pybamm.Scalar) and float(literal.value) == 2.0: + return x * sqrt(inner) * (scale * scale) + inner_exponent = (exponent - constant(1.0)) * constant(0.5) + return x * power(inner, inner_exponent) * power(scale, exponent) def _function_evaluate(self, evaluated_children): """See :meth:`pybamm.Function._function_evaluate()`.""" @@ -990,6 +1004,15 @@ def _function_diff(self, children, idx): base, exponent, scale = children delta = self.delta + if delta == 0: + # Partials of the ``base ** exponent`` fallback + if idx == 0: + return exponent * base ** (exponent - 1) + if idx == 1: + return base**exponent * pybamm.log(base) + if idx == 2: + return pybamm.Scalar(0) + x = base / scale x2_d2 = x**2 + delta**2 scale_factor = scale**exponent @@ -1025,6 +1048,8 @@ def _function_diff(self, children, idx): def _function_jac(self, children_jacs): """See :meth:`pybamm.Function._function_jac()`.""" + if self.delta == 0: + return super()._function_jac(children_jacs) children = self.orphans base, exponent, scale = children delta = self.delta @@ -1082,30 +1107,11 @@ def _function_new_copy(self, children, perform_simplifications=True): return new_function def _casadi_evaluate(self, *converted_children): - """See :meth:`pybamm.Function._casadi_evaluate()`. - - Evaluates ``y = x * r^(a-1) * scale^a`` with ``x = base / scale`` and - ``r = hypot(x, delta)``. The ``a == 0.5`` case (hot for Butler-Volmer) - is specialised to avoid a general ``power`` call. - """ + """See :meth:`pybamm.Function._casadi_evaluate()`.""" base, exponent, scale = converted_children - delta = self.delta - if delta == 0: - return casadi.power(base, exponent) - - x = base / scale - r = casadi.hypot(x, delta) - - if _is_constant_value(exponent, 0.5): - # ``exponent**0`` acts as a ``ones_like(exponent)`` that works - # across casadi.MX / numpy / float. - scale_factor = casadi.sqrt(scale) * exponent**0 - inner_factor = 1 / casadi.sqrt(r) - else: - scale_factor = casadi.power(scale, exponent) - inner_factor = casadi.power(r, exponent - 1) - - return x * inner_factor * scale_factor + return self._decomposition( + base, exponent, scale, casadi.sqrt, casadi.power, lambda v: v + ) def _sympy_operator(self, *children): """Convert to SymPy expression.""" diff --git a/packages/pybamm/src/pybamm/expression_tree/interpolant.py b/packages/pybamm/src/pybamm/expression_tree/interpolant.py index 46149efc17..43b5f933c7 100644 --- a/packages/pybamm/src/pybamm/expression_tree/interpolant.py +++ b/packages/pybamm/src/pybamm/expression_tree/interpolant.py @@ -15,11 +15,6 @@ import pybamm -def _is_uniform_grid(x: npt.NDArray[np.float64]) -> bool: - # Try seeing if the grid was computed using np.linspace - return bool(np.array_equal(x, np.linspace(x[0], x[-1], len(x)))) - - class Interpolant(pybamm.Function): """ Interpolate data in 1D, 2D, or 3D. Interpolation in 3D requires the input data to be @@ -298,7 +293,8 @@ def _linear_to_casadi(self, converted_children): v = self.y.flatten() else: v = self.y.ravel(order="F") - lookup_modes = ["exact" if _is_uniform_grid(xi) else "binary" for xi in self.x] + # CasADi's "exact" lookup can pick a neighbouring segment beside a knot + lookup_modes = ["binary"] * len(self.x) result = casadi.MX.interpn_linear( self.x, v, converted_children, {"lookup_mode": lookup_modes} ) @@ -344,7 +340,6 @@ def _pchip_to_casadi(self, converted_children): n = len(x_np) - 1 h = np.diff(x_np) - uniform = _is_uniform_grid(x_np) is_vector_valued = y_np.ndim > 1 y_2d = y_np if is_vector_valued else y_np[:, np.newaxis] d_2d = d_np if d_np.ndim > 1 else d_np[:, np.newaxis] @@ -352,11 +347,7 @@ def _pchip_to_casadi(self, converted_children): # Precompute Hermite cubic coefficients per interval, pre-divided by h^k. # p(dx) = c0 + c1*dx + c2*dx² + c3*dx³, dx = x - x_lo - # For non-uniform grids x_lo is stored as a 5th table entry; - # for uniform grids x_lo is computed from the index via scalar arithmetic. - stride = 4 - if not uniform: - stride += 1 + stride = 5 coeffs = np.empty((n, stride, m)) y0 = y_2d[:-1] # (n, m) @@ -369,8 +360,7 @@ def _pchip_to_casadi(self, converted_children): coeffs[:, 1] = hd0 * inv_h coeffs[:, 2] = (-3 * y0 - 2 * hd0 + 3 * y1 - hd1) * inv_h**2 coeffs[:, 3] = (2 * y0 + hd0 - 2 * y1 + hd1) * inv_h**3 - if not uniform: - coeffs[:, 4, :] = x_np[:-1, np.newaxis] + coeffs[:, 4, :] = x_np[:-1, np.newaxis] # Honour _num_derivatives (#5582): since dx = x - x_lo, d/dx == d/d(dx), # so differentiating just shifts/scales the coeffs (x_lo row untouched). @@ -385,16 +375,12 @@ def _pchip_to_casadi(self, converted_children): # 1. Single interval index lookup (clamp handles extrapolation) x = converted_children[0] - lookup_mode = "exact" if uniform else "binary" - idx = casadi.low(casadi.MX(x_np), x, {"lookup_mode": lookup_mode}) + idx = casadi.low(casadi.MX(x_np), x, {"lookup_mode": "binary"}) idx = casadi.fmin(casadi.fmax(idx, 0), n - 1) base = idx * stride # 2. Compute dx = x - x_lo - if uniform: - dx = x - (x_np[0] + idx * h[0]) - else: - dx = x - coeffs_flat[base + 4, 0] + dx = x - coeffs_flat[base + 4, 0] # 3. Coefficient lookup and Horner evaluation c0 = coeffs_flat[base, :] diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_functions.py b/packages/pybamm/tests/unit/test_expression_tree/test_functions.py index d6c8525193..2e424003bd 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_functions.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_functions.py @@ -3,15 +3,14 @@ # import itertools +from decimal import Decimal, localcontext -import casadi import numpy as np import pytest import sympy from scipy import special import pybamm -from pybamm.expression_tree.functions import _is_constant_value from tests import ( function_test, multi_var_function_test, @@ -298,6 +297,22 @@ def test_reg_pow(self): result = result.toarray() assert np.isfinite(result).all() + def test_reg_power_half_is_within_its_rounding_bound(self): + # Nine roundings (x/scale, x*x, delta^2, +, sqrt, sqrt, /, sqrt(scale), *) bound + # the relative error by (6 - w/2) u, w = x^2 / (x^2 + delta^2) and u = 2^-53. + delta, scale = pybamm.settings.tolerances["reg_power"], 33133.0 + base = scale * np.geomspace(1e-12, 2.0, 3000) + y = pybamm.StateVector(slice(0, base.size)) + values = pybamm.reg_power(y, 0.5, scale=scale).evaluate(y=base).ravel() + unit_roundoff = Decimal(2) ** -53 + with localcontext() as context: + context.prec = 60 + delta_sq = Decimal(delta) ** 2 + for b, value in zip(base, values, strict=True): + x = Decimal(b) / Decimal(scale) + exact = x / (x * x + delta_sq).sqrt().sqrt() * Decimal(scale).sqrt() + assert abs(Decimal(value) - exact) <= 6 * unit_roundoff * exact, b + class TestSpecificFunctions: @pytest.mark.parametrize("perform_simplifications", [True, False]) @@ -449,26 +464,3 @@ def test_normal_cdf(self): assert fun.evaluate(inputs={"x": -1, "mu": 0, "sigma": 1}) < fun.evaluate( inputs={"x": -1, "mu": 0, "sigma": 2} ) - - -class TestIsConstantValue: - def test_all_branches(self): - - # casadi.MX with size > 2 -> False - large_mx = casadi.MX.sym("x", 10) - assert not _is_constant_value(large_mx, 0) - - # numpy array - assert _is_constant_value(np.array([1.0, 1.0]), 1.0) - assert not _is_constant_value(np.array([1.0, 2.0]), 1.0) - - # scalar - assert _is_constant_value(0.5, 0.5) - assert not _is_constant_value(0.5, 1.0) - - # exception branch - use object that raises on comparison - class BadCompare: - def __eq__(self, other): - raise TypeError("can't compare") - - assert not _is_constant_value(BadCompare(), 0) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_interpolant.py b/packages/pybamm/tests/unit/test_expression_tree/test_interpolant.py index 702b4fed42..91763ff1d4 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_interpolant.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_interpolant.py @@ -7,6 +7,7 @@ import casadi import numpy as np import pytest +from scipy import interpolate import pybamm @@ -393,6 +394,34 @@ def test_diff_to_casadi(self, grid, interpolator, assert_casadi_matches_evaluate assert_casadi_matches_evaluate(interp.diff(y), casadi_y, y_test) assert_casadi_matches_evaluate(interp.diff(y).diff(y), casadi_y, y_test) + @pytest.mark.parametrize(("interpolator", "order"), [("linear", 1), ("pchip", 2)]) + def test_to_casadi_uses_the_segment_holding_x_on_a_uniform_grid( + self, interpolator, order + ): + # At and beside each knot, only the segment holding x has the right slope + # (linear) or curvature (pchip); a neighbour's differs. + grid = np.linspace(0.0, 1.0, 101) + knots = grid[1:-1] + x = np.concatenate( + [np.nextafter(knots, -np.inf), knots, np.nextafter(knots, np.inf)] + ) + child = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant(grid, grid**3, child, interpolator=interpolator) + casadi_x = casadi.MX.sym("x", 1) + derivative = interp.to_casadi(y=casadi_x) + for _ in range(order): + derivative = casadi.jacobian(derivative, casadi_x) + f = casadi.Function("f", [casadi_x], [derivative]) + got = np.array([float(f(v)) for v in x]) + + segment = np.searchsorted(grid, x, side="right") - 1 + if interpolator == "linear": + y = grid**3 + want = (y[segment + 1] - y[segment]) / (grid[segment + 1] - grid[segment]) + else: + want = interpolate.PchipInterpolator(grid, grid**3).derivative(2)(x) + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-9) + @pytest.mark.parametrize("interpolator", ["cubic", "pchip"]) def test_diff_to_casadi_vector_valued( self, interpolator, assert_casadi_matches_evaluate diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py index 13dcc5e7c9..6509f939c9 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py @@ -275,6 +275,55 @@ def test_reg_power_casadi_general_exponent(self): err_msg=f"exponent={exponent}, scale={scale}, x={x_val}", ) + @pytest.mark.parametrize("exponent", [0.5, 2.0]) + @pytest.mark.parametrize("scale", [1.0, 33133.0]) + def test_reg_power_casadi_is_bitwise_identical_to_evaluate(self, exponent, scale): + # These exponents use only operations IEEE 754 rounds correctly + magnitudes = np.concatenate([[0.0, 5e-324], np.geomspace(1e-12, 2.0, 2000)]) + base = scale * np.concatenate([magnitudes, -magnitudes]) + y = pybamm.StateVector(slice(0, base.size)) + expr = pybamm.reg_power(y, exponent, scale=scale) + + reference = expr.evaluate(y=base).ravel() + y_casadi = casadi.MX.sym("y", base.size) + f = casadi.Function("f", [y_casadi], [expr.to_casadi(y=y_casadi)]) + for name, function in {"MX": f, "SX": f.expand()}.items(): + np.testing.assert_array_equal( + np.asarray(function(base), dtype=np.float64).ravel().view(np.uint64), + reference.view(np.uint64), + err_msg=name, + ) + + @pytest.mark.parametrize("exponent", [0.5, 2.0]) + def test_reg_power_with_zero_delta_is_power_in_evaluate_and_casadi(self, exponent): + base = np.array([0.0, -4.0, 4.0]) + y = pybamm.StateVector(slice(0, base.size)) + expr = pybamm.RegPower(y, exponent, delta=0) + y_casadi = casadi.MX.sym("y", base.size) + f = casadi.Function("f", [y_casadi], [expr.to_casadi(y=y_casadi)]) + + with np.errstate(invalid="ignore"): + expected = np.power(base, exponent) + evaluated = expr.evaluate(y=base).ravel() + np.testing.assert_array_equal(evaluated, expected) + np.testing.assert_array_equal(np.asarray(f(base)).ravel(), expected) + + @pytest.mark.parametrize("exponent", [0.5, 2.0]) + def test_reg_power_with_zero_delta_jacobian_matches_casadi(self, exponent): + base = np.array([0.0, -4.0, 4.0]) + y = pybamm.StateVector(slice(0, base.size)) + expr = pybamm.RegPower(y, exponent, delta=0) + y_casadi = casadi.MX.sym("y", base.size) + casadi_jac = casadi.Function( + "J", [y_casadi], [casadi.jacobian(expr.to_casadi(y=y_casadi), y_casadi)] + ) + + with np.errstate(divide="ignore", invalid="ignore"): + jac = expr.jac(y).evaluate(y=base) + np.testing.assert_array_equal( + np.diag(jac.toarray()), np.diag(np.asarray(casadi_jac(base))) + ) + def test_kronecker_product(self): a = np.array([1.0, 2.0, 3.0]) b = np.array([[1, 2, 3], [4, 5, 6]])