lloca.framesnet.equi_frames.LearnedPDFrames

class lloca.framesnet.equi_frames.LearnedPDFrames(*args, gamma_max=None, gamma_hardness=None, deterministic_boost=None, compile=False, **kwargs)[source]

Bases: LearnedFrames

Frames as learnable polar decompositions.

This is our default approach. LearnedSO13Frames works similarly well, but is less flexible.

Parameters:

  • *args – Passed to LearnedFrames

  • **kwargs – Passed to LearnedFrames

  • gamma_max (float | None) – Maximum gamma factor for boost regularization. If None, no regularization is applied.

  • gamma_hardness (float | None) – Hardness, i.e. beta factor in the softplus regularization. If None, a hard clamp is applied.

  • deterministic_boost (str or None) – Deprecated option

  • 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