Accelerating Drug Discovery with Graph Neural Networks (GNNs)

Introduction

Traditional drug discovery is a slow, expensive process. By representing molecules as graphs (where atoms are nodes and bonds are edges), we can leverage Graph Neural Networks (GNNs) to predict molecular properties, toxicity, and binding affinities orders of magnitude faster. In this tutorial, we will use PyTorch Geometric and RDKit to train a Graph Convolutional Network (GCN) on molecular data.

Prerequisites

• Python 3.9+
• PyTorch, PyTorch Geometric
• RDKit (for cheminformatics)
• A molecular dataset (e.g., MoleculeNet / ESOL)

Step 1: Parsing Molecules with RDKit

First, we use RDKit to parse SMILES strings (a text representation of molecules) and extract atom features and adjacency matrices.

from rdkit import Chem
import torch
from torch_geometric.data import Data

def smiles_to_graph(smiles, target):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    
    # Node features: Atomic number
    node_features = []
    for atom in mol.GetAtoms():
        node_features.append([atom.GetAtomicNum()])
    x = torch.tensor(node_features, dtype=torch.float)
    
    # Edge indices: Bonds
    edges = []
    for bond in mol.GetBonds():
        i = bond.GetBeginAtomIdx()
        j = bond.GetEndAtomIdx()
        edges.append([i, j])
        edges.append([j, i]) # Undirected graph
        
    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
    y = torch.tensor([[target]], dtype=torch.float)
    
    return Data(x=x, edge_index=edge_index, y=y)

Step 2: Defining the GNN Architecture

We build a Graph Convolutional Network (GCN) using PyTorch Geometric. The GCN layers perform message passing between atoms, and a global pooling layer aggregates the entire molecule into a single vector for prediction.

import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool

class MolecularGCN(torch.nn.Module):
    def __init__(self, num_node_features, hidden_channels):
        super(MolecularGCN, self).__init__()
        self.conv1 = GCNConv(num_node_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)
        self.lin = torch.nn.Linear(hidden_channels, 1)

    def forward(self, x, edge_index, batch):
        # 1. Obtain node embeddings 
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = self.conv2(x, edge_index)
        x = x.relu()
        x = self.conv3(x, edge_index)

        # 2. Readout layer (Pool node embeddings into graph embedding)
        x = global_mean_pool(x, batch)  # [batch_size, hidden_channels]

        # 3. Apply a final classifier
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)
        
        return x

Step 3: Training the GNN

We train the model using standard PyTorch loops, leveraging the DataLoader from PyTorch Geometric which handles batching graphs of variable sizes.

from torch_geometric.loader import DataLoader

# Assuming `graph_dataset` is a list of Data objects created in Step 1
loader = DataLoader(graph_dataset, batch_size=64, shuffle=True)

model = MolecularGCN(num_node_features=1, hidden_channels=64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.MSELoss()

model.train()
for epoch in range(50):
    for data in loader:
        optimizer.zero_grad()
        out = model(data.x, data.edge_index, data.batch)
        loss = criterion(out, data.y)
        loss.backward()
        optimizer.step()

Conclusion

GNNs provide a mathematically elegant way to learn representations directly from molecular structures. While this tutorial focused on a simple GCN predicting a single property, modern architectures like Graph Attention Networks (GATs) or Message Passing Neural Networks (MPNNs) are setting state-of-the-art records in high-throughput virtual screening.