lloca.framesnet.equi_frames.LearnedSO2Frames

class lloca.framesnet.equi_frames.LearnedSO2Frames(*args, compile=False, **kwargs)[source]

Bases: LearnedFrames

Frames from SO(2) rotations around the beam axis.

This is a special case of LearnedPolarDecompositionFrames where the firsts two vectors are trivial (1,0,0,0) and (0,0,0,1).

Parameters:

  • *args – Passed to LearnedFrames

  • **kwargs – Passed to LearnedFrames

  • compile (bool) – Option to compile the orthonormalization procedure. Does not yet give significant speedups in our tests.

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

  • fourmomenta (torch.Tensor) – Tensor of shape (…, 4) containing the four-momenta

  • scalars (torch.Tensor or None) – Optional tensor of shape (…, n_scalars) containing additional scalar features

  • ptr (torch.Tensor or None) – Pointer for sparse tensors, or None for dense tensors

  • return_tracker (bool) – If True, return a tracker dictionary with regularization information

Returns:

  • Frames – Local frames constructed from the polar decomposition of the four-momenta

  • tracker (dict (optional)) – Dictionary containing regularization information, if return_tracker is True