graphnet.models.components.embedding module

Classes for performing embedding of input data.

class graphnet.models.components.embedding.SinusoidalPosEmb(dim, n_freq, scaled)[source]

Bases: LightningModule

Sinusoidal positional embeddings module.

This module is from the kaggle competition 2nd place solution (see arXiv:2310.15674): It performs what is called Fourier encoding or it’s used in the Attention is all you need arXiv:1706.03762. It can be seen as a soft digitization of the input data

Construct SinusoidalPosEmb.

Parameters:

  • dim (int, default: 16) – Embedding dimension.

  • n_freq (int, default: 10000) – Number of frequencies.

  • scaled (bool, default: False) – Whether or not to scale the output.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.embedding.FourierEncoder(seq_length, mlp_dim, output_dim, scaled, n_features)[source]

Bases: LightningModule

Fourier encoder module.

This module incorporates sinusoidal positional embeddings and auxiliary embeddings to process input sequences and produce meaningful representations. The module assumes that the input data is in the format of (x, y, z, time, charge, auxiliary), being the first four features mandatory.

Construct FourierEncoder.

Parameters:

  • seq_length (int, default: 128) – Dimensionality of the base sinusoidal positional embeddings.

  • mlp_dim (Optional) – Size of hidden, latent space of MLP. If not given, mlp_dim is set automatically as multiples of seq_length (in consistent with the 2nd place solution), depending on n_features.

  • output_dim (int, default: 384) – Dimension of the output (I.e. number of columns).

  • scaled (bool, default: False) – Whether or not to scale the embeddings.

  • n_features (int, default: 6) – The number of features in the input data.

forward(x, seq_length)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • seq_length (Tensor)

class graphnet.models.components.embedding.SpacetimeEncoder(seq_length)[source]

Bases: LightningModule

Spacetime encoder module.

Construct SpacetimeEncoder.

This module calculates space-time interval between each pair of events and generates sinusoidal positional embeddings to be added to input sequences.

Parameters:

seq_length (int, default: 32) – Dimensionality of the sinusoidal positional embeddings.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)