# SPDX-License-Identifier: Apache-2.0 # The GNN architecture and pretrained weights derive from the "Suiren" molecular # graph model (released here under Apache-2.0 with the authors' permission). # Depends on RDKit (BSD-3-Clause) and PyTorch-Geometric (MIT). # See THIRD_PARTY_LICENSES.md at the repo root. """ Molecular Graph Neural Network (GNN) Models for 2D Property Prediction This module provides implementation of Graph Attention Network (GAT) based GNN for molecular property prediction tasks. It includes pre-training and fine-tuning capabilities with multi-head attention mechanisms. """ import torch import torch.nn as nn from torch_geometric.nn import MessagePassing from torch_geometric.utils import add_self_loops, softmax from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set import torch.nn.functional as F from typing import Any, Dict, List, Optional, Sequence, Tuple from torch_geometric.data import Data import pandas as pd from .org_mol2d import from_smiles # Elements with atomic number < 100 (H through Es / Z=99) are allowed; this # matches the size of ``x_embedding1`` (NUM_ATOM_TYPE=100 rows). # Z >= 100 (Fm/Md/No/Lr/transactinides) are rejected at the # build_graph filter (would index past x_embedding1's 100 rows and cause # CUDA illegal memory access). They are not chemically meaningful for # any biological / pharmaceutical dataset. ALLOWED_ELEMENTS = set(range(1, 100)) # ============================================================================ # Global Configuration Constants for Molecular Features # ============================================================================ NUM_ATOM_TYPE = 100 # Including extra mask tokens for pre-training NUM_CHIRALITY_TAG = 12 # Chirality information tags NUM_BOND_TYPE = 28 # Including aromatic and self-loop edges, and extra masked tokens NUM_BOND_DIRECTION = 10 # Bond direction types class GATConv(nn.Module): """ Multi-head Graph Attention Convolution Layer with Full-Connect Graph Processing This layer implements a full-connect graph attention mechanism that: 1. Processes messages through local graph edges 2. Applies attention on the full-connect graph 3. Aggregates messages with residual connections 4. Uses layer normalization and activation functions for stability Args: emb_dim (int): Embedding dimension of node features heads (int): Number of attention heads (default: 2) drop_ratio (float): Dropout ratio for regularization (default: 0.0) negative_slope (float): Negative slope for LeakyReLU in attention (default: 0.2) reduce_ratio (int): Channel reduction ratio for dimensionality reduction (default: 16) """ def __init__(self, emb_dim, heads=2, drop_ratio=0.0, negative_slope=0.2, reduce_ratio=16): super(GATConv, self).__init__() # Store configuration parameters self.emb_dim = emb_dim self.heads = heads self.reduce_ratio = reduce_ratio self.negative_slope = negative_slope # Activation function and normalization layers self.act = nn.SiLU() self.norm1 = nn.LayerNorm(emb_dim) # Normalization for local message computation self.norm2 = nn.LayerNorm(emb_dim) # Normalization for full-connect message computation self.norm3 = nn.LayerNorm(emb_dim) # Normalization for attention mechanism self.norm4 = nn.LayerNorm(emb_dim) # Normalization before feed-forward network # Linear projections for local graph message passing self.weight_linear1 = torch.nn.Linear(emb_dim, heads * emb_dim) # Source node projection self.weight_linear2 = torch.nn.Linear(emb_dim, heads * emb_dim) # Target node projection # Dimensionality reduction for full-connect graph processing self.reduce_channel = emb_dim // self.reduce_ratio self.weight_linear3 = torch.nn.Linear(emb_dim, heads * self.reduce_channel) # Source node in FC graph self.weight_linear4 = torch.nn.Linear(emb_dim, heads * self.reduce_channel) # Target node in FC graph self.weight_linear5 = torch.nn.Linear(self.reduce_channel, emb_dim) # Projection back to full dimension # Attention mechanism projections for full-connect graph self.weight_linear_a1 = torch.nn.Linear(emb_dim, heads) # Source attention weights self.weight_linear_a2 = torch.nn.Linear(emb_dim, heads) # Target attention weights self.fc_attn = torch.nn.Linear(heads, heads) # Attention aggregation # Edge attribute embeddings self.edge_embedding1 = torch.nn.Embedding(NUM_BOND_TYPE, heads * emb_dim) # Bond type embedding self.edge_embedding2 = torch.nn.Embedding(5, heads * emb_dim) # Bond stereo embedding self.edge_embedding3 = torch.nn.Embedding(NUM_BOND_DIRECTION, heads * emb_dim) # Bond direction embedding # Feed-forward network for residual update self.FFN = nn.Sequential(nn.Linear(emb_dim, 4 * emb_dim), nn.SiLU(), nn.Linear(emb_dim * 4, emb_dim)) # Scaling factor for full-connect graph contribution self.scale = 0.01 # Dropout layer for regularization self.dropout = torch.nn.Dropout(drop_ratio) # Initialize embedding weights with Xavier uniform distribution torch.nn.init.xavier_uniform_(self.edge_embedding1.weight.data) torch.nn.init.xavier_uniform_(self.edge_embedding2.weight.data) torch.nn.init.xavier_uniform_(self.edge_embedding3.weight.data) def forward(self, x, edge_index, edge_index_all, edge_attr): """ Forward pass of GAT convolution with full-connect graph processing. Args: x (Tensor): Node feature matrix of shape [num_nodes, emb_dim] edge_index (LongTensor): Edge indices for local graph of shape [2, num_edges] edge_index_all (LongTensor): Edge indices for full-connect graph of shape [2, num_edges_full] edge_attr (Tensor): Edge attributes of shape [num_edges, num_attr_features] Returns: Tensor: Updated node representations of shape [num_nodes, emb_dim] """ # Get tensor dimensions num_node, n_dim = x.shape shortcut1 = x # Store input for residual connection # ======================================================================== # Stage 1: Local Graph Message Passing # ======================================================================== x_norm = self.norm1(x) # Add self-loops to the edge index edge_index = add_self_loops(edge_index, num_nodes=x.size(0))[0] # Create edge attributes for self-loop edges self_loop_attr = torch.zeros(x.size(0), 3, device=edge_attr.device, dtype=edge_attr.dtype) self_loop_attr[:, 0] = 25 # Bond type ID 25 is reserved for self-loops edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0) # Embed edge attributes using multiple embedding layers edge_emb = (self.edge_embedding1(edge_attr[:, 0].long()) + self.edge_embedding2(edge_attr[:, 1].long()) + self.edge_embedding3(edge_attr[:, 2].long())) edge_attr_embedded = edge_emb.view(-1, self.heads, self.emb_dim) # Compute source and target node representations for messages x_src = self.weight_linear1(x_norm).view(-1, self.heads, self.emb_dim) x_tgt = self.weight_linear2(x_norm).view(-1, self.heads, self.emb_dim) x_i = x_src[edge_index[0]] # Source nodes for all edges x_j = x_tgt[edge_index[1]] # Target nodes for all edges # Compute messages with edge embeddings and activation message = self.act(x_i + x_j + edge_attr_embedded) # Combine multi-head messages message_combined = message.mean(dim=1) # Aggregate messages to target nodes new_x = torch.zeros_like(x) new_x = new_x.index_add_(0, edge_index[1], message_combined) # ======================================================================== # Stage 2: Full-Connect Graph Processing with Attention # ======================================================================== x_norm2 = self.norm2(x) # Compute dimensionality-reduced representations x_src_fc = self.weight_linear3(x_norm2).view(-1, self.heads, self.reduce_channel)[edge_index_all[0]] x_tgt_fc = self.weight_linear4(x_norm2).view(-1, self.heads, self.reduce_channel)[edge_index_all[1]] # Compute full-connect graph messages fc_message = self.act(x_src_fc + x_tgt_fc) fc_message = self.weight_linear5(fc_message) # ======================================================================== # Stage 3: Attention Mechanism on Full-Connect Graph # ======================================================================== x_norm3 = self.norm3(x) # Compute attention logits attn_src = self.weight_linear_a1(x_norm3)[edge_index_all[0]] attn_tgt = self.weight_linear_a2(x_norm3)[edge_index_all[1]] fc_attn_logits = F.leaky_relu(attn_src + attn_tgt, self.negative_slope) fc_attn_logits = self.fc_attn(fc_attn_logits) # Normalize attention weights using softmax attn_weights = softmax(fc_attn_logits, edge_index_all[1]).view(-1, self.heads, 1) # Apply attention weights to messages fc_message_weighted = fc_message * attn_weights fc_message_final = fc_message_weighted.mean(dim=1) # Aggregate attention-weighted messages to target nodes new_x_fc = torch.zeros(num_node, self.emb_dim, dtype=x.dtype, device=x.device) new_x_fc = new_x_fc.index_add_(0, edge_index_all[1], fc_message_final) # ======================================================================== # Stage 4: Residual Connection and Feed-Forward Network # ======================================================================== # Combine local and full-connect contributions with residual connection new_x = new_x + self.scale * new_x_fc + shortcut1 shortcut2 = new_x # Apply normalization and feed-forward network new_x = self.norm4(new_x) new_x = self.FFN(new_x) new_x = self.dropout(new_x) # Return output with residual connection return new_x + shortcut2 class GNN(torch.nn.Module): """ Graph Neural Network for Molecular Property Prediction A multi-layer GNN architecture that supports both pre-training and fine-tuning modes. In fine-tuning mode, it can leverage pre-trained embeddings from a larger pre-trained model. Args: num_layer (int): Number of GAT layers in the network emb_dim (int): Embedding dimension for node features drop_ratio (float): Dropout rate for regularization (default: 0) output_type (str): Output mode - "last" returns final layer output, "layers" returns outputs from all layers (default: "last") model_mode (str): Training mode - "pretrain" or "finetune" (default: "pretrain") pretrain_emb_dim (int): Embedding dimension of pre-trained model (default: 126) pretrain_num_layer (int): Number of layers in pre-trained model (default: 12) Attributes: gnns (ModuleList): List of GAT convolution layers x_embedding1 (Embedding): Atom type embedding x_embedding2 (Embedding): Atom chirality embedding x_embedding3-5 (Embedding): Additional atom feature embeddings fc_condition (ModuleList): Linear layers for conditioning in fine-tuning mode """ def __init__(self, num_layer, emb_dim, drop_ratio=0, output_type="last", model_mode="pretrain", pretrain_emb_dim=126, pretrain_num_layer=12): super(GNN, self).__init__() # Store configuration parameters self.num_layer = num_layer self.drop_ratio = drop_ratio self.pretrain_num_layer = pretrain_num_layer self.model_mode = model_mode self.output_type = output_type # ======================================================================== # Input Validation # ======================================================================== if self.num_layer < 2: raise ValueError("Number of GNN layers must be greater than or equal to 2.") if self.model_mode not in ["pretrain", "finetune"]: raise ValueError(f"model_mode must be 'pretrain' or 'finetune', got {self.model_mode}") if self.model_mode == "finetune" and self.num_layer < self.pretrain_num_layer: raise ValueError( f"In finetune mode, num_layer ({self.num_layer}) must be >= " f"pretrain_num_layer ({self.pretrain_num_layer})" ) # ======================================================================== # Node Feature Embeddings # ======================================================================== # Primary atom feature embeddings self.x_embedding1 = torch.nn.Embedding(NUM_ATOM_TYPE, emb_dim) # Atom type self.x_embedding2 = torch.nn.Embedding(NUM_CHIRALITY_TAG, emb_dim) # Atom chirality # Auxiliary atom feature embeddings self.x_embedding3 = torch.nn.Embedding(20, emb_dim) # Additional feature 1 self.x_embedding4 = torch.nn.Embedding(10, emb_dim) # Additional feature 2 self.x_embedding5 = torch.nn.Embedding(10, emb_dim) # Additional feature 3 # Feed-forward network for combining auxiliary embeddings self.fc_embedding = nn.Sequential( nn.Linear(emb_dim, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim) ) # Initialize embedding weights with Xavier uniform distribution torch.nn.init.xavier_uniform_(self.x_embedding1.weight.data) torch.nn.init.xavier_uniform_(self.x_embedding2.weight.data) torch.nn.init.xavier_uniform_(self.x_embedding3.weight.data) torch.nn.init.xavier_uniform_(self.x_embedding4.weight.data) torch.nn.init.xavier_uniform_(self.x_embedding5.weight.data) # ======================================================================== # Graph Attention Convolution Layers # ======================================================================== self.gnns = torch.nn.ModuleList() for _ in range(num_layer): self.gnns.append(GATConv(emb_dim, drop_ratio=self.drop_ratio)) # ======================================================================== # Fine-tuning Mode Specific Components # ======================================================================== if self.output_type == "last" and self.model_mode == "finetune": # Linear layers for conditioning the fine-tune model with pre-trained embeddings self.fc_condition = nn.ModuleList() for _ in range(pretrain_num_layer): condition_layer = nn.Linear(pretrain_emb_dim, emb_dim, bias=False) if pretrain_emb_dim != emb_dim else nn.Identity() self.fc_condition.append( condition_layer ) def forward(self, node_atom, edge_index, edge_index_all, edge_attr, batch, extra_embedding=None): """ Forward pass through the GNN layers. Args: node_atom (Tensor): Node atom features of shape [num_nodes, num_features] Expected format: [atom_type, chirality, feature1, feature2, feature3] edge_index (LongTensor): Edge indices for local graph [2, num_edges] edge_index_all (LongTensor): Edge indices for full-connect graph [2, num_edges_full] edge_attr (Tensor): Edge attributes [num_edges, 3] batch (LongTensor): Batch assignment for each node [num_nodes] extra_embedding (list of Tensor, optional): Pre-trained embeddings for fine-tuning. List of layer outputs from pre-trained model. Returns: Tensor or List[Tensor]: - If output_type == "last": Node representations after final layer [num_nodes, emb_dim] - If output_type == "layers": List of node representations from all layers """ # ======================================================================== # Stage 1: Node Feature Embedding # ======================================================================== # Combine primary atom features (atom type + chirality) x = self.x_embedding1(node_atom[:, 0]) + self.x_embedding2(node_atom[:, 1]) # Process auxiliary atom features extra_features = (self.x_embedding3(node_atom[:, 2]) + self.x_embedding4(node_atom[:, 3]) + self.x_embedding5(node_atom[:, 4])) extra_features = self.fc_embedding(extra_features) # Combine all embeddings x = x + extra_features # Initialize list for storing layer outputs if needed if self.output_type == "layers": h = [] # ======================================================================== # Stage 2: Multi-layer Graph Convolution # ======================================================================== for layer_idx in range(self.num_layer): # In fine-tuning mode, condition the network with pre-trained embeddings if self.model_mode == "finetune": if extra_embedding is not None and layer_idx < self.pretrain_num_layer: # Project pre-trained embedding to current dimension and add as residual condition_feature = self.fc_condition[layer_idx](extra_embedding[layer_idx]) x = x + condition_feature # Apply graph attention convolution x = self.gnns[layer_idx](x, edge_index=edge_index, edge_index_all=edge_index_all, edge_attr=edge_attr) # Store layer output if returning all layers if self.output_type == "layers": h.append(x) # ======================================================================== # Return Results # ======================================================================== if self.output_type == "layers": return h return x def get_suiren_arch_pretrained(): return GNN( # these param are preset by suiren team num_layer=12, emb_dim=256, drop_ratio=0.0, output_type="layers" # Return outputs from all layers ) def build_graph(smiles: str) -> Tuple[Optional[Data], Optional[str]]: if pd.isna(smiles): return None, "empty_smiles" smiles = str(smiles).strip() if not smiles: return None, "empty_smiles" graph_tuple, mol_flag = from_smiles(smiles) if not mol_flag: return None, "invalid_smiles" x, edge_index, edge_attr, edge_index_all = graph_tuple atom_types = set(x[:, 0].tolist()) if not atom_types.issubset(ALLOWED_ELEMENTS): return None, "unsupported_elements" data = Data( x=x.to(torch.long), edge_index=edge_index.to(torch.long), edge_attr=edge_attr.to(torch.long), edge_index_all=edge_index_all.to(torch.long), smiles=smiles, ) return data, None