lloca.framesnet.equi_frames.LearnedFrames

class lloca.framesnet.equi_frames.LearnedFrames(equivectors, n_vectors, is_global=False, random=False, fix_params=False, mass_reg=None, ortho_kwargs=None)

Bases: FramesPredictor

Abstract base class for learnable local frames

Parameters:

• equivectors (nn.Module) – Network that equivariantly predicts vectors

• n_vectors (int) – Number of vectors to predict

• is_global (bool) – If True, average the predicted vectors to construct a global frame

• random (bool) – If True, re-initialize the equivectors at each forward pass. This is equivalent to data augmentation for is_global=True.

• fix_params (bool) – Fix the Frames-Net parameters. This is equivalent to random, but without the resampling. We find that this can be useful to avoid overfitting.

• mass_reg (float | None)

• ortho_kwargs (dict) – Keyword arguments for orthogonalization

globalize_vecs_or_not(vecs, ptr)

init_weights_or_not()

mass_regularize(fourmomenta)