lloca.equivectors.mlp

Edge convolution with a simple MLP.

Functions

get_edge_index_and_batch(fourmomenta, ptr[, ...])

get_nonlinearity(nonlinearity)

get_operation(operation)

softmax(src[, index, ptr, dim])

Adapted version of the torch_geometric softmax function https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/utils/_softmax.html.

Classes

EquiEdgeConv(in_vectors, out_vectors, ...[, ...])

Equivariant edge convolution, implemented using torch_geometric's MessagePassing class.

MLPVectors(n_vectors, *args, **kwargs)

Edge convolution with a simple MLP.
class lloca.equivectors.mlp.EquiEdgeConv(in_vectors, out_vectors, num_scalars, hidden_channels, num_layers_mlp, include_edges=True, operation='add', nonlinearity='softmax', fm_norm=True, layer_norm=True, use_amp=False, dropout_prob=None, aggr='sum')[source]

Equivariant edge convolution, implemented using torch_geometric’s MessagePassing class.

The choice of the parameters operation, nonlinearity, fm_norm, aggr, layer_norm is critical to the stability of the approach. Bad combinations initialize the framesnet to predict strongly boosted vectors, leading to strongly boosted frames and unstable training. We recommend to stick to the default parameters, which worked for all our experiments.

Parameters:

  • in_vectors (int) – Number of input vectors.

  • out_vectors (int) – Number of output vectors.

  • num_scalars (int) – Number of scalar features per particle.

  • hidden_channels (int) – Number of hidden channels in the MLP.

  • num_layers_mlp (int) – Number of hidden layers in the MLP.

  • include_edges (bool) – Whether to include edge attributes in the message passing. If True, edge attributes will be calculated from fourmomenta and standardized. Default is True.

  • operation (str) – Operation to perform on the fourmomenta. Options are “add”, “diff”, or “single”. Default is “add”.

  • nonlinearity (str) – Nonlinearity to apply to the output of the MLP. Options are “exp”, “softplus”, and “softmax”. Default is “softmax”.

  • fm_norm (bool) – Whether to normalize the relative fourmomentum. Default is True.

  • layer_norm (bool) – Whether to apply Lorentz-equivariant layer normalization to the output vectors. Default is True.

  • use_amp (bool) – Whether to use automatic mixed precision (AMP) for the MLP. Default is False.

  • dropout_prob (float) – Dropout probability for the MLP. If None, no dropout will be applied. Default is None.

  • aggr (str) – Aggregation method for message passing. Options are “add”, “mean”, or “max”. Default is “sum”.

forward(fourmomenta, scalars, edge_index, ptr, batch=None)[source]
Parameters:

  • fourmomenta (torch.Tensor) – Tensor of shape (num_particles, in_vectors*4) containing the fourmomenta of the particles.

  • scalars (torch.Tensor) – Tensor of shape (num_particles, num_scalars) containing scalar features for each particle.

  • edge_index (torch.Tensor) – Edge index tensor containing the indices of the source and target nodes, shape (2, num_edges).

  • ptr (torch.Tensor) – Pointer tensor indicating the start of each batch for sparse tensors, shape (num_batches+1,).

  • batch (torch.Tensor, optional) – Batch tensor indicating the batch each particle belongs to. If None, all particles are assumed to belong to the same batch.

Returns:

Tensor of shape (num_particles, out_vectors*4) containing the predicted vectors for each edge.

Return type:

torch.Tensor

message(edge_index, s_i, s_j, fm_i, fm_j, node_ptr, node_batch, edge_attr=None)[source]
Parameters:

  • edge_index (torch.Tensor) – Edge index tensor containing the indices of the source and target nodes, shape (2, num_edges).

  • s_i (torch.Tensor) – Scalar features of the source nodes, shape (num_edges, num_scalars).

  • s_j (torch.Tensor) – Scalar features of the target nodes, shape (num_edges, num_scalars).

  • fm_i (torch.Tensor) – Fourmomentum of the source nodes, shape (num_edges, 4*in_vectors).

  • fm_j (torch.Tensor) – Fourmomentum of the target nodes, shape (num_edges, 4*in_vectors).

  • node_ptr (torch.Tensor) – Pointer tensor indicating the start of each batch for sparse tensors, shape (num_batches+1,).

  • edge_attr (torch.Tensor, optional) – Edge attributes tensor. If None, no edge attributes will be used, shape (num_edges, num_edge_attributes).

Returns:

Tensor of shape (num_edges, out_vectors*4) containing the predicted vectors for each edge.

Return type:

torch.Tensor

class lloca.equivectors.mlp.MLPVectors(n_vectors, *args, **kwargs)[source]

Edge convolution with a simple MLP.

Parameters:

  • n_vectors (int) – Number of output vectors per particle. Different FramesPredictor’s need different n_vectors, so this parameter should be set dynamically.

  • *args

  • **kwargs

forward(fourmomenta, scalars=None, ptr=None, **kwargs)[source]
Parameters:

  • fourmomenta (torch.Tensor) – Tensor of shape (…, 4) containing the fourmomenta of the particles.

  • scalars (torch.Tensor, optional) – Tensor of shape (…, num_scalars) containing scalar features for each particle. If None, a tensor of zeros will be created.

  • ptr (torch.Tensor, optional) – Pointer tensor indicating the start and end of each batch for sparse tensors.

Returns:

Tensor of shape (…, n_vectors, 4) containing the predicted vectors for each particle.

Return type:

torch.Tensor

lloca.equivectors.mlp.get_nonlinearity(nonlinearity)[source]
Parameters:

nonlinearity (str) – Nonlinearity to apply to the output of the MLP. Options are “exp”, “softplus”, “softmax”. We enforce the prediction of timelike vectors.

Returns:

A function that applies the specified nonlinearity to the input tensor.

Return type:

callable

lloca.equivectors.mlp.get_operation(operation)[source]
Parameters:

operation (str) – Operation to perform on the fourmomenta. Options are “add”, “diff”, or “single”.

Returns:

A function that performs the specified operation on two fourmomenta tensors.

Return type:

callable

lloca.equivectors.mlp.softmax(src, index=None, ptr=None, dim=0)[source]

Adapted version of the torch_geometric softmax function https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/utils/_softmax.html. Use the index argument in output_size of torch.repeat_interleave to avoid GPU/CPU sync.

Parameters:

  • src (torch.Tensor) – Source tensor of shape (N,) where N is the number of elements.

  • index (torch.Tensor, optional) – Index tensor indicating the batch index for each element. Tensor of shape (N,) where N is the number of elements.

  • ptr (torch.Tensor) – Pointer tensor indicating the start of each batch. Tensor of shape (B+1,) where B is the number of batches.

  • dim (int, optional) – Dimension along which to apply the softmax. Default is 0.