lloca.framesnet.nonequi_frames.RandomFrames

class lloca.framesnet.nonequi_frames.RandomFrames(transform_type='lorentz', is_global=True, std_eta=0.1, n_max_std_eta=3.0)[source]

Bases: FramesPredictor

Random frames for data augmentation.

Parameters:

  • transform_type (str) – Type of random transformation. One of “lorentz”, “rotation”, “xyrotation”, “ztransform”.

  • is_global (bool) – Global or local data augmentations, the default is global. Local data augmentations are a weird thing, we implemented them because we can.

  • std_eta (float) – Standard deviation of the rapidity eta for the random boost. Only relevant if transform_type is “lorentz” or “ztransform”.

  • n_max_std_eta (float) – Maximum rapidity in units of std_eta. Only relevant if transform_type is “lorentz” or “ztransform”.

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

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

transform(shape, device, dtype)[source]