Derivative Models

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

They are most often combined with a numerical solver.

Port-Hamiltonian systems

dynax.ISOPHS

Input-State-Output port-Hamiltonian Systems (ISOPHS).

Implementation of an input-state-output port-Hamiltonian system:

\[ \begin{align} \dot{\boldsymbol{x}} &= (\boldsymbol{J}-\boldsymbol{R})\frac{\partial\mathcal{H}}{\partial \boldsymbol{x}} + \boldsymbol{G}\boldsymbol{u}, \\ \bigl[\boldsymbol{y} &= \boldsymbol{G}^\intercal \frac{\partial\mathcal{H}}{\partial \boldsymbol{x}}\bigr] \end{align} \]

where

• \(\boldsymbol{x}(t)\) is the state,
• \(\boldsymbol{u}(t)\) is the input,
• \(\boldsymbol{y}(t)\) is the output,
• \(\mathcal{H}(\boldsymbol{x}, \text{args})\) is the Hamiltonian function given by hamiltonian,
• \(\boldsymbol{J}(\boldsymbol{x}, \text{args})\) is the structure matrix given by structure_matrix,
• \(\boldsymbol{R}(\boldsymbol{x}, \text{args})\) is the dissipation matrix given by dissipation_matrix,
• \(\boldsymbol{G}(\boldsymbol{x}, \text{args})\) is the input matrix given by input_matrix,
• and \(\text{args}\) is an optional set of parameters.

Note

The usage of the output equation is optional, but requires that the input matrix (input_matrix) is defined.

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

Initialize the ISOPHS.

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

__call__(t, x, u=None, args=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: Real[ArrayLike, '']
x	State vector at time t. This should be a 1D array of shape (n,) where n is the number of state dimensions. TYPE: Float[Array, n]
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: Float[Array, m] | None DEFAULT: None
args	Additional arguments passed to \(\mathcal{H}, \boldsymbol{J}, \boldsymbol{R}\) and \(\boldsymbol{G}\). TYPE: PyTree DEFAULT: None

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

RETURNS	DESCRIPTION
Float[Array, n]	Time derivative of the state vector \(\boldsymbol{x}\) as a 1D array of shape (n,).

dynax.MatrixWrapper

Compatibility layer for using klax matrices with the call interface required by ISOPHS.

Matrix-valued functions implemented in klax currently only support one input argument x -> A. However, the matrix-valued functions required for the input-state-output port-Hamiltonian system demand the call signature (x, args) -> A. This class can wrap the klax matrix-valued functions to comply with the latter call signature. It supports 4 modes depending on state_dependent and parameter_dependent:

• If both state_dependent=parameter_dependent=True, then the x and args are combined using jax.flatten_util.ravel_pytree([x, args]) and the result is passed to the wrapped matrix-valued function.
• If only state_dependent=True, then x is passed to the matrix-valued function.
• If only parameter_dependent=True, then args is passed to the matrix-valued function.
• If both state_dependent=parameter_dependent=False, then None is passed to the matrix-valued function.

__init__(matrix, state_dependent=False, parameter_dependent=False)

Initialize self. See help(type(self)) for accurate signature.

__call__(x, args)

Physics-agnostic systems

dynax.NeuralODE

Derivative model for a neural ODE.

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

__init__(state_size, input_size=None, parameter_size=None, *, time_dependent=False, state_dependent=True, width_sizes, weight_init=init, bias_init=zeros, activation=softplus, final_activation=<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 elements in the state vector x. TYPE: int
input_size	Number of elements in the input vector u. Can be None to indicate no input dependence. TYPE: int | None DEFAULT: None
parameter_size	Number of elements in the raveled parameter PyTree args; Essentially, the number of scalar parameters. Can be None to indicate no parameter dependence. TYPE: int | None DEFAULT: None
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: Initializer DEFAULT: init
bias_init	The bias initializer of type jax.nn.initializers.Initializer. (Defaults to zeros.) TYPE: Initializer DEFAULT: zeros
activation	The activation function after each hidden layer. (Defaults to jax.nn.softplus). TYPE: Callable DEFAULT: softplus
final_activation	The activation function after the output layer. (Defaults to the identity.) TYPE: Callable DEFAULT: <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[Constraint] | type[Unwrappable[Array]] | None DEFAULT: None
bias_wrap	An optional wrapper that is passed to all biases. TYPE: type[Constraint] | type[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 initialization. TYPE: PRNGKeyArray

__call__(t, x, u=None, args=None)

Evaluate the neural ODE's derivative.