lloca.framesnet.equi_frames.LearnedZFrames

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

Bases: LearnedFrames

Frames from ztransform, i.e. combination of boost along z and rotation around z axis, or SO(1,1)_z x SO(2)_z.

This is a special case of LearnedPolarDecompositionFrames where the boost vector is constrained to point along the z-axis and one of the rotation references is (0,0,0,1).

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.

  • 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