aquakin.esdirk_adjoint_solve#

aquakin.esdirk_adjoint_solve(rhs, y0, params, t_span, t_eval=None, *, solver=None, order=5, rtol=1e-06, atol=1e-09, dt0=1e-06, max_steps=200000, time_dependent=False, primal_rhs=None, jacobian_builder=None, forward_root_finder=None, factormax=None, low_memory=False, newton_iters=12)[source]#

Cap-free reverse-mode gradient through a high-order ESDIRK solve.

Like implicit_euler_adjoint_solve() but the forward uses a high-order ESDIRK method (default diffrax.Kvaerno5, matching the reactors), and the backward is the transposed-stage discrete adjoint of that method. The forward saves each step’s stage derivatives via diffrax dense output, so the backward reconstructs the stage values exactly by the Butcher linear combination Y_i = y_n + sum_j A[i,j]*k_j (the dense-output k is the dt-scaled stage increment) – no Newton recompute, which was the dominant backward cost – then the per-stage transposed solves accumulate the gradient. Finite for stiff models with no dtmax cap.

Trades memory for compute with low_memory=True (see below): the forward then stores only the step states (not the dense-output stage buffer, ~``s``x the trajectory) and the backward re-solves each step’s stages by Newton.

Parameters:

  • rhs (Callable) – As for implicit_euler_adjoint_solve().

  • y0 (Array) – As for implicit_euler_adjoint_solve().

  • params (Array) – As for implicit_euler_adjoint_solve().

  • t_span (tuple[float, float]) – As for implicit_euler_adjoint_solve().

  • t_eval (Array | None) – As for implicit_euler_adjoint_solve().

  • rtol (float) – As for implicit_euler_adjoint_solve().

  • atol (float) – As for implicit_euler_adjoint_solve().

  • dt0 (float) – As for implicit_euler_adjoint_solve().

  • max_steps (int) – As for implicit_euler_adjoint_solve().

  • solver (diffrax.AbstractSolver, optional) – The ESDIRK forward solver; must expose a Butcher tableau. Defaults to diffrax.Kvaerno5.

  • time_dependent (bool, optional) – If False (default) the right-hand side is taken to be autonomous (the reaction RHS is, for fixed conditions), so the backward pass evaluates it with a fixed time argument and the stage times do not enter. If True the field’s explicit time dependence (e.g. a time-varying influent) is handled exactly by carrying time in the state (_autonomize()), so the gradient is exact through a transient solve.

  • primal_rhs (callable, optional) – An alternate right-hand side used for the forward solve and the ``df/dy`` stage Jacobians – everything except the df/dtheta parameter vjp, which always uses rhs. It must produce the same values and the same ``df/dy`` as rhs (so the trajectory and the state-cotangent recurrence are unchanged); only the parameter derivative is taken from rhs. The use case is a RHS with a state-invariant but parameter-dependent sub-computation that is expensive to repeat – e.g. a plant’s recycle map M(params) – which can be evaluated once and reused for every Jacobian/stage call here, while rhs (which recomputes it) still supplies the exact dM/dtheta in the one place it is needed (the parameter vjp). Because the discrete adjoint takes its entire parameter gradient from that vjp and uses the stages / Jacobians only to propagate the state cotangent, the result is the exact gradient – bit-identical to using rhs everywhere when the cached sub-computation equals the recomputed one. The caller MUST jax.lax.stop_gradient() any params-derived value it closes over in primal_rhs (otherwise the closed-over tracer escapes the custom VJP). None (default) uses rhs for everything (the historic path).

  • jacobian_builder (callable, optional) – Builder (f, y) -> J for the per-stage df/dy Jacobian used in the backward pass (built once per reconstructed stage, for the transposed- stage solves), where f is the (autonomized) primal right-hand at fixed parameters. None (default) builds the dense Jacobian with jax.jacfwd (the historic path, bit-identical). A sparsity-colored builder – one Jacobian-vector product per color instead of one per state – cuts the dominant backward cost for a large, block-sparse plant: the backward is ~80% Jacobian builds. The builder MUST return a J equal to the dense Jacobian (the colored construction is exact when its sparsity pattern is a superset of the true nonzeros; the caller guards this at the start state and falls back to dense on a mismatch), so the discrete adjoint – and the gradient – is unchanged. Affects only the backward Jacobian build; the forward solve and the parameter vjp are untouched.

  • low_memory (bool, optional) – Trade memory for compute in the backward pass. False (default) saves each forward step’s dense-output stage increments and reconstructs the stages by the Butcher linear combination – fast, but the saved buffer is ~s``x the trajectory (shape ``(max_steps, s, n)). True instead stores only the step states and recomputes each step’s stage values in the backward pass by re-solving the ESDIRK stage equations (a fixed newton_iters-iteration Newton scan per implicit stage), so the (max_steps, s, n) dense buffer is never allocated. The recompute is a contraction through the same well-conditioned I - dt*gamma*J the forward inverts, so the reconstructed stages – and the gradient – match the saved-stage path to machine precision; the cost is roughly a second per-step stage solve. Use it when the dense-stage buffer is the binding memory constraint (a long, large-state plant solve). Guarded to the singly-diagonal ESDIRK shape it assumes (explicit first stage, constant implicit gamma – Kvaerno3 / Kvaerno5); any other tableau falls back to the saved-stage path with a RuntimeWarning.

  • newton_iters (int, optional) – Newton iterations used to recompute each implicit stage when low_memory=True (ignored otherwise). The default converges the well-conditioned stage equation to machine precision at the step sizes the adaptive forward selects.

  • order (int)

  • forward_root_finder (object | None)

  • factormax (float | None)

Returns:

Final state (n,) if t_eval is None, else states at t_eval (len(t_eval), n); carries the discrete-adjoint VJP w.r.t. y0 and params.

Return type:

jnp.ndarray