API Reference

Equivariant vector prediction

The Frames-Net uses a small Lorentz-equivariant network to predict a list of vectors. We currently have only one option for this equivectors network, but we plan to add more options in the future.

lloca.equivectors.mlp	Edge convolution with a simple MLP.
lloca.equivectors.lgatr	Edge convolution with L-GATr.
lloca.equivectors.pelican	Edge convolution with PELICAN.

Frames-Net

The equivariant vectors can be used in a range of Frames-Net procedures to construct local frames. We support three Lorentz-equivariant Frames-Net approaches, an SO(3)-equivariant approach and an SO(2)-equivariant approach. In addition, we implement non-equivariant networks as identity frames and data augmentation as random frames.

lloca.framesnet.equi_frames.LearnedFrames(...)	Abstract base class for learnable local frames
lloca.framesnet.equi_frames.LearnedPDFrames(*args)	Frames as learnable polar decompositions.
lloca.framesnet.equi_frames.LearnedSO13Frames(*args)	Frames as orthonormal set of Lorentz vectors.
lloca.framesnet.equi_frames.LearnedRestFrames(*args)	Rest frame transformation with learnable rotation.
lloca.framesnet.equi_frames.LearnedSO3Frames(*args)	Frames from SO(3) rotations.
lloca.framesnet.equi_frames.LearnedZFrames(*args)	Frames from ztransform, i.e. combination of boost along z and rotation around z axis, or SO(1,1)_z x SO(2)_z.
lloca.framesnet.equi_frames.LearnedSO2Frames(*args)	Frames from SO(2) rotations around the beam axis.
lloca.framesnet.nonequi_frames.IdentityFrames()	Identity frames for non-equivariant networks
lloca.framesnet.nonequi_frames.RandomFrames([...])	Random frames for data augmentation.

The resulting frames are stored in the Frames bookkeeping class. A range of derived class can be used for efficient access in the backbone architecture.

lloca.framesnet.frames.Frames([matrices, ...])	Bookkeeping class for local frames.
lloca.framesnet.frames.InverseFrames(frames)	Inverse of a collection of frames.
lloca.framesnet.frames.IndexSelectFrames(...)	Index-controlled collection of frames.
lloca.framesnet.frames.ChangeOfFrames(...)	Change of frames between two Frames objects.
lloca.framesnet.frames.LowerIndicesFrames(frames)	Frames with lower indices, obtained by multiplying with the metric.

Backbone networks

The LLoCa framework can be used to make generic backbone architectures Lorentz-equivariant.

1. Transform the network inputs into their local frames to make them invariant.

2. For message-passing architectures, transform the messages from the sender frame to the receiver frame using a non-trivial message representation, i.e. not only scalars.

3. Transform the network outputs back to the global frame to obtain a Lorentz-equivariant output. This step is trivial in the case of Lorentz-invariant outputs.

The MLP does not require any modifications to be used in the LLoCa framework. For message-passing architectures, we provide the LLoCaMessagePassing class to conveniently adapt graph networks based on the torch_geometric.nn.conv.MessagePassing class to the LLoCa framework. For transformers, we provide the LLoCaAttention class as a drop-in replacement for torch.nn.functional.scaled_dot_product_attention and other attention backends. We demonstrate how to use these tools with a baseline GraphNet and a Transformer. For ParticleNet and ParticleTransformer, we demonstrate how to use LLoCa with established architectures.

lloca.backbone.mlp.MLP(in_shape, out_shape, ...)	A simple MLP.
lloca.backbone.lloca_message_passing.LLoCaMessagePassing(...)	Adaptation of the torch_geometric MessagePassing class using the LLoCa formalism.
lloca.backbone.attention.LLoCaAttention(...)	Attention with frame-to-frame transformations.
lloca.backbone.graphnet.GraphNet(...[, ...])	Baseline LLoCa-GNN.
lloca.backbone.transformer.Transformer(...)	Baseline LLoCa-Transformer.
lloca.backbone.particlenet.ParticleNet(...)	ParticleNet with local frame transformations.
lloca.backbone.particletransformer.ParticleTransformer(...)	Particle Transformer (ParT) with local frame transformations.

Lorentz group representations

The LLoCa framework supports arbitrary Lorentz group representations for inputs, outputs and messages. The TensorReps class organizes the properties of these representations, while the TensorRepsTransform class implements the actual transformations of features.

lloca.reps.tensorreps.TensorReps(input[, ...])	Generic tensor representations
lloca.reps.tensorreps_transform.TensorRepsTransform(reps)	Tensor representation transformation module.

Utilities

Finally, we provide a range of utility functions for Lorentz transformations, random transformations and orthogonalization.

lloca.utils.utils	Tools for graph construction and manipulation.
lloca.utils.lorentz	Basic operations in Minkowski space.
lloca.utils.rand_transforms	Random Lorentz transformations.
lloca.utils.orthogonalize_3d	Orthogonalization of euclidean vectors.
lloca.utils.orthogonalize_4d	Orthogonalization of Minkowski vectors.
lloca.utils.polar_decomposition	Lorentz transformations from boosts and rotations.