Integration Models

Integration models perform numerical integration of derivative models while supporting automatic differentiation through time.

dynax.ODESolver

ODE solver module.

Solver for continuous-time ODE systems of the form

\[ \dot{\boldsymbol{x}} = \boldsymbol{f}\bigl(t, \boldsymbol{x}, \boldsymbol{u}; \text{args}\bigr) \]

with time \(t \in \mathbb{R}\), state \(\boldsymbol{x}(t) \in \mathbb{R}^n\), input \(\boldsymbol{u}(t) \in \mathbb{R}^m\) and optional static parameters \(\text{args}\).

The system is integrated using diffrax.diffeqsolve. To evaluate the input \(\boldsymbol{u}(t)\) at arbitrary points in time Hermite cubic splines with backward differences are used to interpolate the discrete input signal \(\boldsymbol{u}_i\) given at times \(t_i\).

Notation

n: State dimension. m: Input dimension. l: Number of time steps in the simulation horizon.

__init__(func, *, solver=Tsit5(), stepsize_controller=PIDController(rtol=1e-06, atol=1e-06), max_steps=4096, dt0=None, adjoint=RecursiveCheckpointAdjoint())

Initialize the ODE solver.

PARAMETER	DESCRIPTION
func	Dynamics function (right hand side of the ODE) f(t, x, u, args) -> dx/dt. TYPE: ODEFunc
solver	Diffrax ODE solver used by diffrax.diffeqsolve. TYPE: AbstractSolver DEFAULT: Tsit5()
stepsize_controller	Diffrax step-size controller. TYPE: AbstractStepSizeController DEFAULT: PIDController(rtol=1e-06, atol=1e-06)
max_steps	Maximum number of integration steps accepted by Diffrax. If None, Diffrax uses its default behavior. TYPE: int | None DEFAULT: 4096
dt0	The step size to use for the first step. If using fixed step sizes then this will also be the step size for all other steps. If set as None then the initial step size will be determined automatically. TYPE: Shaped[Array, ''] | None DEFAULT: None
adjoint	How to differentiate diffrax.diffeqsolve. Defaults to discretize-then-optimize, which is usually the best option for most problems. See the diffrax page on Adjoints for more information. TYPE: AbstractAdjoint DEFAULT: RecursiveCheckpointAdjoint()

__call__(ts, x0, us=None, args=None)

Solve the ODE.

PARAMETER	DESCRIPTION
ts	Monotonic simulation time grid. TYPE: Array
x0	Initial state at ts[0]. TYPE: PyTree
us	Input trajectory aligned with ts. If provided, cubic interpolation is used to evaluate \(u(t)\) during integration. Use None for unforced systems. TYPE: PyTree | None DEFAULT: None
args	Additional arguments passed to func. TYPE: PyTree DEFAULT: None

RETURNS	DESCRIPTION
PyTree	State trajectory sampled at ts.

get_diffrax_solution(ts, x0, us=None, args=None)

Integrate the dynamics and return the raw Diffrax solution object.

PARAMETER	DESCRIPTION
ts	Monotonic simulation time grid. TYPE: Array
x0	Initial state at ts[0]. TYPE: PyTree
us	Input trajectory aligned with ts. If provided, cubic interpolation is used to evaluate \(u(t)\) during integration. Use None for unforced systems. TYPE: PyTree | None DEFAULT: None
args	Additional arguments passed to func. TYPE: PyTree DEFAULT: None

RETURNS	DESCRIPTION
Solution	Diffrax integration result (diffrax.Solution) with states stored
Solution	at sampling times ts.

---

For convenience, some light wrappers around the ODESolver are provided:

dynax.AugmentedODE

Augmented ODE module.

Extends a given initial condition with additional augmented states, then passes it to an ODESolver. Finally, the augmentation is removed from the solution before returning.

\[ \begin{align} \boldsymbol{z}_0 &= \begin{bmatrix} \boldsymbol{x}_0 \\ \tilde{\boldsymbol{x}}_0 \end{bmatrix}, \\ \dot{\boldsymbol{z}} &= \boldsymbol{f}\bigl(t, \boldsymbol{z}, \boldsymbol{u}; \text{args}\bigr), \\ \boldsymbol{x}(t) &= \boldsymbol{z}(t)[:n] \end{align} \]

with time \(t \in \mathbb{R}\), state \(\boldsymbol{x}(t) \in \mathbb{R}^n\), augmented state \(\boldsymbol{z}(t) \in \mathbb{R}^{n+n_\text{aug}}\), input \(\boldsymbol{u}(t) \in \mathbb{R}^m\) and optional static parameters \(\text{args}\).

__init__(ode_solver, x0_auxiliary, concat_axes=0)

Initialize the ODE solver.

PARAMETER	DESCRIPTION
ode_solver	ODESolver for the augmented ODE f(t, z, u, args) -> dz/dt. TYPE: ODESolver
x0_auxiliary	The initial auxiliary state, that is concatenated to x0 to compute the initial augmented state. Should be a prefix of x0. To ensure that the auxiliary initial state is not updated during training, wrap it in klax.non_trainable. Note Typically, your state is a 1-d array. Then x0_auxiliary is also simply a 1-d array, containing the entries that will be concatenated to x0. TYPE: PyTree
concat_axes	The axes along which to concatenate the state and auxiliary state. Should be a prefix of x0_auxiliary. Note TYPE: PyTree[int] DEFAULT: 0

__call__(ts, x0, us=None, args=None)

Solve the augmented ODE.

PARAMETER	DESCRIPTION
ts	Monotonic simulation time grid. TYPE: Array
x0	Initial (non-augmented) state at ts[0]. TYPE: PyTree
us	Input trajectory aligned with ts, or None for unforced systems. TYPE: PyTree | None DEFAULT: None
args	Additional arguments passed to ODESolver.func TYPE: PyTree DEFAULT: None

RETURNS	DESCRIPTION
PyTree	State trajectory sampled at ts with the auxiliary states removed.

dynax.StateSpaceSystem

Module for general state-space systems.

Module for continuous-time state-space systems of the form

\[ \begin{align} \dot{\boldsymbol{x}} &= \boldsymbol{f}\bigl(t, \boldsymbol{x}, \boldsymbol{u}; \text{args}\bigr), \\ \boldsymbol{y} &= \boldsymbol{g}\bigl(t, \boldsymbol{x}, \boldsymbol{u}; \text{args}\bigr), \end{align} \]

with time \(t \in \mathbb{R}\), state \(\boldsymbol{x}(t) \in \mathbb{R}^n\), input \(\boldsymbol{u}(t) \in \mathbb{R}^m\), output \(\boldsymbol{y}(t) \in \mathbb{R}^p\) and optional static parameters \(\text{args}\).

The state equation is integrated using the ODESolver module.

__init__(ode_solver, output_equation)

Initialize a state-space simulation model.

PARAMETER	DESCRIPTION
ode_solver	ODESolver for the state equation. TYPE: ODESolver
output_equation	Output function g(t, x, u, args) -> y. TYPE: Callable[[Real[ArrayLike, ''], PyTree, PyTree, PyTree], PyTree]

__call__(ts, x0, us=None, args=None)

Simulate the system and return output trajectories.

PARAMETER	DESCRIPTION
ts	Monotonic simulation time grid. TYPE: Array
x0	Initial state at ts[0]. TYPE: PyTree
us	Input trajectory aligned with ts, or None for unforced systems. Each leaf is expected to have a leading time axis of length len(ts). Passed sample-wise to output_equation (i.e. u(ts[i]), not interpolated). TYPE: PyTree | None DEFAULT: None
args	Additional arguments passed to output_equation. TYPE: PyTree DEFAULT: None

RETURNS	DESCRIPTION
PyTree	Output trajectory sampled at ts.