register_method

Contents

register_method#

class braintools.quad.register_method(name, fn, kind='ode', *, overwrite=False)#

Register a one-step integrator under a short name.

After registration the stepper can be selected by name in the drivers, e.g. odeint(f, y0, t0, t1, dt, method=name), and looked up with get_method(). The stepper must follow the calling convention of the built-in steppers of its kind (for 'ode': fn(f, y, t, *args, **kwargs) -> y_next, reading the step size from brainstate.environ.get_dt()).

Parameters:

  • name (str) – Short name of the method, e.g. 'my_rk'. Must be a non-empty string.

  • fn (Callable) – The one-step integrator.

  • kind (str) – Equation kind the stepper solves. Default 'ode'.

  • overwrite (bool) – Replace an existing entry with the same name. Default False.

Returns:

fn unchanged, so the call can be used inline.

Return type:

Callable

Raises:

  • TypeError – If name is not a non-empty string or fn is not callable.

  • ValueError – If kind is unknown, or name is already registered for kind and overwrite is False.

See also

get_method

Look up a registered stepper.

list_methods

List the registered names.

Notes

The registry is a process-wide dictionary; registration is a host-side Python operation and has nothing to do with jit/vmap/grad.

Examples

>>> import brainstate
>>> from braintools.quad import register_method, get_method
>>> def forward_euler(f, y, t, *args, **kwargs):
...     return y + brainstate.environ.get_dt() * f(y, t, *args, **kwargs)
>>> # ``overwrite=True`` lets the example be re-run in the same process.
>>> _ = register_method('my_euler', forward_euler, overwrite=True)
>>> get_method('my_euler') is forward_euler
True
Contents