odeint

Contents

odeint#

class braintools.quad.odeint(f, y0, t0, t1=None, dt=None, *args, method='rk4', save_every=1, checkpoint=False, n_steps=None, **kwargs)#

Integrate an ODE with a fixed step and return the trajectory.

Runs n steps of a one-step integrator (see braintools.quad.get_method()) as one compiled scan inside brainstate.environ.context(dt=dt), saving every save_every-th state. Use it instead of a Python for loop over ode_*_step calls: the step function is traced once and the whole loop compiles into a single XLA program.

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 of the braintools.quad steppers.

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

  • t0 (float or Quantity) – Initial time.

  • t1 (float or Quantity, optional) – Final time. The number of steps is n = round((t1 - t0) / dt), computed on the host, so t0, t1 and dt must be concrete (not traced). Mutually exclusive with n_steps.

  • dt (float or Quantity) – Step size (required). May be negative to integrate backwards (then t1 < t0).

  • *args – Extra positional arguments forwarded to f.

  • method (str | Callable) – Name of a registered 'ode' stepper ('euler', 'rk2', 'rk4', 'dopri5', …; see braintools.quad.list_methods()) or a stepper callable with the step(f, y, t, *args, **kwargs) convention. Default 'rk4'.

  • save_every (int) – Save one state every save_every steps. Default 1.

  • checkpoint (bool) – If True, use brainstate.transform.checkpointed_scan so that reverse-mode gradients through long rollouts rematerialise the intermediate states instead of storing them all. The values and gradients are identical to checkpoint=False. Default False.

  • n_steps (int | None) – Number of steps, given instead of t1. Use this under jax.jit when t0 or dt are traced. Must be a Python int >= 0.

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

Returns:

  • ts (array or Quantity) – Saved times, shape (n // save_every + 1,), starting at t0: ts[k] = t0 + k * save_every * dt.

  • ys (PyTree) – Saved states with the structure of y0; every leaf gains a leading axis of length n // save_every + 1 and ys[0] == y0.

Raises:

  • ValueError – If dt is missing or zero, if neither or both of t1 and n_steps are given, if the step count is negative (t1 lies behind t0 in the direction of dt), if n_steps is not a non-negative int, if save_every < 1, or if t1 is given but t0, t1 or dt is traced (use n_steps under jit).

  • brainunit.UnitMismatchError – If (t1 - t0) / dt is not dimensionless.

  • KeyError – If method is an unknown name.

See also

solve_ivp

Adaptive-step solver with error control.

get_method

Resolve a stepper by name.

Notes

The i-th step is taken at time t0 + i * dt (computed from the step index, so round-off does not accumulate in t).

If save_every does not divide n, the trajectory stops at the last saved state: only (n // save_every) * save_every steps are run, and the last saved time is t0 + (n // save_every) * save_every * dt, which is earlier than t1. Choose save_every dividing n to reach t1. t1 == t0 (or n < save_every) returns ts == [t0] and ys holding only y0.

Since n is rounded, the last time t0 + n * dt can differ from t1 by up to dt / 2.

Transformations: odeint is jax.jit-compatible (with n_steps instead of t1 when times are traced; method, save_every, checkpoint and n_steps are static), jax.vmap-compatible over y0, args and kwargs, and differentiable in forward and reverse mode with respect to y0, args, kwargs, t0 and dt.

Examples

>>> import jax.numpy as jnp
>>> from braintools.quad import odeint
>>> ts, ys = odeint(lambda y, t: -y, jnp.array(1.0), 0.0, 1.0, 0.1)
>>> ts.shape, ys.shape
((11,), (11,))
>>> print(f'{float(ys[-1]):.5f}  exact {float(jnp.exp(-1.0)):.5f}')
0.36788  exact 0.36788
>>>
>>> ts, ys = odeint(lambda y, t: -y, jnp.array(1.0), 0.0, 1.0, 0.1,
...                 method='euler', save_every=5)
>>> ts.shape
(3,)

Pytree and unit-aware states:

>>> import brainunit as u
>>> def neuron(s, t, I):
...     return {'V': (-65. * u.mV - s['V'] + 10. * u.Mohm * I) / (10. * u.ms)}
>>> ts, ys = odeint(neuron, {'V': -65. * u.mV}, 0. * u.ms, 50. * u.ms,
...                 0.1 * u.ms, 0.5 * u.nA)
>>> ys['V'].shape
(501,)
>>> print(u.math.round(ys['V'][-1], 2))
-60.03 mV
Contents