scimba_torch.neural_nets.structure_preserving_nets.coupling_symplectic_layers

Symplectic coupling layers for structure-preserving neural networks.

Classes

GSymplecticLayer(size, conditional_size, ...)	A G-symplectic coupling layer.
LASymplecticLayer(size, conditional_size, ...)	A LA-symplectic coupling layer.
PeriodicGSymplecticLayer(size, ...)	A periodic G-symplectic coupling layer.

class GSymplecticLayer(size, conditional_size, width, **kwargs)

Bases: InvertibleLayer

A G-symplectic coupling layer.

Applies coupling transformations where one variable is kept fixed while transforming the other, preserving the symplectic structure.

Parameters:

• size (int) – Total dimension of the state space (will be split in half).

• conditional_size (int) – Dimension of the conditional input.

• width (int) – Width of the internal layers.

• **kwargs – Additional keyword arguments.

forward(y, mu)

Applies the forward coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size) containing (p, q).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Transformed tensor of shape (batch_size, size).

backward(y, mu)

Applies the inverse coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Inverse transformed tensor of shape (batch_size, size).

log_abs_det_jacobian(y, mu)

Computes the log absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Log determinant of shape (batch_size,).

abs_det_jacobian(y, mu)

Computes the absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Determinant of shape (batch_size,).

class PeriodicGSymplecticLayer(size, conditional_size, width, period, **kwargs)

Bases: InvertibleLayer

A periodic G-symplectic coupling layer.

Applies coupling transformations where one variable is kept fixed while transforming the other, preserving the symplectic structure. The first variable is periodized, the second one is not.

Parameters:

• size (int) – Total dimension of the state space (will be split in half).

• conditional_size (int) – Dimension of the conditional input.

• width (int) – Width of the internal layers.

• period (Tensor) – The period of the potential.

• **kwargs – Additional keyword arguments.

forward(y, mu)

Applies the forward coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size) containing (p, q).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Transformed tensor of shape (batch_size, size).

backward(y, mu)

Applies the inverse coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Inverse transformed tensor of shape (batch_size, size).

log_abs_det_jacobian(y, mu)

Computes the log absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Log determinant of shape (batch_size,).

abs_det_jacobian(y, mu)

Computes the absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Determinant of shape (batch_size,).

class LASymplecticLayer(size, conditional_size, width, **kwargs)

Bases: InvertibleLayer

A LA-symplectic coupling layer.

Applies Linear-Activation coupling transformations preserving symplectic structure.

Parameters:

• size (int) – Total dimension of the state space (will be split in half).

• conditional_size (int) – Dimension of the conditional input.

• width (int) – Width of the internal layers.

• **kwargs – Additional keyword arguments.

forward(y, mu)

Applies the forward coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size) containing (p, q).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Transformed tensor of shape (batch_size, size).

backward(y, mu)

Applies the inverse coupling transformation.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Inverse transformed tensor of shape (batch_size, size).

log_abs_det_jacobian(y, mu)

Computes the log absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Log determinant of shape (batch_size,).

abs_det_jacobian(y, mu)

Computes the absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – Input tensor of shape (batch_size, size).

• mu (Tensor) – Conditional input of shape (batch_size, conditional_size).

Return type:

Tensor

Returns:

Determinant of shape (batch_size,).