layers

Class(es) implementing layers to be used in graphnet models.

class graphnet.models.components.layers.DynEdgeConv(nn, aggr, nb_neighbors, features_subset, **kwargs)[source]

Bases: EdgeConv, LightningModule

Dynamical edge convolution layer.

Construct DynEdgeConv.

Parameters:

  • nn (Callable) – The MLP/torch.Module to be used within the EdgeConv.

  • aggr (str, default: 'max') – Aggregation method to be used with EdgeConv.

  • nb_neighbors (int, default: 8) – Number of neighbours to be clustered after the EdgeConv operation.

  • features_subset (Union[Sequence[int], slice, None], default: None) – Subset of features in Data.x that should be used when dynamically performing the new graph clustering after the EdgeConv operation. Defaults to all features.

  • **kwargs (Any) – Additional features to be passed to EdgeConv.

forward(x, edge_index, batch)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor) –

  • edge_index (Tensor | SparseTensor) –

  • batch (Tensor | None) –

class graphnet.models.components.layers.EdgeConvTito(nn, aggr, **kwargs)[source]

Bases: MessagePassing, LightningModule

Implementation of EdgeConvTito layer used in TITO solution for.

‘IceCube - Neutrinos in Deep’ kaggle competition.

Construct EdgeConvTito.

Parameters:

  • nn (Callable) – The MLP/torch.Module to be used within the EdgeConvTito.

  • aggr (str, default: 'max') – Aggregation method to be used with EdgeConvTito.

  • **kwargs (Any) – Additional features to be passed to EdgeConvTito.

reset_parameters()[source]

Reset all learnable parameters of the module.

Return type:

None

forward(x, edge_index)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor | Tuple[Tensor, Tensor]) –

  • edge_index (Tensor | SparseTensor) –

message(x_i, x_j)[source]

Edgeconvtito message passing.

Return type:

Tensor

Parameters:

  • x_i (Tensor) –

  • x_j (Tensor) –

class graphnet.models.components.layers.DynTrans(layer_sizes, aggr, features_subset, n_head, **kwargs)[source]

Bases: EdgeConvTito, LightningModule

Implementation of dynTrans1 layer used in TITO solution for.

‘IceCube - Neutrinos in Deep’ kaggle competition.

Construct DynTrans.

Parameters:

  • nn – The MLP/torch.Module to be used within the DynTrans.

  • layer_sizes (Optional[List[int]], default: None) – List of layer sizes to be used in DynTrans.

  • aggr (str, default: 'max') – Aggregation method to be used with DynTrans.

  • features_subset (Union[Sequence[int], slice, None], default: None) – Subset of features in Data.x that should be used when dynamically performing the new graph clustering after the EdgeConv operation. Defaults to all features.

  • n_head (int, default: 8) – Number of heads to be used in the multiheadattention models.

  • **kwargs (Any) – Additional features to be passed to DynTrans.

forward(x, edge_index, batch)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor) –

  • edge_index (Tensor | SparseTensor) –

  • batch (Tensor | None) –