Create model_utils.py

model_utils.py

@@ -0,0 +1,128 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import TransformerConv, Linear, HeteroConv
+import torch.nn.functional as F
+
+class GaussianSmearing(nn.Module):
+    """
+    Expands a scalar distance tensor into a 16-dimensional Radial Basis Function (RBF) vector.
+    This creates a rich representation for neural networks out of a single scalar distance.
+    """
+    def __init__(self, start=0.0, stop=8.0, num_gaussians=16):
+        super(GaussianSmearing, self).__init__()
+        offset = torch.linspace(start, stop, num_gaussians)
+        # Calculate the Gaussian coefficient directly handling width scale
+        self.coeff = -0.5 / ((stop - start) / (num_gaussians - 1))**2
+        self.register_buffer('offset', offset)
+
+    def forward(self, dist):
+        # dist shape: [E, 1]
+        # output shape: [E, 16]
+        dist = dist.view(-1, 1) - self.offset.view(1, -1)
+        return torch.exp(self.coeff * torch.pow(dist, 2))
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, hidden_dim, edge_dim=19, dropout=0.1):
+        super(ResidualAttentionBlock, self).__init__()
+        # TransformerConv natively absorbs edge_attr into its attention mechanism.
+        # We split the 128 hidden_dim across 4 attention heads (32 dim per head) to maintain exact tensor shapes.
+        self.conv = TransformerConv(hidden_dim, hidden_dim // 4, heads=4, concat=True, edge_dim=edge_dim)
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x, edge_index, edge_attr):
+        # Handle bipartite graphs (e.g. ['ligand', 'binds', 'protein'])
+        if isinstance(x, tuple):
+            x_src, x_dst = x
+            identity = x_dst
+        else:
+            identity = x
+            
+        out = self.conv(x, edge_index, edge_attr)
+        out = self.norm(out)
+        out = F.relu(out)
+        out = self.dropout(out)
+        return out + identity
+
+class Struct2SeqGNN(nn.Module):
+    def __init__(self, node_features=6, ligand_features=6, hidden_dim=128, num_classes=21, num_layers=4, dropout=0.1):
+        super(Struct2SeqGNN, self).__init__()
+        
+        # Initial node embeddings for distinct node types
+        self.protein_emb = Linear(node_features, hidden_dim)
+        self.ligand_emb = Linear(ligand_features, hidden_dim)
+        
+        # Distinct RBF distance expansions for each edge type
+        # Decoupling these allows the network to scale bonds differently (e.g. covalent backbone vs weak ionic ligand bonds)
+        self.edge_embs = nn.ModuleDict({
+            'protein__interacts_with__protein': GaussianSmearing(start=0.0, stop=8.0, num_gaussians=16),
+            'ligand__binds__protein': GaussianSmearing(start=0.0, stop=8.0, num_gaussians=16),
+            'protein__binds__ligand': GaussianSmearing(start=0.0, stop=8.0, num_gaussians=16)
+        })
+        
+        # Deep module list involving HeteroConv
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            conv = HeteroConv({
+                ('protein', 'interacts_with', 'protein'): ResidualAttentionBlock(hidden_dim, edge_dim=19, dropout=dropout),
+                ('ligand', 'binds', 'protein'): ResidualAttentionBlock(hidden_dim, edge_dim=19, dropout=dropout),
+                ('protein', 'binds', 'ligand'): ResidualAttentionBlock(hidden_dim, edge_dim=19, dropout=dropout)
+            }, aggr='sum')
+            self.layers.append(conv)
+        
+        # Final layer normalization before classification (only applied to protein output)
+        self.norm_out = nn.LayerNorm(hidden_dim)
+        
+        # Sequence prediction: linear classification layer for standard amino acids
+        self.fc = nn.Linear(hidden_dim, num_classes)
+        
+    def forward(self, data):
+        x_dict = data.x_dict
+        edge_index_dict = data.edge_index_dict
+        edge_attr_dict = data.edge_attr_dict
+        
+        # 1. Expand features
+        x_dict_expanded = {
+            'protein': self.protein_emb(x_dict['protein'])
+        }
+        if 'ligand' in x_dict:
+            x_dict_expanded['ligand'] = self.ligand_emb(x_dict['ligand'])
+        x_dict = x_dict_expanded
+        
+        edge_attr_dict_expanded = {}
+        for edge_type, edge_attr in edge_attr_dict.items():
+            src_type, rel_type, dst_type = edge_type
+            
+            # Generate 3D direction vectors symmetrically on-the-fly!
+            # We recalculate the distance dynamically from pos so that it structurally 
+            # matches the vectors even when Gaussian noise is injected into pos during training.
+            src_pos = data[src_type].pos
+            dst_pos = data[dst_type].pos
+            src_idx, dst_idx = edge_index_dict[edge_type]
+            
+            vec = dst_pos[dst_idx] - src_pos[src_idx]
+            dist_new = torch.norm(vec, dim=-1)
+            
+            # Convert raw coordinates into purely directional unit vectors
+            vec_norm = vec / (dist_new.unsqueeze(-1) + 1e-7)
+
+            key = f"{src_type}__{rel_type}__{dst_type}"
+            if key in self.edge_embs:
+                # Expand standard scalar distance to [E, 16] 
+                dist_smeared = self.edge_embs[key](dist_new)
+                
+                # Combine length and direction! Output -> [E, 19]
+                edge_attr_dict_expanded[edge_type] = torch.cat([dist_smeared, vec_norm], dim=-1)
+            else:
+                edge_attr_dict_expanded[edge_type] = edge_attr
+        
+        # 2. Iterative Message Passing through HeteroConv
+        for layer in self.layers:
+            x_dict = layer(x_dict, edge_index_dict, edge_attr_dict=edge_attr_dict_expanded)
+            
+        # 3. Readout on Protein nodes
+        protein_x = x_dict['protein']
+        protein_x = self.norm_out(protein_x)
+        logits = self.fc(protein_x)
+        
+        return logits