ChannelWeightNetwork

class madspace.ChannelWeightNetwork(self: madspace._madspace_py.ChannelWeightNetwork, channel_count: SupportsInt, particle_count: SupportsInt, hidden_dim: SupportsInt = 32, layers: SupportsInt = 3, activation: madspace._madspace_py.MLP.Activation = <Activation.leaky_relu: 1>, prefix: str = '', include_preprocessing: bool = True)

Bases: FunctionGenerator

Learned multi-channel weights from a neural network.

The MadNIS approach to multi-channeling: an MLP predicts the per-channel weights \(\alpha_i(x)\) from the MomentumPreprocessing features and a prior, replacing the fixed PropagatorChannelWeights [1, 2]. The network weights are compute-graph globals named after prefix. Call initialize_globals(context) once before use.

batch is the leading batch dimension. c is channel_count.

Arguments

• input – float, shape (batch, MomentumPreprocessing::output_dim()) – the invariant features.

• prior – float, shape (batch, c) – the prior per-channel weights.

Returns

• channel_weights – float, shape (batch, c) – the learned per-channel weights.

References

• [1] T. Heimel, O. Mattelaer, R. Winterhalder, “MadSpace”, https://arxiv.org/abs/2602.06895 (Sec. 2.1.1)

• [2] T. Heimel et al., “MadNIS - Neural multi-channel importance sampling”, https://arxiv.org/abs/2212.06172

Parameters:

• channel_count – Number of channels.

• particle_count – Number of external particles.

• hidden_dim – Hidden width of the MLP.

• layers – Number of MLP layers.

• activation – MLP activation.

• prefix – Prefix for the trainable global names.

• include_preprocessing – If true, apply MomentumPreprocessing to the input first.

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

Register the network’s trainable parameters on context.

Parameters:

• context – Context receiving the globals.

• seed – Seed for the random initialization. If unset (the default), the initialization is non-deterministic.

mask_name(self: madspace._madspace_py.ChannelWeightNetwork) → str

Global name of the channel mask.

mlp(self: madspace._madspace_py.ChannelWeightNetwork) → madspace._madspace_py.MLP

The underlying network.

preprocessing(self: madspace._madspace_py.ChannelWeightNetwork) → madspace._madspace_py.MomentumPreprocessing

The feature preprocessing.