solve_ivp

Contents

solve_ivp#

class braintools.quad.solve_ivp(f, t_span, y0, *args, method='dopri5', rtol=1e-06, atol=1e-09, dt0=None, max_steps=4096, t_eval=None, checkpoint=True, **kwargs)#

Solve an initial value problem with an adaptive embedded Runge–Kutta method.

Integrates \(\dot y = f(y, t)\) from t_span[0] to t_span[1] with error control. Each attempt takes one step of an embedded pair (method), measures the error, accepts or rejects the step and picks the next step size with a PI controller. The attempts run in one fixed-length scan of max_steps iterations; once t1 is reached the remaining iterations are skipped through a done mask.

Parameters:

  • f (Callable) – Right-hand side f(y, t, *args, **kwargs) -> dy/dt with the same PyTree structure (and units [y]/[t]) as y. Note the (y, t) argument order (scipy uses (t, y)).

  • t_span (tuple of (float or Quantity)) – (t0, t1). t1 < t0 integrates backwards in time. May be traced.

  • y0 (PyTree) – Initial state. Leaves may be arrays or brainunit Quantities.

  • *args – Extra positional arguments forwarded to f.

  • method (str) – Embedded pair: 'dopri5' (Dormand–Prince 5(4), default), 'rk45' (Cash–Karp 5(4)), 'rkf45' (Fehlberg 4(5)), 'rk23' or 'bs32' (Bogacki–Shampine 3(2)), 'dopri8' (DOP853). The aliases 'rk45_dopri' and 'rk87_dopri' are accepted.

  • rtol (float) – Relative tolerance. Default 1e-6.

  • atol (float, Quantity or PyTree, optional) – Absolute tolerance. A float applies to the mantissas of every leaf; a Quantity is converted to each leaf’s unit; a PyTree matching y0 sets one tolerance per leaf. Default 1e-9.

  • dt0 (float or Quantity, optional) – Initial step size magnitude. If None (default) it is chosen by the Hairer–Wanner starting-step heuristic [1].

  • max_steps (int) – Static number of step attempts (accepted plus rejected). Default 4096.

  • t_eval (array or Quantity, optional) – 1-D array of output times (static length, any order, duplicates allowed) inside t_span. The steps are shortened so that every t_eval point is an accepted step (see Notes).

  • checkpoint (bool) – If True (default), the differentiated replay loop is a brainstate.transform.checkpointed_scan, so reverse-mode memory no longer scales with max_steps (at roughly 30% more gradient compute). False stores the residuals of every one of the max_steps iterations. Values and gradients are identical either way; forward-only evaluation is unaffected.

  • **kwargs – Extra keyword arguments forwarded to f.

Returns:

NamedTuple (ts, ys, n_steps, n_rejected, success). Without t_eval, ts/ys hold the max_steps + 1 slots of accepted steps (ts[0] == t0), padded after n_steps with the last accepted value; ys[n_steps] is the final state. With t_eval, ts is t_eval and ys the interpolated states.

Return type:

ODESolution

Raises:

  • KeyError – If method is not a registered ODE method.

  • ValueError – If method has no embedded error estimate, t_span is not a pair, max_steps < 1, t_eval is not 1-D or (when concrete) lies outside t_span, dt0 is (concretely) zero, or atol does not match y0.

  • brainunit.UnitMismatchError – If times, tolerances or f’s output have units incompatible with t_span and y0.

See also

odeint

Fixed-step trajectory driver.

scipy.integrate.solve_ivp

The SciPy counterpart.

Notes

Error norm. The step error err returned by the embedded pair is measured on the mantissas as

\[\|e\| = \sqrt{\frac{1}{N}\sum_i \left(\frac{e_i}{\mathrm{atol}_i + \mathrm{rtol}\,\max(|y_i|, |y^{new}_i|)}\right)^2},\]

and the step is accepted when \(\|e\| \le 1\).

Step-size control. After an accepted step the PI controller of [1] (Sec. IV.2) proposes \(h_{new} = h \cdot s\,\|e_n\|^{-0.7/k} \|e_{n-1}\|^{0.4/k}\) with \(k\) one more than the order of the embedded error estimate, safety \(s = 0.9\) and the factor bounded to [0.2, 10] (no growth right after a rejection). A rejected step is retried with \(h \cdot \max(0.2, s\,\|e\|^{-1/k})\). Non-finite steps are rejected. If the step size falls below 10 * eps * |t| the solve stops with success=False; the same happens when max_steps attempts are exhausted. In both cases the outputs hold the last accepted state (no NaN padding).

Two passes. Pass 1 runs the adaptive loop (a fixed-length scan of max_steps attempts, with a cond skipping the iterations after done) on stop-gradient copies of all inputs and records, per attempt, whether it was accepted and its step size. Pass 2 replays only the accepted steps (without the error estimate) on the differentiated path. So each accepted step is computed twice, but a rejected attempt, even one that produced NaN, never enters a gradient, and recomputation under checkpoint=True cannot change a decision.

Gradients. The step sizes chosen by the controller are constants, i.e. gradients are those of the solution on the accepted time grid (“discretise-then-optimise”); the steps that land on t1 or a t_eval point carry the dependence on those times.

Output at t_eval. A cubic Hermite interpolant between accepted steps is only \(O(h^4)\) accurate, which for 5th- and 8th-order pairs is far worse than the step error. The solver therefore never steps over a t_eval point: a step that would cross one is shortened to land on it (without shrinking the next proposed step), and that step’s state is the output, with the accuracy of the solver itself. This costs about one extra step per t_eval point, which counts towards max_steps. Points beyond the last accepted time of a failed solve return the last accepted state.

Memory. With t_eval the forward pass holds a few copies of the state plus the (len(t_eval),) + y.shape output, independent of max_steps. Without t_eval the output itself has max_steps + 1 rows, and the forward pass holds about two such arrays; prefer t_eval for large states. Reverse mode with checkpoint=False stores the step residuals (several copies of the state) for each of the max_steps iterations, whatever the number of steps actually taken; e.g. about 0.9 GiB for a 1e4-element float32 state at max_steps=4096, against about 7 MiB with checkpoint=True.

Backwards integration (t1 < t0) is supported: steps are taken with negative dt.

Transformations: jax.jit (method, max_steps and the length of t_eval are static; t_span, y0, tolerances and arguments may be traced), jax.vmap and forward/reverse-mode jax.grad with respect to y0, args, kwargs, t_span and t_eval.

References

Examples

>>> import jax.numpy as jnp
>>> from braintools.quad import solve_ivp
>>> sol = solve_ivp(lambda y, t: -y, (0.0, 2.0), jnp.array(1.0),
...                 rtol=1e-6, atol=1e-8,
...                 t_eval=jnp.array([0.0, 1.0, 2.0]))
>>> bool(sol.success)
True
>>> print([round(float(v), 5) for v in sol.ys])
[1.0, 0.36788, 0.13534]

Without t_eval the accepted steps are returned:

>>> sol = solve_ivp(lambda y, t, k: -k * y, (0.0, 2.0), jnp.array(1.0), 3.0,
...                 max_steps=64)
>>> n = int(sol.n_steps)
>>> print(float(sol.ts[n]), round(float(sol.ys[n]), 6))
2.0 0.002479
Contents