Derivative Models

Derivative models parameterize the right hand side of an ODE in normal form: $$ \dot{\mathbf{y}} = \mathbf{f}(t, \mathbf{y}[, \mathbf{u}, \text{args}]) $$ where \(\mathbf{y}(t)\in\mathbb{R}^n\) is the systems state vector, \(t\) is time, \(\mathbf{u}(t)\in\mathbb{R}^m\) are external inputs (for example forces acting on a system), and \(\text{args}\) are arbitrary parameters of the system.

They are most often combined with a numerical solver.

Port-Hamiltonian systems

dynax.ISPHS

Input-State port-Hamiltonian Systems (ISPHS).

Implementation of the state equation of a input-state-output port-Hamiltonian system: $$ \dot{\mathbf{x}} = (\mathbf{J}(\mathbf{x})-\mathbf{R}(\mathbf{x}))\frac{\partial\mathcal{H}(\mathbf{x})}{\partial \mathbf{x}} + \mathbf{B}(\mathbf{x})\mathbf{u} $$ where: \(\mathbf{x}(t)\) is the state vector, \(\mathbf{u}(t)\) is the input vector, \(\mathcal{H}(\mathbf{x})\) is the Hamiltonian function given by hamiltonian, \(\mathbf{J}\) is the structure matrix given by structure_matrix, \(\mathbf{R}\) is the dissipation matrix given by dissipation_matrix, and \(\mathbf{B}\) is the input matrix given by input_matrix.

__init__(hamiltonian, structure_matrix, dissipation_matrix=None, input_matrix=None)

Initialize the ISPHS.

PARAMETER	DESCRIPTION
hamiltonian	Function or submodel computing the Hamiltonian \(\mathcal{H}(x)\) as a function of the state vector. TYPE: Callable[[Array], Shaped[Array, '']]
structure_matrix	Function or submodel computing the structure matrix \(J(x)\) as a function of the state vector. From the port-Hamiltonian modeling view the structure matrix is required to be skew-symmetric, i.e., \(J = -J^T\). TYPE: Callable[[Float[Array, 'n']], Float[Array, 'n n']]
dissipation_matrix	Function or submodel computing the dissipation matrix \(R(x)\) as a function of the state vector. From the port-Hamiltonian modeling view the dissipation matrix is required to be symmetric positive semi-definite, i.e., \(R = R^T,\,R\succ0\). Alternatively, None can be used to indicate that the system does not have a dissipation matrix \(R\). TYPE: Callable[[Float[Array, 'n']], Float[Array, 'n n']] | None DEFAULT: None
input_matrix	Function or submodel computing the input matrix \(G(x)\) as a function of the state vector. Alternatively, None can be used to indicate that the system does not have external inputs \(u\). Defaults to None. TYPE: Callable[[Float[Array, 'n']], Float[Array, 'n m']] | None DEFAULT: None

Tipp

Consider using klax for convenient implementations of matrix valued functions.

__call__(t, x, u=None)

Return the time derivative of the state vector \(x\).

PARAMETER	DESCRIPTION
t	Scalar time for evaluation. This argument is unused, but required to comply with the ODESolver function signature. TYPE: Shaped[Array, '']
x	State vector at time t. This should be a 1D array of shape (n,) where n is the number of state dimensions. TYPE: Array
u	Input vector at time t. This should be a 1D array of shape (m,) where m is the number of input dimensions. If the system does not have inputs, this can be set to None. TYPE: Array | None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If a input is provided but the system does not have an input matrix.

RETURNS	DESCRIPTION
Array	Time derivative of the state vector \(x\) as a 1D array of shape (n,).

Physics-agnostic systems

dynax.NeuralODE

Derivative model for a neural ODE.

This implements the following equation: $$ \dot{y} = f(t, y, u) $$ where \(f\) is a multi-layer perceptron (MLP), depending on time \(t\in\mathbb{R}\), the state vector \(y(t)\in\mathbb{R}^n\), and an input vector \(u(t)\in\mathbb{R}^m\).

__init__(state_size, input_size, *, time_dependent=False, state_dependent=True, width_sizes, weight_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, activation=<PjitFunction of <function softplus at 0x7fdf0a78fc40>>, final_activation=<function NeuralODE.<lambda>>, use_bias=True, use_final_bias=True, weight_wrap=None, bias_wrap=None, dtype=None, key)

Create a neural ode.

PARAMETER	DESCRIPTION
state_size	Number of state variables. TYPE: int
input_size	Number of input variables. Can be zero to indicate no input dependence. TYPE: int
time_dependent	If true then the time is fed as an additional array entry to the MLP. Defaults to False. TYPE: bool DEFAULT: False
state_dependent	If true then the state is fed as an additional array entry to the MLP. Defaults to True. TYPE: bool DEFAULT: True
width_sizes	The sizes of each hidden layer in a list. TYPE: Sequence[int]
weight_init	The weight initializer of type jax.nn.initializers.Initializer. (Defaults to he_normal().) TYPE: jax.nn.initializers.Initializer DEFAULT: <function variance_scaling.<locals>.init>
bias_init	The bias initializer of type jax.nn.initializers.Initializer. (Defaults to zeros.) TYPE: jax.nn.initializers.Initializer DEFAULT: <function zeros>
activation	The activation function after each hidden layer. (Defaults to jax.nn.softplus). TYPE: Callable DEFAULT: <PjitFunction of <function softplus at 0x7fdf0a78fc40>>
final_activation	The activation function after the output layer. (Defaults to the identity.) TYPE: Callable DEFAULT: <function NeuralODE.<lambda>>
use_bias	Whether to add on a bias to internal layers. (Defaults to True.) TYPE: bool DEFAULT: True
use_final_bias	Whether to add on a bias to the final layer. (Defaults to True.) TYPE: bool DEFAULT: True
weight_wrap	An optional wrapper that is passed to all weights. TYPE: type[klax._wrappers.Constraint] | type[klax._wrappers.Unwrappable[Array]] | None DEFAULT: None
bias_wrap	An optional wrapper that is passed to all biases. TYPE: type[klax._wrappers.Constraint] | type[klax._wrappers.Unwrappable[Array]] | None DEFAULT: None
dtype	The dtype to use for all the weights and biases in this MLP. Defaults to either jax.numpy.float32 or jax.numpy.float64 depending on whether JAX is in 64-bit mode. TYPE: type | None DEFAULT: None
key	A jax.random.PRNGKey used to provide randomness for parameter initialisation. TYPE: PRNGKeyArray

__call__(t, y, u)

Evaluate the neural ODE's derivative.