torchdiffeq

You can trigger callbacks during a solve by defining specific methods on the func object passed to odeint or odeint_adjoint.

Supported Forward Callbacks:

• callback_step(self, t0, y0, dt): Called immediately before taking a step. Supported by all solvers except scipy_solver.
• callback_accept_step(self, t0, y0, dt): Called when a step is accepted. Supported by adaptive solvers (dopri8, dopri5, bosh3, adaptive_heun).
• callback_reject_step(self, t0, y0, dt): Called when a step is rejected. Supported by adaptive solvers.

Adjoint Callbacks: To trigger these during the adjoint pass, append _adjoint to the method name (e.g., callback_step_adjoint).

Use odeint_event to terminate an ODE solution when an event function triggers (i.e., when an element of event_fn(t, y) equals zero).

Arguments:

• func: The ODE callable.
• y0: Initial values.
• t0: Scalar representing the initial time value.
• event_fn: A required keyword argument event_fn(t, y) that returns a tensor. The solve terminates when any element in the returned tensor is zero.
• reverse_time: Boolean (default False) to solve in reverse time.
• odeint_interface: Specifies whether to use odeint or odeint_adjoint for differentiation (default is odeint).
• **kwargs: Passed to the chosen odeint_interface.

Note: To obtain gradients for the event function, its parameters must be contained within the state y itself.

The odeint function is the main interface for solving initial value problems (IVP) of the form dy/dt = f(t, y) with initial condition y(t_0) = y_0.

Arguments:

• func: A callable implementing the ODE f(t, x).
• y0: An any-D Tensor representing the initial values.
• t: A 1-D Tensor containing the evaluation points. The initial time is taken to be t[0].

For backpropagation through ODE solutions with constant memory cost, use odeint_adjoint. It uses $O(1)$ memory by solving an adjoint ODE during the backward pass.

Critical Requirement: When using the adjoint method, func must be a torch.nn.Module so that the solver can collect the parameters of the differential equation.

When using odeint_adjoint, you can specify parameters for the backward pass:

• adjoint_rtol, adjoint_atol, adjoint_method, adjoint_options: Controls for the backward pass. Defaults to the values used in the forward pass.
• adjoint_options: Can include {"norm": "seminorm"} to provide a more efficient adjoint solve when using adaptive step solvers.
• adjoint_params: A tuple of tensors representing the parameters to compute gradients with respect to. Defaults to tuple(func.parameters()). If func has no parameters, you must specify adjoint_params=().

Fixed solvers (euler, midpoint, rk4, explicit_adams, implicit_adams) use the following options:

• step_size: The size of each discrete step. If not provided, the solver steps between values in t. Note: step_size is mutually exclusive with grid_constructor.
• grid_constructor: A callable func(func, y0, t) -> grid that returns a 1D tensor representing the desired step locations.
• perturb: If True, adds small perturbations to the start and end of each step to assist with stepping to discontinuities.

You can pass several keyword arguments to odeint or odeint_adjoint to control the solver behavior:

• rtol: Relative tolerance.
• atol: Absolute tolerance.
• method: The solver algorithm to use.
• options: A dictionary of solver-specific options.

Available Solvers

Adaptive-step solvers:

• dopri5 (Default): Runge-Kutta of order 5 (Dormand-Prince-Shampine).
• dopri8: Runge-Kutta of order 8 (Dormand-Prince-Shampine).
• bosh3: Runge-Kutta of order 3 (Bogacki-Shampine).
• fehlberg2: Runge-Kutta-Fehlberg of order 2.
• adaptive_heun: Runge-Kutta of order 2.

Fixed-step solvers:

• euler: Euler method.
• midpoint: Midpoint method.
• rk4: Fourth-order Runge-Kutta with 3/8 rule.
• explicit_adams: Explicit Adams-Bashforth.
• implicit_adams: Implicit Adams-Bashforth-Moulton.

SciPy solvers:

• scipy_solver: Wraps all solvers available through SciPy.

The Adams solvers have specialized parameters:

explicit_adams

• max_order: The maximum order of the Adams-Bashforth predictor. Note that rtol and atol are ignored for this solver.

implicit_adams

• max_order: The maximum order of the Adams-Bashforth-Moulton predictor-corrector.
• max_iters: Maximum number of iterations for the Adams-Moulton corrector.
• rtol and atol: Correspond to the tolerance for convergence of the Adams-Moulton corrector.