Flow¶

class madspace.Flow(self: madspace._madspace_py.Flow, input_dim: SupportsInt, condition_dim: SupportsInt = 0, prefix: str = '', bin_count: SupportsInt = 10, subnet_hidden_dim: SupportsInt = 32, subnet_layers: SupportsInt = 3, subnet_activation: madspace._madspace_py.MLP.Activation = <Activation.leaky_relu: 1>, invert_spline: bool = True)¶

Bases: Mapping

Normalizing-flow adaptive-sampling mapping.

A coupling-based normalizing flow built from rational-quadratic spline transformations, whose bin positions are predicted by MLP subnetworks [1, 2]. It reparametrizes the unit hypercube into a learned sampling distribution and is trained with MadNIS [3]. The subnetwork weights are compute-graph globals named after prefix; call initialize_globals(context) once before use.

batch is the leading batch dimension.

Inputs

latent – float, shape (batch, input_dim) – uniform coordinates in [0, 1).

Conditions

c – float, shape (batch, condition_dim) – conditioning input. Present only when condition_dim is nonzero.

Outputs

data – float, shape (batch, input_dim) – the transformed coordinates.

In addition every mapping returns a weight (float, shape (batch,)), the Jacobian of the transformation.

References

[1] T. Heimel, O. Mattelaer, R. Winterhalder, “MadSpace”, https://arxiv.org/abs/2602.06895 (Sec. 3.2.2)
[2] C. Durkan et al., “Neural spline flows”, https://arxiv.org/abs/1906.04032
[3] T. Heimel et al., “MadNIS - Neural multi-channel importance sampling”, https://arxiv.org/abs/2212.06172

Parameters:

input_dim – Number of dimensions to transform.
condition_dim – Width of the conditioning input; 0 for an unconditional flow.
prefix – Prefix for the trainable global names.
bin_count – Number of spline bins per dimension.
subnet_hidden_dim – Hidden width of the MLP subnetworks.
subnet_layers – Number of layers in the subnetworks.
subnet_activation – Activation of the subnetworks.
invert_spline – Whether the spline is applied in the inverse direction.

condition_dim(self: madspace._madspace_py.Flow) → int¶: Width of the conditioning input.

initialize_from_vegas(self: madspace._madspace_py.Flow, context: madspace._madspace_py.Context, grid_name: str, seed: SupportsInt | None = None) → None¶

Initialize the flow to reproduce a trained VEGAS grid.

Parameters:

context – Context receiving the globals.
grid_name – Global name of the VEGAS grid to copy.
seed – Seed for the random initialization of the remaining parameters. If unset (the default), that initialization is non-deterministic.

initialize_globals(self: madspace._madspace_py.Flow, context: madspace._madspace_py.Context, seed: SupportsInt | None = None) → None¶

Parameters:

context – Context receiving the globals.
seed – Seed for the random initialization. If unset (the default), the initialization is non-deterministic.

input_dim(self: madspace._madspace_py.Flow) → int¶: Number of transformed dimensions.