Latent space projections

Linear projections

dynax.PODLatentSpace

Proper orthogonal decomposition based latent space.

Linear projections to and from a latent space using the right singular vectors from the SVD of snapshots with as modes/basis.

When projecting into the latent space, the input is first shifted (optional), then scaled and finally projected onto the modes. By default, the scaling is computed such that the latent variables of the snapshots used to calculate the modes have standard deviation one.

For the projection into the full order space, the latent variables are first multiplied with the modes, then scaled and finally shifted.

__init__(snapshots, num_modes, shift=None, scaling='normalize')

Initialize the SVD latent space.

Note

After initialization you can still change the number of modes, i.e., the latent space dimension.

PARAMETER	DESCRIPTION
snapshots	Snapshots of shape (n_batch, n_time, n_space). TYPE: Union[Shaped[ndarray, 'n_batch n_time n_space'], Shaped[Array, 'n_batch n_time n_space']]
num_modes	Number of modes to use in the latent space. TYPE: int
shift	Array or float to shift a snapshot by, before transforming. By setting shift=<special state>, the special state is transformed to the origin of the latent space. TYPE: Union[float, Shaped[ndarray, '#n_space'], Shaped[Array, '#n_space'], NoneType] DEFAULT: None
scaling	Scalar value to scale all latent variables by. Can be "normalize", in which case the scale is automatically chosen, such that the latent variables of snapshots have standard deviation one. TYPE: Union[float, Shaped[ndarray, ''], Shaped[Array, ''], Literal['normalize'], NoneType] DEFAULT: 'normalize'

to_latent(snapshot)

Project snapshot into the latent space.

PARAMETER	DESCRIPTION
snapshot	Data with shape (..., n_space). TYPE: Union[Shaped[ndarray, '... n_space'], Shaped[Array, '... n_space']]

RETURNS	DESCRIPTION
Shaped[Array, '... n_modes']	Array of latent variables with shape (..., n_modes).

from_latent(latent_variables)

Project latent variables into the full order space.

PARAMETER	DESCRIPTION
latent_variables	Data with shape (..., n_modes). TYPE: Union[Shaped[ndarray, '... n_modes'], Shaped[Array, '... n_modes']]

RETURNS	DESCRIPTION
Shaped[Array, '... n_space']	Array with shape (..., n_space).

Non-linear projections

dynax.CALMLayer

Continuous and Adaptive convolution Layer.

Implementation of CALM-PDE by Jan Hagnberger, Daniel Musekamp and Mathias Niepert.

__init__(in_channels, out_channels, *, initial_query_points, receptive_size, temperature=1.0, query_points_learnable=True, modulate_kernel=True, initial_input_points=None, num_candidates=None, weight_init=init, bias_init=zeros, rff_init=init, activation=gelu, kernel_activation=gelu, key)

Initialize CALM layer.

PARAMETER	DESCRIPTION
in_channels	Number of input channels. TYPE: int
out_channels	Number of output channels. TYPE: int
initial_query_points	Initial query points array with shape (num_query_points, num_spatial_dims). Alternatively, it can be a callable that takes a PRNG key and returns such an array. TYPE: Float[Array, 'num_query_points num_spatial_dims'] | Callable[[PRNGKeyArray], Float[Array, 'num_query_points num_spatial_dims']]
receptive_size	Number of points in the input to consider for each query point's receptive field. This is similar to the kernel size in discrete convolutions. TYPE: int
temperature	Temperature parameter for the softmax function computing the distance weighting. Smaller values increase the weighting of contributions of closer input points for the Monte Carlo integration, while larger values tend to treat all input points equally. Defaults to 1. TYPE: float DEFAULT: 1.0
query_points_learnable	Whether query points are learnable. Defaults to True. TYPE: bool DEFAULT: True
modulate_kernel	Whether to use kernel modulation. Kernel modulation enables the use of different kernel functions per query point, by scaling and shifting the intermediate kernel MLP features with query point specific parameters. Defaults to True. TYPE: bool DEFAULT: True
initial_input_points	Optional initial input points array with shape (num_input_points, num_spatial_dims). If provided, this is used to determine the indices of candidate input points per query point for efficient nearest neighbour computation. Instead of the inefficient brute-force nearest neighbour search over all input points, only the candidate input points are considered, when computing the receptive_size nearest neighbours for each query point. Candidates are the num_candidates closest input points to each query point. WARNING: The candidate indices are computed once at initialization and remain fixed during training. Therefore, if the input OR query points change significantly (This includes permutations of their arrays!) during training or inference, the candidate indices may no longer be valid, leading to incorrect results. Defaults to None. TYPE: Float[Array, 'num_input_points num_spatial_dims'] | None DEFAULT: None
num_candidates	Number of input point candidates saved for the efficient nearest neighbour computation. If None, the receptive_size is used. Defaults to None. TYPE: int | None DEFAULT: None
weight_init	Initializer for weights. Defaults to glorot_normal(). TYPE: Initializer DEFAULT: init
bias_init	Initializer for biases. Defaults to zeros. TYPE: Initializer DEFAULT: zeros
rff_init	Initializer for the random fourier feature (RFF) weight. Defaults to normal(6.0). TYPE: Initializer DEFAULT: init
activation	Activation function for the layer output. Defaults to jax.nn.gelu. TYPE: Callable[[Array], Array] DEFAULT: gelu
kernel_activation	Activation function for the kernel MLP. Defaults to jax.nn.gelu. TYPE: Callable[[Array], Array] DEFAULT: gelu
key	PRNG key for initialization. TYPE: PRNGKeyArray

__call__(x, points, *, key=None)

Compute the continuous convolution output.

PARAMETER	DESCRIPTION
x	Input features with shape (..., num_input_points, in_channels). TYPE: Float[Array, '... n c']
points	Input point coordinates with shape (num_input_points, num_spatial_dims). TYPE: Float[Array, 'n d']
key	Unused random key. Defaults to None. TYPE: PRNGKeyArray | None DEFAULT: None

RETURNS	DESCRIPTION
Float[Array, '... m o']	Output features for the query point with shape (..., num_query_points, out_channels).
Float[Array, 'm d']	Query point coordinates with shape (num_query_points, num_spatial_dims).

compute_kernel(rel_pos)

Compute the kernel matrix for a single query point.

PARAMETER	DESCRIPTION
rel_pos	Relative positions with shape (num_query_points, num_neighbours, num_spatial_dims). TYPE: Float[Array, 'm k d']

RETURNS	DESCRIPTION
Float[Array, 'm k c o']	Kernel matrix with shape (num_query_points, num_neighbours, in_channels, out_channels)

Helper classes for building CALM-based autoencoders

dynax.PointArgSequential

A sequence of equinox Modules adhering to the PointArgFunc protocol applied in sequence.

This is a modified version of equinox' Sequential.

__init__(layers)

Initialize the PointArgSequential.

PARAMETER	DESCRIPTION
layers	Sequence of PointArgFuncs. TYPE: Sequence[PointArgFunc]

__call__(features, points, /, *, key=None)

Compute the sequential models output.

PARAMETER	DESCRIPTION
features	Features passed to the first layer. TYPE: Array
points	Points passed to the first layer. TYPE: Array
key	If provided, then it is split by the number of layers, and a subkey passed to each layer. TYPE: PRNGKeyArray | None DEFAULT: None

RETURNS	DESCRIPTION
tuple[Array, Array]	Output of the last layer, which is a tuple of features and points.

dynax.PointArgAE

Autoencoder with PointArgFunc encoder and decoder.

__init__(encoder, decoder)

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

encode(features, points, /, *, key=None)

decode(features, points, /, *, key=None)

dynax.PointArgWrapper

Wraps a function from features to features to be compatible with the PointArgFunc protocol.

__init__(func)

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

__call__(features, points, /, *, key=None)

dynax.PointArgFunc

Function that maps arrays of features and points to features and points.

__call__(features, points, /, *, key=None)