diffrax-bdf 0.1.0__py3-none-any.whl

diffrax_bdf/__init__.py

@@ -0,0 +1,7 @@
+from ._bdf import BDF as BDF, ErrorWithOrder as ErrorWithOrder
+from ._controller import BDFController as BDFController
+from ._coeffs import MAX_ORDER as MAX_ORDER
+from ._term import SemiExplicitDAETerm as SemiExplicitDAETerm
+
+
+__version__ = "0.1.0"

diffrax_bdf/_bdf.py

@@ -0,0 +1,656 @@
+"""A fixed-order BDF/NDF solver for diffrax.
+
+The method is the one used by `scipy.integrate.BDF` and MATLAB's `ode15s`: a
+backward differentiation formula applied to a quasi-constant-step history stored
+as backward differences, with the Shampine & Reichelt NDF correction available.
+
+On top of that sits SUNDIALS CVODE's linear-algebra reuse policy, which is what
+makes a BDF worth using over an ESDIRK. `Kvaerno5` already computes only one
+Jacobian per step and reuses it across its five stages, so "one Newton solve per
+step instead of five stage solves" is worth only a few linear solves. The real
+saving is keeping `J` and the factorisation of `I - c * J` across *steps*, which
+CVODE gates on two separate counters (see `_should_update`).
+
+Three properties of `diffrax`'s integration loop shape the implementation:
+
+* `solver_state` is a `lax.while_loop` carry, so it must have a fixed pytree
+  structure, shape and dtype on every step. The BDF order is therefore a *traced*
+  integer used to mask fixed-size buffers, never a Python-level dimension.
+* `solver_state` is rolled back when a step is rejected (`_integrate.py:423-425`).
+  This is why the reuse counters live here rather than in the controller: they must
+  stay consistent with the stored factorisation, which is also rolled back. It also
+  matches CVODE, whose `nst`-based counters likewise advance only on accepted steps.
+* `t0` and `t1` are only ever passed to `terms.contr`/`terms.vf_prod` and never
+  subtracted. Under backwards integration `diffrax` wraps the terms so that `contr`
+  returns a *negative* step, and every formula below is consistent under that sign
+  provided the step is never recomputed as `t1 - t0`.
+"""
+
+from collections.abc import Callable
+from typing import ClassVar
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.scipy.linalg as jsl
+import jax.tree_util as jtu
+import optimistix as optx
+from diffrax import (
+    AbstractAdaptiveSolver,
+    AbstractImplicitSolver,
+    AbstractTerm,
+    LocalLinearInterpolation,
+    RESULTS,
+    VeryChord,
+    with_stepsize_controller_tols,
+)
+from equinox.internal import ω
+from jax.flatten_util import ravel_pytree
+from jaxtyping import Array, Bool, Float, Int, PyTree
+
+from ._coeffs import D_ROWS, MAX_ORDER, change_D, make_tables, weighted_row_sum
+from ._newton import chord_solve, newton_tolerance, rms_norm
+from ._term import flat_algebraic_mask, split_dae_term
+
+
+class ErrorWithOrder(eqx.Module):
+    """An error estimate together with the BDF order that produced it.
+
+    `diffrax` asks the solver for its error order once, outside the loop
+    (`_integrate.py:388` calls `solver.error_order(terms)`, which sees neither the
+    state nor the time), so a variable-order solver has no way to tell the
+    controller which order the current error estimate belongs to. The step-size
+    controller's only other input from the solver is `y_error`, and the integration
+    loop treats that opaquely -- it only does a structure-agnostic `tree_map` over it
+    (`_integrate.py:386`) -- so the order travels in here alongside the estimate.
+
+    [`diffrax_bdf.BDFController`][] unpacks this. Stock `diffrax.PIDController`
+    cannot: it tree-maps `y_error` against `y0` and will raise on the extra leaf.
+    """
+
+    error: PyTree
+    order: Float[Array, ""]
+
+    def __truediv__(self, other):
+        # `PIDController` tree-maps its scaling function over `(y0, y1, y_error)`,
+        # taking its structure from `y0`, so this whole object arrives where a leaf
+        # was expected and the first thing done to it is a division. Intercepting
+        # that turns an unreadable TypeError into an actionable one.
+        del other
+        raise TypeError(
+            "A variable-order `BDF` reports its error estimate together with the "
+            "order that produced it, and `diffrax.PIDController` cannot read that. "
+            "Either pass `stepsize_controller=diffrax_bdf.BDFController(...)`, which "
+            "can, or pin the order with `BDF(bdf_order=...)` to get a plain error "
+            "estimate that any adaptive controller accepts."
+        )
+
+
+class _BDFState(eqx.Module):
+    """The multistep history, plus the cached linear algebra and its age."""
+
+    d_array: PyTree[Float[Array, "8 ..."]]
+    order: Int[Array, ""]
+    n_equal_steps: Int[Array, ""]
+    h_prev: Float[Array, ""]
+    jac: Float[Array, "n n"]
+    lu: Float[Array, "n n"]
+    piv: Int[Array, " n"]
+    c_at_last_lu: Float[Array, ""]
+    steps_since_lu: Int[Array, ""]
+    steps_since_jac: Int[Array, ""]
+    initial_failed: Bool[Array, ""]
+
+
+def _bdf_residual(d, nonlinear_args):
+    """Residual whose root is the total Newton correction from the predictor.
+
+    Solving for the correction rather than for `y1` avoids forming `y1 - y_pred`,
+    a cancellation between two `O(|y|)` quantities whose difference is
+    `O(h ** (k + 1))`; that difference feeds straight into the error estimate.
+    """
+    vf_prod, t1, y_pred, args, control, inv_alpha, psi = nonlinear_args
+    y = (y_pred**ω + d**ω).ω
+    # `vf_prod(t1, y, args, control)` is `h * f(t1, y)`, so scaling by
+    # `1 / alpha[order]` gives scipy's `c * f` with `c = h / alpha[order]`.
+    c_f = (inv_alpha * vf_prod(t1, y, args, control) ** ω).ω
+    return (c_f**ω - psi**ω - d**ω).ω
+
+
+def _zero_algebraic(x, mask):
+    if mask is None:
+        return x
+    flat, unravel = ravel_pytree(x)
+    return unravel(jnp.where(flat_algebraic_mask(mask, x), 0, flat))
+
+
+def _suppress(x, error_mask):
+    if error_mask is None:
+        return x
+    return jnp.where(error_mask, 0, x)
+
+
+def _row_mask(mask, x):
+    """Broadcast a length-`D_ROWS` row mask against a leaf of the difference array."""
+    return jnp.reshape(mask, (-1,) + (1,) * (jnp.ndim(x) - 1))
+
+
+def _update_difference_array(d_array, order, correction):
+    """Fold an accepted Newton correction into the difference array.
+
+    scipy writes this as
+
+        D[order + 2] = d - D[order + 1]
+        D[order + 1] = d
+        for i in reversed(range(order + 1)): D[i] += D[i + 1]
+
+    The loop is a reverse cumulative sum over rows `0..order + 1`, which is done
+    here by masking rather than slicing so that the shapes stay static.
+    """
+    idx = jnp.arange(D_ROWS)
+
+    def update(x, d):
+        x = x.at[order + 2].set(d - x[order + 1])
+        x = x.at[order + 1].set(d)
+        summable = jnp.where(_row_mask(idx <= order + 1, x), x, 0)
+        reverse_cumsum = jnp.cumsum(summable[::-1], axis=0)[::-1]
+        return jnp.where(_row_mask(idx <= order, x), reverse_cumsum, x)
+
+    return jtu.tree_map(update, d_array, correction)
+
+
+class BDF(AbstractImplicitSolver, AbstractAdaptiveSolver):
+    """Fixed-order backward differentiation formula with Jacobian reuse.
+
+    The order ramps from 1 up to `bdf_order` over the opening steps, since a
+    `k`-step formula needs `k` points of history. Because a rejected step rolls the
+    solver state back, the ramp is automatically correct: a step that did not
+    happen does not advance the order.
+
+    Pair this with [`diffrax_bdf.BDFController`][]. With diffrax's stock
+    `PIDController` the step size, and hence `c`, changes on every step, so the
+    factorisation is rebuilt every step and the reuse below never engages.
+
+    It also solves semi-explicit index-1 differential algebraic equations given as
+    a [`diffrax_bdf.SemiExplicitDAETerm`][].
+
+    !!! warning
+
+        The Jacobian is formed densely, so this is intended for small to medium
+        systems. The backward differences also suffer heavy cancellation at high
+        order: orders above 2 are effectively noise in float32, so enable
+        `jax_enable_x64`.
+    """
+
+    term_structure: ClassVar = AbstractTerm
+    interpolation_cls: ClassVar[Callable[..., LocalLinearInterpolation]] = (
+        LocalLinearInterpolation
+    )
+
+    bdf_order: int | None = eqx.field(static=True, default=None)
+    max_order: int = eqx.field(static=True, default=MAX_ORDER)
+    use_ndf: bool = eqx.field(static=True, default=True)
+    root_finder: optx.AbstractRootFinder = with_stepsize_controller_tols(VeryChord)()
+    root_find_max_steps: int = 4
+    # CVODE's `MSBP` and `MSBJ`, and its `DGMAX_LSETUP`.
+    lu_reuse_steps: int = eqx.field(static=True, default=20)
+    jac_reuse_steps: int = eqx.field(static=True, default=51)
+    max_gamma_change: float = 0.3
+    suppress_algebraic_error: bool = eqx.field(static=True, default=False)
+    correct_initial_algebraic: bool = eqx.field(static=True, default=True)
+    initial_algebraic_max_steps: int = eqx.field(static=True, default=50)
+
+    def __check_init__(self):
+        if self.bdf_order is not None and not 1 <= self.bdf_order <= MAX_ORDER:
+            raise ValueError(f"`bdf_order` must be between 1 and {MAX_ORDER}, or None.")
+        if not 1 <= self.max_order <= MAX_ORDER:
+            raise ValueError(f"`max_order` must be between 1 and {MAX_ORDER}.")
+
+    @property
+    def _variable_order(self) -> bool:
+        return self.bdf_order is None
+
+    @property
+    def _order_cap(self) -> int:
+        return self.max_order if self._variable_order else self.bdf_order
+
+    def order(self, terms):
+        del terms
+        return self._order_cap
+
+    def error_order(self, terms):
+        del terms
+        # The local error of a k-th order BDF step is O(h ** (k + 1)). At variable
+        # order this is only the value used to pick the very first step size; the
+        # per-step order reaches the controller through `ErrorWithOrder`.
+        return self._order_cap + 1
+
+    def func(self, terms, t0, y0, args):
+        terms, _ = split_dae_term(terms)
+        return terms.vf(t0, y0, args)
+
+    def _tolerances(self):
+        rtol = getattr(self.root_finder, "rtol", None)
+        atol = getattr(self.root_finder, "atol", None)
+        if not isinstance(rtol, (int, float)) or not isinstance(atol, (int, float)):
+            raise ValueError(
+                "`BDF` needs concrete tolerances for its Newton iteration. Either use "
+                "an adaptive step size controller, which supplies them, or construct "
+                "the solver with e.g. "
+                "`BDF(root_finder=diffrax.VeryChord(rtol=..., atol=...))`."
+            )
+        return rtol, atol
+
+    def init(self, terms, t0, t1, y0, args) -> _BDFState:
+        terms, mask = split_dae_term(terms)
+        control = terms.contr(t0, t1)
+        if jnp.shape(control) != ():
+            raise ValueError(
+                "`BDF` requires a term whose control is a scalar (e.g. `ODETerm`); "
+                f"got a control of shape {jnp.shape(control)}. The step-size ratio "
+                "used to rescale the difference array is not defined otherwise."
+            )
+        if mask is not None and self.correct_initial_algebraic:
+            y0, initial_failed = self._correct_algebraic(terms, t0, y0, args, mask)
+        else:
+            initial_failed = jnp.asarray(False)
+        # Start at order 1 with `D = [y0, h * f0]`, matching scipy, except that the
+        # algebraic components of `h * f0` are zero: there `f0` is a constraint
+        # residual, not a derivative. The remaining rows are zero rather than
+        # uninitialised so that a high-order formula applied before enough history
+        # exists degenerates gracefully.
+        f0 = _zero_algebraic(terms.vf_prod(t0, y0, args, control), mask)
+        d_array = jtu.tree_map(
+            lambda y, f: jnp.zeros((D_ROWS,) + jnp.shape(y), jnp.result_type(y))
+            .at[0]
+            .set(y)
+            .at[1]
+            .set(f),
+            y0,
+            f0,
+        )
+        flat, _ = ravel_pytree(y0)
+        size = flat.size
+        dtype = flat.dtype
+        return _BDFState(
+            d_array=d_array,
+            order=jnp.asarray(1, dtype=jnp.int32),
+            n_equal_steps=jnp.asarray(0, dtype=jnp.int32),
+            h_prev=control,
+            jac=jnp.zeros((size, size), dtype),
+            lu=jnp.zeros((size, size), dtype),
+            piv=jnp.zeros((size,), jnp.int32),
+            c_at_last_lu=jnp.asarray(0.0, dtype),
+            # Start both counters expired so the first step builds everything.
+            steps_since_lu=jnp.asarray(self.lu_reuse_steps, jnp.int32),
+            steps_since_jac=jnp.asarray(self.jac_reuse_steps, jnp.int32),
+            initial_failed=initial_failed,
+        )
+
+    def _correct_algebraic(self, terms, t0, y0, args, mask):
+        flat_y0, unravel = ravel_pytree(y0)
+        flat_mask = flat_algebraic_mask(mask, y0)
+
+        def residual(flat, _):
+            f = ravel_pytree(terms.vf(t0, unravel(flat), args))[0]
+            return jnp.where(flat_mask, f, flat - flat_y0)
+
+        rtol, atol = self._tolerances()
+        solution = optx.root_find(
+            residual,
+            optx.Newton(rtol=rtol, atol=atol),
+            flat_y0,
+            max_steps=self.initial_algebraic_max_steps,
+            throw=False,
+        )
+        failed = (solution.result != optx.RESULTS.successful) | jnp.invert(
+            jnp.all(jnp.isfinite(solution.value))
+        )
+        corrected = jnp.where(failed, flat_y0, solution.value)
+        return unravel(corrected), failed
+
+    def _should_update(self, state, c):
+        """CVODE's two-tier gate on the Jacobian and its factorisation.
+
+        The asymmetry is the point. `J` costs a full automatic-differentiation pass
+        through the vector field and survives ~51 steps; the factorisation costs an
+        LU and is redone whenever `c` has moved appreciably or 20 steps have passed,
+        reassembling `I - c * J` from the *stored* `J` without touching the vector
+        field. So the expensive quantity is refreshed far less often than the cheap
+        one. See `cvLsSetup` and `cvNls` in SUNDIALS.
+        """
+        need_jac = state.steps_since_jac >= self.jac_reuse_steps
+        safe_previous = jnp.where(state.c_at_last_lu == 0, jnp.inf, state.c_at_last_lu)
+        gamma_change = jnp.abs(c / safe_previous - 1)
+        need_lu = (
+            need_jac
+            | (gamma_change > self.max_gamma_change)
+            | (state.steps_since_lu >= self.lu_reuse_steps)
+        )
+        return need_jac, need_lu, gamma_change
+
+    def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
+        terms, mask = split_dae_term(terms)
+        control = terms.contr(t0, t1)
+        gamma_table, alpha_table, error_const = make_tables(self.use_ndf)
+
+        d_array = solver_state.d_array
+        order = solver_state.order
+        # The difference array is scaled for the previous step size; rescale it
+        # before use, or the history polynomial is evaluated at the wrong times.
+        factor = control / solver_state.h_prev
+        d_array = change_D(d_array, order, factor)
+        n_equal_steps = jnp.where(factor == 1, solver_state.n_equal_steps + 1, 0)
+        steps_since_jac = solver_state.steps_since_jac
+
+        # A jump means the vector field was discontinuous at `t0`, so the history
+        # polynomial is meaningless and we restart at order 1 with a fresh Jacobian.
+        # `PIDController` always reports `False` and `diffrax` keeps it a static bool
+        # where it can, so in the common case this costs nothing at all.
+        if not (isinstance(made_jump, bool) and not made_jump):
+            d_array, order, n_equal_steps, steps_since_jac = self._restart(
+                terms,
+                t0,
+                y0,
+                args,
+                control,
+                d_array,
+                order,
+                n_equal_steps,
+                steps_since_jac,
+                made_jump,
+                mask,
+            )
+
+        flat_y0, unravel = ravel_pytree(y0)
+        dtype = flat_y0.dtype
+        size = flat_y0.size
+        indices = jnp.arange(MAX_ORDER + 1)
+        active = indices <= order
+        inv_alpha = 1 / jnp.asarray(alpha_table, dtype)[order]
+        c = control * inv_alpha
+
+        y_pred = weighted_row_sum(active.astype(dtype), d_array)
+        psi_coeffs = jnp.where(active, jnp.asarray(gamma_table, dtype), 0.0)
+        psi = (inv_alpha * weighted_row_sum(psi_coeffs, d_array) ** ω).ω
+        flat_pred, _ = ravel_pytree(y_pred)
+        flat_psi, _ = ravel_pytree(psi)
+
+        def vector_field(flat):
+            return ravel_pytree(terms.vf(t1, unravel(flat), args))[0]
+
+        if mask is None:
+            flat_mask = None
+            mass = jnp.eye(size, dtype=dtype)
+        else:
+            flat_mask = flat_algebraic_mask(mask, y0)
+            mass = jnp.diag(jnp.where(flat_mask, 0, 1).astype(dtype))
+
+        def residual(correction):
+            y = unravel(flat_pred + correction)
+            scaled = ravel_pytree(terms.vf_prod(t1, y, args, control))[0] * inv_alpha
+            if flat_mask is None:
+                return scaled - flat_psi - correction
+            return (
+                scaled
+                - jnp.where(flat_mask, 0, flat_psi)
+                - jnp.where(flat_mask, 0, correction)
+            )
+
+        state = eqx.tree_at(
+            lambda s: s.steps_since_jac, solver_state, steps_since_jac
+        )
+        need_jac, need_lu, gamma_change = self._should_update(state, c)
+
+        rtol, atol = self._tolerances()
+        scale = atol + rtol * jnp.abs(flat_pred)
+        tol = newton_tolerance(rtol, dtype)
+
+        def factorise(jacobian):
+            return jsl.lu_factor(mass - c * jacobian)
+
+        jac = jax.lax.cond(
+            need_jac, lambda: jax.jacfwd(vector_field)(flat_pred), lambda: state.jac
+        )
+        lu, piv = jax.lax.cond(
+            need_lu, lambda: factorise(jac), lambda: (state.lu, state.piv)
+        )
+        guess = jnp.zeros_like(flat_pred)
+        correction, converged, _ = chord_solve(
+            residual, guess, lu, piv, scale, tol, self.root_find_max_steps
+        )
+
+        # A Newton failure against reused linear algebra is not evidence that the step
+        # is too large, so rejecting would be the wrong response -- and it would not
+        # even help, since a rejected step rolls the solver state back and the retry
+        # would arrive with the same stale matrices. scipy and CVODE both refresh and
+        # retry within the step instead.
+        #
+        # Which part to refresh matters. If the factorisation was stale with respect to
+        # `c` -- because `c` moved less than `max_gamma_change`, or because the order
+        # ramped and changed `alpha[order]` -- then rebuilding it from the *stored*
+        # Jacobian is enough, and costs no vector-field evaluations. Only when the
+        # factorisation was already current at this `c` is the Jacobian itself the
+        # suspect. This is CVODE's `convfail` distinction.
+        retry = jnp.invert(converged)
+        # Is the factorisation we just used already built at this `c`? If so the
+        # Jacobian inside it is the only thing left to blame. If not -- `c` drifted
+        # within the deadband, or the order ramped and changed `alpha[order]` --
+        # rebuilding the factorisation from the stored Jacobian is the cheap fix.
+        lu_is_current = need_lu | (gamma_change < 1e-10)
+        blame_jacobian = retry & lu_is_current
+
+        def refresh():
+            fresh_jac = jax.lax.cond(
+                blame_jacobian,
+                lambda: jax.jacfwd(vector_field)(flat_pred),
+                lambda: jac,
+            )
+            fresh_lu, fresh_piv = factorise(fresh_jac)
+            retried, retried_converged, _ = chord_solve(
+                residual, guess, fresh_lu, fresh_piv, scale, tol, self.root_find_max_steps
+            )
+            return fresh_jac, fresh_lu, fresh_piv, retried, retried_converged
+
+        jac, lu, piv, correction, converged = jax.lax.cond(
+            retry,
+            refresh,
+            lambda: (jac, lu, piv, correction, converged),
+        )
+        did_jac = need_jac | blame_jacobian
+        did_lu = need_lu | retry
+
+        flat_y1 = flat_pred + correction
+        y1 = unravel(flat_y1)
+        if self.suppress_algebraic_error:
+            error_mask = flat_mask
+        else:
+            error_mask = None
+        y_error = unravel(
+            _suppress(jnp.asarray(error_const, dtype)[order] * correction, error_mask)
+        )
+        # A failed Newton solve is reported as an infinite error estimate so that
+        # the controller rejects the step, rather than as a non-successful
+        # `RESULTS`, which `diffrax` would treat as a failure of the whole solve.
+        y_error = jtu.tree_map(lambda e: jnp.where(converged, e, jnp.inf), y_error)
+        # If `init` failed, accept this step so that diffrax sees its result;
+        # diffrax ignores the result of a rejected step.
+        y_error = jtu.tree_map(
+            lambda e: jnp.where(state.initial_failed, 0, e), y_error
+        )
+
+        d_array = _update_difference_array(d_array, order, unravel(correction))
+
+        if self._variable_order:
+            next_order, n_equal_steps = self._select_order(
+                d_array,
+                order,
+                correction,
+                flat_y1,
+                error_const,
+                rtol,
+                atol,
+                n_equal_steps,
+                dtype,
+                error_mask,
+            )
+            # The order that produced this error estimate is the one the controller
+            # needs, not the one we have just chosen for the next step.
+            y_error = ErrorWithOrder(error=y_error, order=order.astype(dtype))
+        else:
+            # At fixed order the only movement is the opening ramp: a k-step formula
+            # needs k points of history. A rejected step rolls the state back, so the
+            # ramp is automatically correct -- a step that did not happen does not
+            # advance the order.
+            next_order = jnp.minimum(order + 1, self.bdf_order)
+
+        new_state = _BDFState(
+            d_array=d_array,
+            order=next_order,
+            n_equal_steps=n_equal_steps,
+            h_prev=control,
+            jac=jac,
+            lu=lu,
+            piv=piv,
+            c_at_last_lu=jnp.where(did_lu, c, state.c_at_last_lu),
+            steps_since_lu=jnp.where(did_lu, 0, state.steps_since_lu + 1).astype(
+                jnp.int32
+            ),
+            steps_since_jac=jnp.where(did_jac, 0, state.steps_since_jac + 1).astype(
+                jnp.int32
+            ),
+            initial_failed=state.initial_failed,
+        )
+        dense_info = dict(y0=y0, y1=y1)
+        result = RESULTS.where(
+            state.initial_failed, RESULTS.nonlinear_divergence, RESULTS.successful
+        )
+        return y1, y_error, dense_info, new_state, result
+
+    def _select_order(
+        self,
+        d_array,
+        order,
+        correction,
+        flat_y1,
+        error_const,
+        rtol,
+        atol,
+        n_equal_steps,
+        dtype,
+        error_mask,
+    ):
+        """scipy's order heuristic: compare the error at this order, one below, one above.
+
+        The two rows of the difference array beyond `order` exist for exactly this.
+        After `_update_difference_array`, row `order` carries what the error would
+        have been one order lower and row `order + 2` what it would be one order
+        higher, so all three candidates are available without re-solving anything.
+        """
+        table = jnp.asarray(error_const, dtype)
+        scale = atol + rtol * jnp.abs(flat_y1)
+
+        def row(index):
+            flat_row = ravel_pytree(jtu.tree_map(lambda x: x[index], d_array))[0]
+            return _suppress(flat_row, error_mask)
+
+        correction = _suppress(correction, error_mask)
+        error = rms_norm(table[order] * correction / scale)
+        error_lower = rms_norm(table[order - 1] * row(order) / scale)
+        error_higher = rms_norm(table[order + 1] * row(order + 2) / scale)
+        # Candidates outside 1..cap are ruled out with an infinite error, which maps
+        # to a zero factor and so can never win the argmax below.
+        error_lower = jnp.where(order > 1, error_lower, jnp.inf)
+        error_higher = jnp.where(order < self._order_cap, error_higher, jnp.inf)
+
+        exponents = jnp.stack([order, order + 1, order + 2]).astype(dtype)
+        candidates = jnp.stack([error_lower, error, error_higher])
+        factors = candidates ** (-1 / exponents)
+        delta = jnp.argmax(factors) - 1
+
+        # scipy holds the order still until `order + 1` steps of equal size have been
+        # taken. Both the error constants and the difference array assume an equally
+        # spaced history, so changing order before that compares invalid estimates.
+        #
+        # scipy additionally freezes the step size between order changes, so that the
+        # count reliably accumulates. That was tried here and measured worse than
+        # letting `BDFController`'s deadband govern the step size (160ms against 124ms
+        # on enzax's methionine model from a bad guess), so the two stay decoupled: the
+        # deadband already holds the step size still often enough for the gate to open.
+        allowed = n_equal_steps >= order + 1
+        next_order = jnp.where(
+            allowed, jnp.clip(order + delta, 1, self._order_cap), order
+        )
+        # Restart the count whenever the gate opens, whether or not the order actually
+        # moved. This is scipy's behaviour and it measured slightly better than only
+        # restarting on a real change (423 steps against 454 on enzax's methionine
+        # model from a bad guess).
+        n_equal_steps = jnp.where(allowed, 0, n_equal_steps)
+        return next_order.astype(jnp.int32), n_equal_steps.astype(jnp.int32)
+
+    def _restart(
+        self,
+        terms,
+        t0,
+        y0,
+        args,
+        control,
+        d_array,
+        order,
+        n_equal_steps,
+        steps_since_jac,
+        made_jump,
+        mask,
+    ):
+        """Rebuild the history from scratch after a discontinuity."""
+        f0 = _zero_algebraic(terms.vf_prod(t0, y0, args, control), mask)
+
+        def reset(x, y, f):
+            fresh = jnp.zeros_like(x).at[0].set(y).at[1].set(f)
+            return jnp.where(made_jump, fresh, x)
+
+        d_array = jtu.tree_map(reset, d_array, y0, f0)
+        order = jnp.where(made_jump, 1, order)
+        n_equal_steps = jnp.where(made_jump, 0, n_equal_steps)
+        # The vector field changed, so the cached Jacobian is suspect too.
+        steps_since_jac = jnp.where(
+            made_jump, self.jac_reuse_steps, steps_since_jac
+        ).astype(jnp.int32)
+        return d_array, order, n_equal_steps, steps_since_jac
+
+
+BDF.__init__.__doc__ = """**Arguments:**
+
+- `bdf_order`: pin the formula to a fixed order between 1 and 5, or leave it as
+    `None` (the default) to vary the order between 1 and `max_order` as scipy and
+    CVODE do. Variable order needs [`diffrax_bdf.BDFController`][]; a fixed order
+    works with any adaptive step size controller. Note that a fixed high order is
+    less stable: BDF4 is A(73.35 degrees)-stable and BDF5 only
+    A(51.84 degrees)-stable, so both can go unstable on strongly oscillatory systems.
+- `max_order`: the highest order the variable-order heuristic may select.
+- `use_ndf`: whether to apply the Shampine & Reichelt NDF correction. `False`
+    gives the classical BDF coefficients, which is the family SUNDIALS CVODE
+    implements.
+- `root_finder`: only its `rtol` and `atol` are used, to set the Newton convergence
+    tolerance. The iteration itself is a chord iteration against a stored
+    factorisation, which is what allows that factorisation to be reused between
+    steps.
+- `root_find_max_steps`: maximum chord iterations per step. scipy uses 4, CVODE 3.
+- `lu_reuse_steps`: refactorise `I - c * J` after this many accepted steps even if
+    `c` has not moved. CVODE's `MSBP`.
+- `jac_reuse_steps`: re-evaluate the Jacobian after this many accepted steps.
+    CVODE's `MSBJ`. Much larger than `lu_reuse_steps` because a Jacobian costs an
+    automatic-differentiation pass where a refactorisation does not.
+- `max_gamma_change`: refactorise when `c` has moved by more than this relative
+    amount since the last factorisation. CVODE's `DGMAX_LSETUP`.
+- `suppress_algebraic_error`: leave the algebraic components of a
+    `SemiExplicitDAETerm` out of the local error estimate and the order
+    selection, as SUNDIALS IDA's `suppressalg` does. Each step's Newton
+    iteration still has to converge on them.
+- `correct_initial_algebraic`: before the first step, solve the constraints of
+    a `SemiExplicitDAETerm` for its algebraic components, holding the others
+    fixed. If this fails, the solve stops with `RESULTS.nonlinear_divergence`.
+    Turn it off when `y0` is known to be consistent.
+- `initial_algebraic_max_steps`: the maximum number of Newton steps in that
+    correction.
+"""

diffrax_bdf/_coeffs.py

@@ -0,0 +1,101 @@
+"""BDF/NDF coefficient tables and backward-difference array manipulation.
+
+The representation follows `scipy.integrate.BDF` (in turn following the MATLAB
+`ode15s` of Shampine & Reichelt 1997): the solution history is stored as an array
+of backward differences `D`, with `D[j]` the `j`-th backward difference of `y`.
+
+Everything here is written for static shapes: arrays are always sized for
+`MAX_ORDER` and the active block is selected by masking against a *traced* order,
+never by slicing with it. This is what makes the solver usable inside the
+`lax.while_loop` that `diffrax` runs its integration in.
+"""
+
+import jax.numpy as jnp
+import jax.tree_util as jtu
+import numpy as np
+from jaxtyping import Array, Float, PyTree
+
+
+MAX_ORDER = 5
+
+# Rows of the difference array. Rows `0..order` hold the differences themselves;
+# row `order + 1` holds the accepted Newton correction `d`, and row `order + 2`
+# holds `d - D_prev[order + 1]`. The last two feed the order-up error estimate.
+D_ROWS = MAX_ORDER + 3
+
+# Shampine & Reichelt's NDF correction. `kappa = 0` recovers the classical BDF
+# formulae (which is what SUNDIALS CVODE uses); see `make_tables`.
+KAPPA_NDF = np.array([0.0, -0.1850, -1 / 9, -0.0823, -0.0415, 0.0])
+KAPPA_BDF = np.zeros(MAX_ORDER + 1)
+
+
+def make_tables(use_ndf: bool = True):
+    """Return `(gamma, alpha, error_const)` as float64 numpy arrays.
+
+    With `use_ndf=False` this reduces to the classical BDF coefficients,
+    `alpha[k] == sum_{j=1..k} 1/j` and `error_const[k] == 1/(k + 1)`, which is the
+    formula family CVODE implements (its `gamma` is our `h / alpha[order]`).
+    """
+    kappa = KAPPA_NDF if use_ndf else KAPPA_BDF
+    gamma = np.hstack((0.0, np.cumsum(1 / np.arange(1, MAX_ORDER + 1))))
+    alpha = (1 - kappa) * gamma
+    error_const = kappa * gamma + 1 / np.arange(1, MAX_ORDER + 2)
+    return gamma, alpha, error_const
+
+
+def compute_R(factor, order) -> Float[Array, "6 6"]:
+    """The `(MAX_ORDER + 1, MAX_ORDER + 1)` difference-rescaling matrix.
+
+    The top-left `(order + 1, order + 1)` block equals scipy's
+    `compute_R(order, factor)`; outside that block the result is the identity, so
+    that multiplying by it leaves the inactive rows of `D` untouched.
+
+    Note the unmasked entries `M[i, j] = (i - 1 - factor * j) / i` do not depend on
+    `order` at all, and `cumprod` along axis 0 only ever reads rows `<= i`. That is
+    why building the full matrix and masking gives exactly the same active block as
+    scipy's order-sized construction.
+    """
+    idx = jnp.arange(MAX_ORDER + 1)
+    i = jnp.arange(1, MAX_ORDER + 1)[:, None]
+    j = jnp.arange(1, MAX_ORDER + 1)[None, :]
+    # `jnp.result_type(float)` respects the x64 config, so this is float64 when
+    # x64 is enabled rather than being dragged down to float32 by a literal.
+    dtype = jnp.result_type(factor, jnp.result_type(float))
+    body = (i - 1 - factor * j) / i
+    m = jnp.zeros((MAX_ORDER + 1, MAX_ORDER + 1), dtype=dtype)
+    m = m.at[1:, 1:].set(body.astype(dtype))
+    m = m.at[0].set(1.0)
+    r = jnp.cumprod(m, axis=0)
+    active = (idx[:, None] <= order) & (idx[None, :] <= order)
+    return jnp.where(active, r, jnp.eye(MAX_ORDER + 1, dtype=dtype))
+
+
+def change_D(
+    d_array: PyTree[Float[Array, "8 ..."]], order, factor
+) -> PyTree[Float[Array, "8 ..."]]:
+    """Rescale the difference array for a change of step size by `factor`.
+
+    This is scipy's `change_D`, and is exactly SUNDIALS' `cvRescale` expressed in
+    the backward-difference basis rather than the Nordsieck one: writing `S` for the
+    (fixed, order-only) change of basis, `(R @ U).T == inv(S) @ diag(factor**j) @ S`.
+    """
+    ru = compute_R(factor, order) @ compute_R(1.0, order)
+    return jtu.tree_map(lambda x: _apply_rows(ru.T, x), d_array)
+
+
+def _apply_rows(mat: Float[Array, "6 6"], x: Float[Array, "8 ..."]):
+    """Left-multiply the first `MAX_ORDER + 1` rows of `x` by `mat`."""
+    x = jnp.asarray(x)
+    head = jnp.tensordot(mat.astype(x.dtype), x[: MAX_ORDER + 1], axes=1)
+    return x.at[: MAX_ORDER + 1].set(head)
+
+
+def weighted_row_sum(coeffs: Float[Array, " 6"], d_array: PyTree) -> PyTree:
+    """Contract `coeffs` against rows `0..MAX_ORDER` of each leaf of `d_array`."""
+    def contract(x):
+        x = jnp.asarray(x)
+        return jnp.tensordot(
+            jnp.asarray(coeffs).astype(x.dtype), x[: MAX_ORDER + 1], axes=1
+        )
+
+    return jtu.tree_map(contract, d_array)

diffrax_bdf/_controller.py

@@ -0,0 +1,73 @@
+"""Step-size control with CVODE's deadband.
+
+A BDF only pays for itself if the factorisation of `I - c * J` can be reused across
+steps, and that requires `c`, and hence the step size, to stay put. diffrax's stock
+`PIDController` adjusts the step on every single step, so `c` changes every step and
+the factorisation has to be redone every step -- at which point a BDF is doing the
+same amount of linear algebra as an ESDIRK and has given up its advantage.
+
+CVODE solves this with a deadband (`cvSetEta`): a successful step never changes the
+step size unless the proposed growth factor is at least `eta_max_fx` (1.5 by
+default), and never shrinks it at all. Rejected steps are unaffected. This gives up
+a little step-size optimality and buys long runs at constant `c`.
+"""
+
+import diffrax
+import jax.numpy as jnp
+from jaxtyping import PyTree
+
+from ._bdf import ErrorWithOrder
+
+
+class BDFController(diffrax.PIDController):
+    """`diffrax.PIDController` with CVODE's step-size deadband and variable-order support.
+
+    All of `PIDController`'s arguments apply. The added `eta_max_fx` is the growth
+    factor below which an accepted step keeps its step size unchanged.
+
+    This controller is also required by a variable-order [`diffrax_bdf.BDF`][].
+    `diffrax` asks a solver for its error order only once, so a solver whose order
+    changes from step to step has to send the current one along with the error
+    estimate; this controller knows how to read that.
+    """
+
+    eta_max_fx: float = 1.5
+
+    def adapt_step_size(
+        self,
+        t0,
+        t1,
+        y0: PyTree,
+        y1_candidate: PyTree,
+        args: PyTree,
+        y_error: PyTree,
+        error_order,
+        controller_state,
+    ):
+        if isinstance(y_error, ErrorWithOrder):
+            # The step-size exponent is `1 / (p + 1)` for a method of order `p`, so at
+            # variable order it has to follow the order that actually produced this
+            # estimate rather than a value fixed before the solve began.
+            error_order = y_error.order + 1
+            y_error = y_error.error
+        keep_step, next_t0, next_t1, made_jump, state, result = super().adapt_step_size(
+            t0, t1, y0, y1_candidate, args, y_error, error_order, controller_state
+        )
+        # Work in signed step sizes throughout so that backwards integration, where
+        # `diffrax` hands us a negative step, needs no special case.
+        step = t1 - t0
+        proposed = next_t1 - next_t0
+        growth = proposed / jnp.where(step == 0, 1.0, step)
+        frozen = keep_step & (growth < self.eta_max_fx)
+        next_t1 = jnp.where(frozen, next_t0 + step, next_t1)
+        return keep_step, next_t0, next_t1, made_jump, state, result
+
+
+BDFController.__init__.__doc__ = """**Arguments:**
+
+As [`diffrax.PIDController`][], plus:
+
+- `eta_max_fx`: an accepted step keeps its step size unless the controller wants to
+    grow it by at least this factor. Raising it reuses factorisations for longer at
+    the cost of less optimal steps; setting it to 1 recovers `PIDController`.
+"""

diffrax_bdf/_newton.py

@@ -0,0 +1,96 @@
+"""Chord iteration against a stored LU factorisation.
+
+This replaces the Optimistix root find used elsewhere in diffrax for one reason: we
+need to keep the factorisation of `I - c * J` in the solver state and reuse it
+across *steps*. Optimistix's root finders build their linearisation internally and
+hold it as a `lineax` operator with a captured jaxpr, which cannot be carried
+through a `lax.while_loop`. A pair of plain `(lu, piv)` arrays can.
+
+The iteration and its convergence test follow `scipy.integrate.BDF`'s
+`solve_bdf_system` closely, including the divergence test, so that behaviour is
+comparable step by step.
+"""
+
+import equinox.internal as eqxi
+import jax.numpy as jnp
+import jax.scipy.linalg as jsl
+from jaxtyping import Array, Float, Int
+
+
+def rms_norm(x: Float[Array, " n"]) -> Float[Array, ""]:
+    """scipy's `norm`: the 2-norm scaled by the square root of the size."""
+    return jnp.linalg.norm(x) / jnp.sqrt(x.size)
+
+
+def newton_tolerance(rtol, dtype) -> Float[Array, ""]:
+    """scipy's `newton_tol`.
+
+    The `eps` must come from the *runtime* dtype. Hard-coding the float64 value
+    would, in float32 at `rtol=1e-6`, give a tolerance above 1 and make the
+    iteration "converge" immediately on its first step.
+    """
+    eps = jnp.finfo(dtype).eps
+    return jnp.maximum(10 * eps / rtol, jnp.minimum(0.03, jnp.sqrt(rtol)))
+
+
+def chord_solve(
+    residual,
+    guess: Float[Array, " n"],
+    lu: Float[Array, "n n"],
+    piv: Int[Array, " n"],
+    scale: Float[Array, " n"],
+    tol,
+    max_iter: int,
+):
+    """Solve `residual(d) == 0` by chord iteration, holding `lu` fixed.
+
+    Returns `(correction, converged, num_iterations)`.
+    """
+    zero = jnp.zeros((), dtype=guess.dtype)
+
+    def cond(carry):
+        _, _, _, converged, diverged = carry
+        return jnp.invert(converged | diverged)
+
+    def body(carry):
+        step, correction, previous_norm, _, _ = carry
+        residual_value = residual(correction)
+        delta = jsl.lu_solve((lu, piv), residual_value)
+        delta_norm = rms_norm(delta / scale)
+
+        # On the first iteration there is no previous norm to compare against, so
+        # neither the divergence nor the convergence rate test applies.
+        has_history = step > 0
+        safe_previous = jnp.where(previous_norm > 0, previous_norm, 1.0)
+        rate = jnp.where(has_history, delta_norm / safe_previous, zero)
+        # Guard the `1 - rate` denominators: they are only meaningful when the
+        # iteration is contracting, and the `rate >= 1` test below catches the rest.
+        safe_rate = jnp.minimum(rate, 0.99)
+        remaining = jnp.maximum(max_iter - step, 1)
+
+        diverging = has_history & (
+            (rate >= 1)
+            | (safe_rate**remaining / (1 - safe_rate) * delta_norm > tol)
+        )
+        diverging = diverging | jnp.invert(jnp.all(jnp.isfinite(residual_value)))
+        # scipy breaks out *before* applying a diverging update.
+        correction = jnp.where(diverging, correction, correction + delta)
+        converged = jnp.invert(diverging) & (
+            (delta_norm == 0)
+            | (has_history & (safe_rate / (1 - safe_rate) * delta_norm < tol))
+        )
+        return step + 1, correction, delta_norm, converged, diverging
+
+    init = (
+        jnp.asarray(0, jnp.int32),
+        guess,
+        zero,
+        jnp.asarray(False),
+        jnp.asarray(False),
+    )
+    # `kind="bounded"` keeps this reverse-mode differentiable, which a plain
+    # `lax.while_loop` would not be.
+    steps, correction, _, converged, _ = eqxi.while_loop(
+        cond, body, init, max_steps=max_iter, kind="bounded"
+    )
+    return correction, converged, steps

diffrax_bdf/_term.py

@@ -0,0 +1,72 @@
+from collections.abc import Callable
+
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from diffrax import AbstractTerm, ODETerm
+from diffrax._term import WrapTerm
+from jax.flatten_util import ravel_pytree
+from jaxtyping import Array, Bool, PyTree
+
+
+class SemiExplicitDAETerm(AbstractTerm):
+    """A semi-explicit differential algebraic equation, for
+    [`diffrax_bdf.BDF`][].
+
+    Where `algebraic_mask` is false, the vector field gives derivatives,
+    `dy/dt = f(t, y)`; where it is true, it gives constraints, `0 = g(t, y)`.
+    The constraints must have index 1: their Jacobian with respect to the
+    algebraic components must be nonsingular.
+
+    Only `BDF` can solve this term. Other diffrax solvers would integrate the
+    constraints as if they were derivatives, so they raise an error instead.
+    """
+
+    vector_field: Callable
+    algebraic_mask: PyTree
+
+    def _unsupported(self, *args, **kwargs):
+        raise TypeError(
+            "A `SemiExplicitDAETerm` can only be solved with `diffrax_bdf.BDF`, "
+            "which treats its algebraic components as constraints. Other solvers "
+            "would integrate them as if they were derivatives."
+        )
+
+    def vf(self, t, y, args):
+        return self.vector_field(t, y, args)
+
+    vf_prod = _unsupported
+    prod = _unsupported
+
+    def contr(self, t0, t1, **kwargs):
+        return t1 - t0
+
+
+SemiExplicitDAETerm.__init__.__doc__ = """**Arguments:**
+
+- `vector_field`: a function `(t, y, args) -> f` returning a pytree with the
+    structure of `y`.
+- `algebraic_mask`: which components of `y` are algebraic. Each leaf is
+    broadcast against the matching leaf of `y`, so `(False, True)` marks the
+    second element of `y = (z, u)` as algebraic.
+"""
+
+
+def split_dae_term(terms) -> tuple[AbstractTerm, PyTree | None]:
+    """Return an ODE view of `terms` and its algebraic mask, or `None` for an ODE
+    term."""
+    if isinstance(terms, SemiExplicitDAETerm):
+        return ODETerm(terms.vector_field), terms.algebraic_mask
+    if isinstance(terms, WrapTerm) and isinstance(terms.term, SemiExplicitDAETerm):
+        ode_term = ODETerm(terms.term.vector_field)
+        return WrapTerm(ode_term, terms.direction), terms.term.algebraic_mask
+    return terms, None
+
+
+def flat_algebraic_mask(mask: PyTree, y: PyTree) -> Bool[Array, " n"]:
+    """Broadcast the mask against `y` and flatten it in `ravel_pytree`'s order."""
+    broadcast = jtu.tree_map(
+        lambda m, leaf: jnp.broadcast_to(jnp.asarray(m, bool), jnp.shape(leaf)),
+        mask,
+        y,
+    )
+    return ravel_pytree(broadcast)[0]

diffrax_bdf-0.1.0.dist-info/METADATA

@@ -0,0 +1,201 @@
+Metadata-Version: 2.3
+Name: diffrax-bdf
+Version: 0.1.0
+Summary: Diffrax-compatible backward differentiation formula ODE solver
+Author-email: groves.teddy@gmail.com
+License: MIT
+Requires-Dist: diffrax>=0.7.2
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# diffrax-bdf
+
+A [backward differentiation formula](https://en.wikipedia.org/wiki/Backward_differentiation_formula)
+solver for [diffrax](https://github.com/patrick-kidger/diffrax), aimed at stiff
+steady-state problems — in particular the kinetic reaction networks in
+[enzax](https://github.com/dtu-qmcm/enzax).
+
+diffrax has no BDF ([#8](https://github.com/patrick-kidger/diffrax/issues/8)). This
+package supplies one as an ordinary `AbstractSolver` subclass, with no patching of
+diffrax itself.
+
+```python
+import diffrax
+import jax.numpy as jnp
+from diffrax_bdf import BDF, BDFController
+
+sol = diffrax.diffeqsolve(
+    diffrax.ODETerm(vector_field),
+    BDF(),                       # variable order 1-5, as scipy and CVODE do
+    t0=0.0,
+    t1=jnp.inf,
+    dt0=1e-6,
+    y0=guess,
+    args=parameters,
+    stepsize_controller=BDFController(rtol=1e-8, atol=1e-10, dtmax=1e6),
+    event=diffrax.Event(diffrax.steady_state_event()),
+    adjoint=diffrax.ImplicitAdjoint(),
+    saveat=diffrax.SaveAt(t1=True),
+)
+```
+
+Pin the order with `BDF(bdf_order=3)` to use any adaptive step size controller,
+including diffrax's own `PIDController`.
+
+## Status
+
+Variable order 1-5, adaptive step size, reuse of the Jacobian and its
+factorisation across steps, and semi-explicit index-1 DAEs.
+
+## Method
+
+The formulation follows `scipy.integrate.BDF` and MATLAB's `ode15s`: the history is
+a backward-difference array rescaled on each change of step size, with the Shampine
+& Reichelt NDF correction available via `use_ndf` (`False` gives the classical BDF
+coefficients, the family SUNDIALS CVODE implements). The order is chosen each step
+by scipy's heuristic, comparing the error estimate at the current order against what
+it would have been one order lower and one order higher.
+
+The linear-algebra reuse policy is CVODE's, because that is where CVODE's advantage
+over scipy actually lives. `Kvaerno5` already computes one Jacobian per step and
+reuses it across its five stages, so "one Newton solve per step instead of five
+stage solves" buys only a few linear solves. The real saving is keeping `J` and the
+factorisation of `I - c * J` across *steps*, gated on two separate counters:
+
+- the factorisation is rebuilt when `c` has moved by more than `max_gamma_change`
+  (0.3), or after `lu_reuse_steps` (20) accepted steps;
+- the Jacobian itself is rebuilt only after `jac_reuse_steps` (51) accepted steps,
+  so a refactorisation reassembles `I - c * J` from the stored `J` without touching
+  the vector field.
+
+`BDFController` adds CVODE's step-size deadband on top of `diffrax.PIDController`:
+an accepted step keeps its step size unless the controller wants to grow it by at
+least 1.5x, and never shrinks it. Holding `h` still holds `c` still, which is what
+lets the factorisation survive.
+
+It also carries the order. diffrax asks a solver for its error order once, before
+the solve (`solver.error_order(terms)` sees neither state nor time), so a
+variable-order solver cannot report a per-step order through that route. The only
+other channel to the controller is `y_error`, which the integration loop treats
+opaquely, so the order travels inside it.
+
+Everything is written for static shapes -- the order is a *traced* integer used to
+mask fixed-size buffers -- because diffrax carries the solver state through a
+`lax.while_loop`.
+
+## Order selection on a steady-state solve
+
+Integrating the `n=3` problem below to steady state, the order climbs through the
+transient and then falls back as the residual collapses and the step grows:
+
+```
+order:  1 2 2 3 4 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 4 4 4 4 4 4 4 2 1 1 1 1 1 1 1
+step:   1.2e-07  ------------------------------------------------>  7.9e+05
+```
+
+The tail is the useful part: BDF1 with an enormous step *is* Newton's method on
+`f(y) = 0`, so a correct variable-order BDF turns into a steady-state solver by
+itself as it converges.
+
+## Differential algebraic equations
+
+`BDF` also solves semi-explicit index-1 DAEs, where some components of `y` obey differential equations and the rest obey algebraic constraints:
+
+```python
+term = SemiExplicitDAETerm(vector_field, algebraic_mask=(False, True))
+```
+
+Here `vector_field` returns `(f, g)` for `y = (z, u)`, meaning `dz/dt = f` and `0 = g`. The corrector solves the constraints in the same chord iteration as the differential equations, with iteration matrix `M − cJ` where `M` is 1 on differential components and 0 on algebraic ones, and reuses its factorisation across steps as for an ODE.
+
+- By default `init` corrects an inconsistent initial value by a Newton solve on the constraints. diffrax keeps its own copy of `y0`, so `SaveAt(t0=True)` shows the value you passed, not the corrected one.
+- `suppress_algebraic_error=True` drops the algebraic components from the error estimate, which suits problems where only `z` is of interest. On a rapid-equilibrium steady state it saved 29% of steps.
+- The steady-state event and `ImplicitAdjoint` work unchanged: the steady state is the root of `(f, g)`.
+
+## Caveats
+
+- **Enable `jax_enable_x64`.** High-order backward differences suffer heavy
+  cancellation; above order 2 they are noise in float32.
+- **Set `dtmax` when integrating to `t1=inf`.** As the residual collapses the
+  controller's growth factor saturates and the step grows without bound until it
+  overflows the time variable.
+- Variable order requires `BDFController`; a fixed `bdf_order` works with any
+  adaptive controller.
+- The Jacobian is formed densely, so this suits small to medium systems.
+- Requires a term whose control is a scalar, i.e. `ODETerm` or `SemiExplicitDAETerm`.
+
+## Performance
+
+Steady-state solves to `rtol=1e-8`, CPU, float64, medians of interleaved rounds.
+Compare within a block, not across them -- absolute timings drift between runs.
+
+**Whether BDF beats `Kvaerno5` depends entirely on the problem.** Where `Kvaerno5`'s
+step control struggles, BDF wins by a wide margin; where it does not, `Kvaerno5`
+wins on step count.
+
+| problem | solver | steps | rejected | median |
+|---|---|---|---|---|
+| fast equilibrium, n=3 | `Kvaerno5` | 7784 | 4563 | 192.4 ms |
+| fast equilibrium, n=3 | `BDF(bdf_order=5)` | 276 | 14 | 3.5 ms |
+| fast equilibrium, n=3 | `BDF()` variable | 231 | 6 | **3.3 ms** |
+| stiff random network, n=18 | `Kvaerno5` | 212 | 14 | **23.6 ms** |
+| stiff random network, n=18 | `BDF(bdf_order=5)` | 593 | 11 | 28.8 ms |
+| stiff random network, n=18 | `BDF()` variable | 589 | 12 | 31.7 ms |
+| stiff random network, n=60 | `Kvaerno5` | 193 | 7 | **73.6 ms** |
+| stiff random network, n=60 | `BDF(bdf_order=5)` | 586 | 6 | **145.2 ms** |
+| stiff random network, n=60 | `BDF()` variable | 580 | 7 | 153.7 ms |
+
+The first problem is the interesting one: `Kvaerno5` rejects 59% of its steps, where
+BDF rejects 2-5%.
+
+Variable order helps where the order genuinely needs to move -- 231 steps against
+276 for the best fixed order on the steady-state problem. On the random networks,
+where order 5 is right throughout, it matches fixed order 5 on step count and pays
+about 8% per step for the heuristic. Reuse is worth a further 16-18%.
+
+## On enzax's models
+
+The motivating case. Measured against
+[enzax](https://github.com/dtu-qmcm/enzax)'s own steady-state configuration
+(`t1=inf`, `steady_state_event`, `ImplicitAdjoint`, `rtol=atol=1e-9`), each solver
+with the step size controller that suits it. Medians of interleaved rounds.
+
+| model | n | regime | `Kvaerno5` (enzax today) | `BDF()` | speedup |
+|---|---|---|---|---|---|
+| methionine | 5 | bad guess | 555 ms (649 steps, 145 rejected) | 158 ms (423, 52) | **3.5x** |
+| methionine | 5 | warm guess, 2% parameter jitter | 30.7 ms (38 steps) | 23.3 ms (88) | 1.3x |
+| methionine | 5 | `jacrev` through the solve | 27.9 ms | 22.4 ms | 1.25x |
+| conserved moiety | 5 | bad guess | 191 ms (225, 30) | 53.7 ms (293, 4) | **3.6x** |
+| conserved moiety | 5 | warm guess, 2% jitter | 58.3 ms (68) | 31.4 ms (165) | 1.9x |
+| conserved moiety | 5 | `jacrev` | 53.0 ms | 27.6 ms | 1.9x |
+| glycolysis | 18 | bad guess | 2302 ms (1007, 247) | 271 ms (821, 70) | **8.5x** |
+| glycolysis | 18 | guess perturbed 20% | 491 ms (212, 38) | 108 ms (277, 18) | **4.6x** |
+| glycolysis | 18 | `jacrev` | 427 ms | 104 ms | **4.1x** |
+
+"bad guess" is enzax's own `BAD_GUESS = full(shape, 0.01)`. Steady states agree with
+`Kvaerno5` to 2e-7 relative or better and gradients to 1e-5 or better.
+
+Note the step counts: BDF often takes *more* steps and is still several times
+faster, because a step is one Newton solve against a reused factorisation rather
+than five stage solves. It also rejects far fewer steps -- 70 against 247 on
+glycolysis from a bad guess. Compile time is 30-45% lower throughout, which matters
+when every leapfrog step of a NUTS chain pays it.
+
+The gap widens with model size, which is the encouraging part: 1.3x on the 5-state
+methionine model against 4-8x on the 18-state glycolysis one.
+
+### Tuning
+
+enzax's `PIDController(pcoeff=0.1, icoeff=0.3)` is tuned for `Kvaerno5`, which
+rejects heavily; it damps step size growth to `0.4 / error_order` where the default
+pure-I controller uses `1 / error_order`. BDF rejects very few steps and does not
+need the damping, so it should be left on `BDFController`'s defaults. Using enzax's
+tuning costs BDF roughly a factor of two -- 174 steps and 45.6 ms against 88 steps
+and 26.9 ms on methionine.
+
+## Testing
+
+`uv run pytest`. scipy's BDF is pure Python and is used as a white-box oracle: the
+coefficient tables, the difference-array rescaling, and a single pinned step are all
+asserted against it directly. The reuse policy, the deadband and the order-selection
+heuristic are tested directly as well, since without them the solver has no
+advantage over `Kvaerno5` at all.

diffrax_bdf-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+diffrax_bdf/__init__.py,sha256=FTsR3R3gYwjI6-uNlDMyM69pPDJ-o0DTsIaow5zNizU,249
+diffrax_bdf/_bdf.py,sha256=otW-MNeGdaTPgoUYK9F6hIYYJvY0XAvih2hJsuWbM4w,28757
+diffrax_bdf/_coeffs.py,sha256=1Zpm0S5mEa2QmHoLIPhmPbYzAxr_iC9r7HeQxrD2624,4406
+diffrax_bdf/_controller.py,sha256=yMm4WLsYpOo7Uc9FzVFTFnKTXRg6HuFxSOjU_TdXTaE,3072
+diffrax_bdf/_newton.py,sha256=0NFOwtZB2TePNNomm7aMBJTEpiwJkQbLwZ4rKyA7Jcc,3663
+diffrax_bdf/_term.py,sha256=IBqSvXnpbgt_albZizeB_pOqzsDkSRro7n0IdwVoV5A,2581
+diffrax_bdf-0.1.0.dist-info/WHEEL,sha256=3OmCxFA0ffBIsX4x0o8CdfSvZAOvO3mzEG8U_z5TvV4,80
+diffrax_bdf-0.1.0.dist-info/METADATA,sha256=iif1JYi7Hx18wypw6t4EGr_AX8DXb_vnl6k8sor284M,9985
+diffrax_bdf-0.1.0.dist-info/RECORD,,

diffrax_bdf-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.13.8
+Root-Is-Purelib: true
+Tag: py3-none-any