Spaces:
Sleeping
Sleeping
File size: 7,371 Bytes
8f66d04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """
drug_encoder.py — GINE-based molecular graph encoder + Morgan fingerprint fusion.
Uses Graph Isomorphism Network with Edge features (GINE) for learned
molecular representations, combined with Morgan fingerprints (ECFP4)
for knowledge-driven chemical similarity encoding.
v2: adds DropNode/DropEdge graph augmentation for cold-start robustness,
and a MorganFPEncoder for dual-channel drug representation.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GINEConv, BatchNorm
from torch_geometric.data import Batch
class GINELayer(nn.Module):
"""Single GINE layer with batch norm and residual connection."""
def __init__(self, hidden_dim: int, edge_dim: int, dropout: float = 0.2):
super().__init__()
# GINE uses an MLP as the update function
mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim * 2),
nn.ReLU(),
nn.Linear(hidden_dim * 2, hidden_dim),
)
self.conv = GINEConv(nn=mlp, edge_dim=edge_dim)
self.bn = BatchNorm(hidden_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x, edge_index, edge_attr):
h = self.conv(x, edge_index, edge_attr)
h = self.bn(h)
h = F.relu(h)
h = self.dropout(h)
# residual connection
return h + x
class DrugEncoder(nn.Module):
"""
GINE-based drug molecular graph encoder.
Produces per-atom representations suitable for cross-attention
with protein residue representations.
Architecture:
Input projection → N × GINE layers → atom-level output
We intentionally do NOT apply global readout (mean/sum pooling)
here — the cross-attention module needs individual atom vectors
to compute residue-level interaction maps.
"""
def __init__(self,
num_atom_features: int = 52,
edge_dim: int = 16,
hidden_dim: int = 256,
num_layers: int = 3,
dropout: float = 0.2,
jk: str = 'last'):
"""
Args:
num_atom_features: input atom feature dimension
edge_dim: bond feature dimension
hidden_dim: GINE hidden dimension
num_layers: number of GINE layers
dropout: dropout rate
jk: jumping knowledge mode ('last' or 'cat')
"""
super().__init__()
self.num_layers = num_layers
self.jk = jk
# input projection
self.input_proj = nn.Sequential(
nn.Linear(num_atom_features, hidden_dim),
nn.ReLU(),
nn.Dropout(dropout),
)
# edge feature projection (match edge_attr dim to hidden_dim for GINE)
self.edge_proj = nn.Linear(edge_dim, hidden_dim)
# GINE layers
self.layers = nn.ModuleList([
GINELayer(hidden_dim, hidden_dim, dropout)
for _ in range(num_layers)
])
# jumping knowledge output
if jk == 'cat':
self.jk_proj = nn.Linear(hidden_dim * num_layers, hidden_dim)
self.output_dim = hidden_dim
def forward(self, drug_batch: Batch) -> tuple:
"""
Args:
drug_batch: PyG Batch of molecular graphs
Returns:
atom_repr: (total_atoms, hidden_dim) — per-atom representations
batch_index: (total_atoms,) — which graph each atom belongs to
"""
x = drug_batch.x
edge_index = drug_batch.edge_index
edge_attr = drug_batch.edge_attr
batch_index = drug_batch.batch
# project inputs
h = self.input_proj(x)
edge_features = self.edge_proj(edge_attr) if edge_attr.size(0) > 0 \
else edge_attr
# message passing
layer_outputs = []
for layer in self.layers:
h = layer(h, edge_index, edge_features)
layer_outputs.append(h)
# jumping knowledge
if self.jk == 'cat':
h = torch.cat(layer_outputs, dim=-1)
h = self.jk_proj(h)
else:
h = layer_outputs[-1]
return h, batch_index
class GraphAugmentation(nn.Module):
"""
Stochastic graph augmentation for regularisation during training.
DropNode: randomly masks node features (zero-out) with probability p_node.
DropEdge: randomly removes edges with probability p_edge.
These augmentations prevent the GNN from memorising specific molecular
graphs, which is critical for cold-start generalisation where test
molecules are unseen during training.
Reference: Rong et al., "DropEdge: Towards Deep Graph ConvNets", ICLR 2020
"""
def __init__(self, p_node: float = 0.1, p_edge: float = 0.15):
super().__init__()
self.p_node = p_node
self.p_edge = p_edge
def forward(self, drug_batch: Batch) -> Batch:
if not self.training:
return drug_batch
# DropNode: zero-out whole node feature vectors
if self.p_node > 0:
mask = torch.rand(drug_batch.x.size(0), 1,
device=drug_batch.x.device) > self.p_node
drug_batch.x = drug_batch.x * mask.float()
# DropEdge: remove random edges
if self.p_edge > 0 and drug_batch.edge_index.size(1) > 0:
n_edges = drug_batch.edge_index.size(1)
keep = torch.rand(n_edges, device=drug_batch.edge_index.device) > self.p_edge
drug_batch.edge_index = drug_batch.edge_index[:, keep]
if drug_batch.edge_attr is not None and drug_batch.edge_attr.size(0) == n_edges:
drug_batch.edge_attr = drug_batch.edge_attr[keep]
return drug_batch
class MorganFPEncoder(nn.Module):
"""
Encoder for Morgan (ECFP4) molecular fingerprints.
Morgan fingerprints capture circular substructure patterns and provide
a chemistry-prior-based global drug representation that inherently
generalises to unseen molecules — unlike GNN features that depend on
message passing over specific graph topologies.
This encoder projects the binary fingerprint to a dense vector and
is fused with the GNN representation for dual-channel drug encoding.
"""
def __init__(self, input_dim: int = 1024, hidden_dim: int = 256,
dropout: float = 0.2):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim * 2),
nn.LayerNorm(hidden_dim * 2),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim * 2, hidden_dim),
nn.LayerNorm(hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
)
self.output_dim = hidden_dim
def forward(self, fp: torch.Tensor) -> torch.Tensor:
"""
Args:
fp: (B, input_dim) binary Morgan fingerprint
Returns:
(B, hidden_dim) dense drug representation
"""
return self.encoder(fp)
|