aquakin.implicit_euler_adjoint_solve#

aquakin.implicit_euler_adjoint_solve(rhs, y0, params, t_span, t_eval=None, *, rtol=1e-06, atol=1e-09, dt0=1e-06, max_steps=200000, time_dependent=False, primal_rhs=None)[source]#

Integrate over t_span with a cap-free discrete-adjoint reverse-mode rule.

The forward pass is a robust adaptive implicit-Euler diffrax solve. The custom VJP is the exact discrete adjoint of that solve, evaluated by a backward scan of bounded transposed linear solves – finite for stiff models at any step size, with no dtmax cap. The first-order s=1 special case of esdirk_adjoint_solve(); both share the _discrete_adjoint_solve() harness.

Parameters:

  • rhs (callable) – Vector field rhs(t, y, params) -> dy (the reactor RHS), differentiable in y and params.

  • y0 (jnp.ndarray) – Initial state, shape (n,).

  • params (jnp.ndarray) – Parameter vector, shape (n_params,).

  • t_span (tuple of float) – (t0, t1) integration interval.

  • t_eval (jnp.ndarray, optional) – Observation times at which to return the state. Must be strictly ascending and lie in (t0, t1] (a time equal to t0 returns y0). The forward is forced to step exactly onto these times so the adjoint is exact. None (default) returns only the final state.

  • rtol (float) – Forward-solve tolerances (the adaptive controller meets these; the adjoint is exact for whatever step sequence results).

  • atol (float) – Forward-solve tolerances (the adaptive controller meets these; the adjoint is exact for whatever step sequence results).

  • dt0 (float) – Initial step.

  • max_steps (int) – Maximum number of accepted steps (also the size of the saved trajectory the backward scan walks).

  • time_dependent (bool, optional) – If False (default) the right-hand side is assumed autonomous, as the reaction RHS is for fixed conditions. If True the field’s explicit time dependence (e.g. a time-varying influent) is handled exactly by carrying time in the state, so the gradient is exact through a transient solve. See _autonomize().

  • primal_rhs (callable, optional) – Fast alternate RHS for the forward solve and the df/dy Jacobian, while rhs supplies the df/dtheta vjp. See esdirk_adjoint_solve() for the full contract (it must match rhs in value and df/dy; only the parameter derivative is taken from rhs; stop_gradient any params-derived value closed over). None uses rhs throughout.

Returns:

If t_eval is None, the final state y(t1), shape (n,). Otherwise the states at t_eval, shape (len(t_eval), n). Either way the result carries the discrete-adjoint VJP w.r.t. y0 and params.

Return type:

jnp.ndarray