fix: make CasADi's RegPower and uniform-grid interpolants match evaluate

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
 

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."""

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,19 +340,14 @@ 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]
         m = y_2d.shape[1]
 
         # 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, :]

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)

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

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]])