scimba_torch.flows.deep_flows

Discrete and continuous flow-based approximation spaces.

It includes functionality for evaluating the networks, setting/retrieving degrees of freedom, and computing Jacobians.

Functions

identity(input)

Identity function that returns the input as is.

Classes

ContinuousFlowSpace(size, param_dim, **kwargs)

A continuous-time flow model that can be used for generating flows.

DiscreteFlowSpace(nb_unknowns, ...)

A nonlinear approximation space using a neural network model.
identity(input)[source]

Identity function that returns the input as is.

Parameters:

input (Any) – Any input.

Return type:

Any

Returns:

The input unchanged.

class DiscreteFlowSpace(nb_unknowns, nb_parameters, flow_type, **kwargs)[source]

Bases: AbstractApproxSpace, Module

A nonlinear approximation space using a neural network model.

This class represents a parametric approximation space, where the solution is modeled by a neural network. It integrates functionality for evaluating the network, setting/retrieving degrees of freedom, and computing the Jacobian.

Parameters:

  • nb_unknowns (int) – Number of unknowns in the approximation problem.

  • nb_parameters (int) – Number of parameters in the input.

  • flow_type (Module) – The neural network class used to approximate the solution.

  • **kwargs – Additional arguments passed to the neural network model.

in_size: int

Size of the input to the neural network (spatial dimension + parameters).

param_dim: int

Number of parameters in the input.

out_size: int

Size of the output of the neural network (number of unknowns).

rollout: int

Number of rollout steps.

post_processing: Callable

Post-processing function applied to outputs.

apply_at_each_step: Callable

Function applied at each step during inference.

network: Module

Neural network used for the approximation.

ndof: int

Total number of degrees of freedom in the network.

type_space: str

Type of the approximation space.

forward(x, mu, with_last_layer=True)[source]

Evaluate the parametric model for given input features.

Parameters:

  • x (Tensor) – Input tensor representing the state.

  • mu (Tensor) – Input tensor representing the parameters.

  • with_last_layer (bool) – Whether to include the last layer in evaluation.

Return type:

Tensor

Returns:

Output tensor from the neural network.

evaluate(x, mu, with_last_layer=True, **kwargs)[source]

Evaluate the model with optional rollout over multiple steps.

Parameters:

  • x (Tensor) – Input tensor representing the initial state.

  • mu (Tensor) – Input tensor representing the parameters.

  • with_last_layer (bool) – Whether to include the last layer in evaluation.

  • **kwargs – Additional arguments (not used here).

Return type:

Tensor

Returns:

Output tensor after applying the model for the specified number of rollout

steps.

inference(x, mu, n=1, with_last_layer=True, **kwargs)[source]

Perform inference over multiple steps and return the trajectory.

Parameters:

  • x (Tensor) – Input tensor representing the initial state.

  • mu (Tensor) – Input tensor representing the parameters.

  • n (int) – Number of steps.

  • with_last_layer (bool) – Whether to include the last layer in evaluation (default: True).

  • **kwargs – Additional arguments (not used here).

Return type:

Tensor

Returns:

Tensor containing the trajectory of states over the steps.

set_dof(theta, flag_scope='all')[source]

Sets the degrees of freedom (DoF) for the neural network.

Parameters:

  • theta (Tensor) – A vector containing the network parameters.

  • flag_scope (str) – Scope flag for setting degrees of freedom..

Return type:

None

get_dof(flag_scope='all', flag_format='list')[source]

Retrieves the degrees of freedom (DoF) of the neural network.

Parameters:

  • flag_scope (str) – Scope flag for getting degrees of freedom.

  • flag_format (str) – The format for returning the parameters. Options are “list” or “tensor”.

Return type:

Tensor

Returns:

The network parameters in the specified format.

jacobian(x, mu)[source]

Compute the Jacobian of the network with respect to its parameters.

Parameters:

  • x (Tensor) – Input tensor from the spatial domain.

  • mu (Tensor) – Input tensor representing the parameters.

Return type:

Tensor

Returns:

Jacobian matrix of shape (num_samples, out_size, num_params).

class ContinuousFlowSpace(size, param_dim, **kwargs)[source]

Bases: Module

A continuous-time flow model that can be used for generating flows.

This model computes a forward pass by combining time, state, and parameters using different architectures, such as MLP or SympNet.

Parameters:

  • size (int) – The size of the input/output tensor.

  • param_dim (int) – The size of the parameters tensor.

  • **kwargs (Any) – Additional keyword arguments that specify the type of flow model (e.g., “mlp” or “sympnet”).

size: int

The size of the input/output tensor.

param_dim: int

The size of the parameters tensor.

flowtype: str

Type of the flow model, can be “mlp”, “sympnet”, etc.

net: Module

The neural network model used for transformations, either a MLP or a SympNet.

forward(t, x, mu)[source]

Computes the forward pass of the ContinuousFlow model.

Parameters:

  • t (Tensor) – The input tensor representing time, of shape (batch_size, 1).

  • x (Tensor) – The input tensor representing state, of shape (batch_size, size).

  • mu (Tensor) – The input tensor representing parameters, of shape (batch_size, param_dim).

Return type:

Tensor

Returns:

The output tensor after applying the transformation.