lloca.framesnet.nonequi_frames.COMRandomFrames

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

Bases: RandomFrames

Special random frame for data augmentation in amplitude regression.

Modifies the forward function of RandomFrames such that an additional boost is applied to the whole event. Only applicable to amplitude regression, the boost changes the reference frame to the center of mass of the incoming particles.

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.