matcha.torch.encoders.gps

General, Powerful, Scalable (GPS) Graph Transformer encoder.

Classes

MPNNPlusConv

Minimal reimplementation of MPNNPlus convolution layer for GPS.

GPSBlock

General, Powerful, Scalable (GPS) Graph Transformer Block.

GPS

General, Powerful, Scalable (GPS) Graph Transformer encoder.

Module Contents

class matcha.torch.encoders.gps.MPNNPlusConv(in_dim: int, out_dim: int, in_dim_edges: int, out_dim_edges: int, activation: str = 'gelu', dropout: float = 0.1, edge_dropout: float = 0.0, mlp_expansion_ratio: int = 4)[source]

Bases: torch.nn.Module

Minimal reimplementation of MPNNPlus convolution layer for GPS.

GPS++ style message passing layer that updates both node and edge features. Uses variance-preserving aggregation for better gradient flow. No internal normalization - relies on pre-norm pattern in GPSBlock.

Reference: https://arxiv.org/abs/2212.02229

Parameters:

  • in_dim (int) – Input node feature dimension

  • out_dim (int) – Output node feature dimension

  • in_dim_edges (int) – Input edge feature dimension

  • out_dim_edges (int) – Output edge feature dimension

  • activation (str) – Activation function name

  • dropout (float) – Node dropout rate

  • edge_dropout (float) – Edge dropout rate

  • mlp_expansion_ratio (int) – Hidden dim expansion factor

aggregator
edge_model
node_model
forward(x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor) → tuple[torch.Tensor, torch.Tensor][source]

Forward pass.

Parameters:

  • x (torch.Tensor) – Node features [num_nodes, in_dim]

  • edge_index (torch.Tensor) – Edge indices [2, num_edges]

  • edge_attr (torch.Tensor) – Edge features [num_edges, in_dim_edges]

Return tuple:

Updated node and edge features

class matcha.torch.encoders.gps.GPSBlock(atom_feats: int, edge_feats: int, dropout: float = 0.0, norm: str = 'adarmsn', activation: str = 'swish', num_heads: int = 4, expansion_k: int = 2, distance_k: int | None = 5)[source]

Bases: torch.nn.Module

General, Powerful, Scalable (GPS) Graph Transformer Block.

Combines local message passing (MPNNPlus) with global self-attention. Uses PyTorch Geometric for graph operations.

References: - GPS: https://arxiv.org/abs/2205.12454 - GPS++: https://arxiv.org/abs/2212.02229

Parameters:

  • atom_feats (int) – Node feature dimension

  • edge_feats (int) – Edge feature dimension

  • dropout (float) – Dropout ratio

  • norm (str) – Normalization type

  • activation (str) – Activation function

  • num_heads (int) – Number of attention heads

  • expansion_k (int) – FFN expansion factor

  • distance_k (int) – Maximum distance for spatial encoding (unused, kept for API compatibility)

mp
att
norm1_local
norm1_global
norm2
mlp
forward(graph: torch_geometric.data.Batch, feat: torch.Tensor, edge_feat: torch.Tensor, graph_id: torch.Tensor, dist_bias: torch.Tensor | None) → tuple[torch.Tensor, torch.Tensor][source]

Forward pass.

Parameters:

  • graph (Batch) – PyG batched graph

  • feat (torch.Tensor) – Node features [num_nodes, atom_feats]

  • edge_feat (torch.Tensor) – Edge features [num_edges, edge_feats]

  • graph_id (torch.Tensor) – Batch assignment [num_nodes]

  • dist_bias (torch.Tensor) – Distance bias for attention or None

Return tuple:

Updated node and edge features

class matcha.torch.encoders.gps.GPS(num_layers: int, atom_input_dim: int, bond_input_dim: int, atom_hidden_dim: int, num_heads: int, expansion_k: int, distance_k: float | None, activation: str, dropout: float, norm: str | None, jk: str, readout: str, laplacian_k: int, rwse_k: int, elstatic_k: int, distmat_k: int, rrwp_k: int)[source]

Bases: matcha.torch.encoders.base_graph_encoder.BaseGraphEncoder, lightning.pytorch.core.mixins.HyperparametersMixin

General, Powerful, Scalable (GPS) Graph Transformer encoder.

Combines local message passing with global self-attention for learning molecular representations. Uses PyTorch Geometric for all graph operations.

It inherits from BaseGraphEncoder for common graph encoding routines (e.g. jk-related routines) and from lightning.pytorch.core.mixins for saving its hyperparameters.

References: - https://arxiv.org/abs/2404.11568 - https://arxiv.org/abs/2205.12454 - https://proceedings.mlr.press/v202/ma23c.html

It is intended to be used inside a BaseClassicModel instance. Check the docs of matcha.torch.models.classic.GPSModel for further details.

Parameters:

  • num_layers (int) – number of message passing layers

  • atom_input_dim (int) – number of input atom features from GraphFeaturizer

  • bond_input_dim (int) – number of input bond features from GraphFeaturizer

  • atom_hidden_dim (int) – number of hidden atom (and bond) features in message passing layers

  • num_heads (int) – number of attention heads, must divide atom_hidden_dim evenly

  • expansion_k (int) – expansion factor for the feed-forward network in GPSBlock

  • distance_k (float | None) – Upper bound for the shortest path distance to encode

  • activation (str) – activation function to be used in all layers

  • dropout (float) – dropout noise level

  • norm (str | None) – which norm to use inside LnBnDr layers

  • jk (str) – jumping knowledge strategy to use when returning molecular representations after forward pass

  • readout (str) – readout function to aggregate all atom representations

layers
atom_projection
bond_projection
forward(graph: torch_geometric.data.Batch) → torch.Tensor[source]

Converts a batched PyG graph into a (batch_size, fp_dim) tensor for further processing.

Parameters:

graph (Batch) – batched PyG graph from the dataloader

Return torch.Tensor:

learned representation