embedding

Classes for performing embedding of input data.

class graphnet.models.components.embedding.LinearEdgeEncoder(dim_emb)[source]

Bases: LightningModule

Linear encoding for edge attributes.

Construct LinearEdgeEncoder.

Parameters:

dim_emb (int) – Embedding dimension.

forward(data)[source]

Forward pass.

Return type:

Data

Parameters:

data (Data)

class graphnet.models.components.embedding.LinearNodeEncoder(dim_in, dim_emb)[source]

Bases: LightningModule

Linear encoding for nodes.

Construct LinearNodeEncoder.

Parameters:

  • dim_in (int) – Input dimension.

  • dim_emb (int) – Embedding dimension.

forward(data)[source]

Forward pass.

Return type:

Data

Parameters:

data (Data)

class graphnet.models.components.embedding.SinusoidalPosEmb(dim, n_freq, scaled, phase_dtype)[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. The ladder holds dim / 2 frequencies, each contributing a sine and a cosine.

  • n_freq (float, default: 10000.0) – Span of the frequency ladder. The frequencies run from 1 down to n_freq ** -((dim/2 - 1)/(dim/2)), so this sets how far the embedding reaches beyond its finest scale, not how many frequencies there are.

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

  • phase_dtype (Optional[dtype], default: torch.float32) – Precision in which the phase and its sine and cosine are evaluated, with the result cast back to the input dtype. A large multiplier makes the phase large enough that bf16 or fp16 can no longer separate neighbouring positions, so it defaults to float32; None keeps the input dtype.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

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

Bases: LightningModule

Fourier encoder of the IceMix Kaggle solution.

The graphnet implementation of the encoder as presented in the EPJ-C publication (arXiv:2310.15674). It incorporates sinusoidal positional embeddings and auxiliary embeddings to process input sequences and produce meaningful representations.

It carries assumptions from that competition that make it hard to use elsewhere. Use FourierEncoder unless you are using the model and data from the EPJ-C paper.

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.

  • phase_dtype (Optional[dtype], default: torch.float32) – Precision for the sinusoidal phase computation; see SinusoidalPosEmb.

forward(x, seq_length)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • seq_length (Tensor)

class graphnet.models.components.embedding.FourierEncoder(schema, seq_length, scaled, add_sequence_length, n_freq, phase_dtype)[source]

Bases: LightningModule

Apply sinusoidal positional encodings to sequence representations.

Embeds [B, K, D] sequences – batch size, padded sequence length, and features per step – into sinusoidal positions, optionally alongside the unpadded length of each sequence.

Construct FourierEncoder.

Parameters:

  • schema (Dict[int, Union[float, Tuple[float, float]]]) –

    Maps an input column index to how it is embedded, either as a multiplier or as a (multiplier, n_freq) pair, e.g. {1: 1024, 3: (4096, 500)}. Columns absent from the schema are not embedded.

    The two together set the band of separations a column can resolve: wavelengths from 2 * pi * unit / multiplier up to that times n_freq ** ((dim/2 - 1)/(dim/2)), where unit is the physical size of one unit of the (normalised) column. So the multiplier places the band and n_freq sets its width, and both are properties of the data rather than of the model – values carried over from a different normalisation resolve a different physical range.

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

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

  • add_sequence_length (bool, default: True) – If True, the unpadded length of each sequence is embedded as well, at half width.

  • n_freq (float, default: 10000.0) – Ladder span for columns whose schema entry gives only a multiplier, and for the sequence-length embedding.

  • phase_dtype (Optional[dtype], default: torch.float32) – Precision for the sinusoidal phase computation; see SinusoidalPosEmb.

forward(x, seq_length)[source]

Apply positional encoding of x.

Parameters:

  • x (Tensor) – [B, K, D]-dimensional sequence representation of an event.

  • seq_length (Optional[Tensor], default: None) – [B]-dimensional unpadded length of each sequence in x, one entry per sequence rather than per step. Required when add_sequence_length is True, in which case each step of a sequence receives a copy of its sequence’s embedded length.

Return type:

Tensor

Returns:

Embedded [B, K, J]-dimensional sequence.

graphnet.models.components.embedding.signed_four_distance(x, time_scale, columns=(0, 1, 2, 3))[source]

Signed spacetime four-distance between every pair of steps.

sign(s) * sqrt(|s|) of the interval dx^2 - (c dt)^2. Positive is spacelike.

time_scale converts the time column into the same length unit as the position columns, i.e. t_scale * c / pos_scale for whatever normalisation produced the features. It is a property of that normalisation rather than of the physics, so a value carried over from other data silently reduces the interval to a spatial distance.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • time_scale (float)

  • columns (Tuple[int, int, int, int])

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

Bases: LightningModule

Spacetime encoder of the IceMix Kaggle solution.

The graphnet implementation as presented in the EPJ-C publication (arXiv:2310.15674). It computes the spacetime interval between each pair of pulses and embeds it sinusoidally.

It carries assumptions from that competition: position must be columns 0-2 and time column 3, and the constant converting time into a length, the multiplier before the sinusoidal ladder and the clip are all fixed to the Kaggle dataset’s normalisation. See SpacetimeEncoder for a version without them.

Construct SpacetimeEncoderEPJC.

Parameters:

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

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.embedding.SpacetimeEncoder(seq_length, output_dim, columns=(0, 1, 2, 3), time_scale, scale, clip, n_freq)[source]

Bases: LightningModule

Embed the spacetime interval between every pair of steps.

A refactor of SpacetimeEncoderEPJC that drops the assumptions tying it to the Kaggle dataset, so the same encoder can be reused on other data: which columns carry the coordinates, how time converts into a length, and which band of separations the embedding resolves are all stated by the caller rather than fixed.

Construct SpacetimeEncoder.

Parameters:

  • seq_length (int, default: 32) – Dimensionality of the sinusoidal embedding.

  • output_dim (Optional[int], default: None) – Width of the projected per-pair feature. Defaults to seq_length.

  • columns (Tuple[int, int, int, int], default: (0, 1, 2, 3)) – Input columns holding (x, y, z, t), so the encoder reads the right ones whatever order the data arrives in.

  • time_scale (float, default: 1.0) – Factor converting the time column into the position columns’ length unit, t_scale * c / pos_scale for the normalisation in use. The default assumes time already is a length, as it is when the interval needs no constant.

  • scale (float, default: 1024.0) – Multiplier applied before the sinusoidal ladder. With n_freq it sets the resolved band: wavelengths from 2 * pi / scale up to that times n_freq ** ((seq_length/2 - 1)/(seq_length/2)), in the unit the interval is measured in.

  • clip (Optional[float], default: 4.0) – Bound on the four-distance before embedding, or None to leave it unbounded. The tail is heavy, so without a bound a few far-separated pairs dominate.

  • n_freq (float, default: 10000.0) – Span of the frequency ladder.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.embedding.RRWPLinearNodeEncoder(emb_dim, out_dim, use_bias, apply_norm, norm_layer=<class 'torch.nn.modules.batchnorm.BatchNorm1d'>, pe_name)[source]

Bases: LightningModule

Relative random walk probability node encoder.

Original code: https://github.com/LiamMa/GRIT/blob/main/grit/encoder/rrwp_encoder.py

Construct RRWPLinearNodeEncoder.

Parameters:

  • emb_dim (int) – Embedding dimension.

  • out_dim (int) – Output dimension.

  • use_bias (bool, default: False) – Apply bias to linear layer.

  • apply_norm (bool, default: True) – Apply normalization layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.batchnorm.BatchNorm1d'>) – Normalization layer.

  • pe_name (str, default: 'rrwp') – Positional encoding name.

forward(data)[source]

Forward pass.

Return type:

Data

Parameters:

data (Data)

class graphnet.models.components.embedding.RRWPLinearEdgeEncoder(emb_dim, out_dim, use_bias, apply_norm, norm_layer=<class 'torch.nn.modules.batchnorm.BatchNorm1d'>, pad_to_full_graph, fill_value, add_node_attr_as_self_loop, overwrite_old_attr)[source]

Bases: LightningModule

Relative random walk probability edge encoder.

Original code: https://github.com/LiamMa/GRIT/blob/main/grit/encoder/rrwp_encoder.py

Construct RRWPLinearEdgeEncoder.

Parameters:

  • emb_dim (int) – Embedding dimension.

  • out_dim (int) – Output dimension.

  • use_bias (bool, default: False) – Apply bias to linear layer.

  • apply_norm (bool, default: True) – Apply normalization layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.batchnorm.BatchNorm1d'>) – Normalization layer.

  • pad_to_full_graph (bool, default: True) – Pad edges to fully-connected graph.

  • fill_value (float, default: 0.0) – Fill value for padding.

  • add_node_attr_as_self_loop (bool, default: False) – Add self loop edges with node attr.

  • overwrite_old_attr (bool, default: False) – Overwrite old edge attr.

forward(data)[source]

Forward pass.

Return type:

Data

Parameters:

data (Data)

class graphnet.models.components.embedding.RWSELinearNodeEncoder(emb_dim, out_dim, use_bias)[source]

Bases: LightningModule

Random walk structural node encoding.

Construct RWSELinearEdgeEncoder.

Parameters:

  • emb_dim (int) – Embedding dimension.

  • out_dim (int) – Output dimension.

  • use_bias (bool, default: False) – Apply bias to linear layer.

forward(data)[source]

Forward pass.

Return type:

Data

Parameters:

data (Data)