Bio Applied Molecular Gnn

Graph Neural Networks for Molecular Property Prediction

When to Use

• Predicting a molecular property (blood-brain barrier penetration, solubility, toxicity, binding) directly from SMILES/graph structure rather than hand-crafted descriptors.
• Building or debugging a SMILES → graph (nodes=atoms, edges=bonds) pipeline for PyTorch Geometric.
• Training/evaluating a GCN or MPNN on a MoleculeNet-style dataset (BBBP, ESOL, Tox21, HIV).
• Deciding whether a GNN is worth the extra complexity vs a Morgan-fingerprint + Random Forest baseline.
• Needing a scaffold split (not random split) to get an honest generalization estimate.

Version Compatibility

• Python ≥ 3.10, PyTorch ≥ 2.1, torch-geometric ≥ 2.5, RDKit ≥ 2023.09, scikit-learn ≥ 1.3.

Prerequisites

• pip install torch torch-geometric rdkit scikit-learn numpy
• Familiarity with bio-chemoinformatics-molecular-io (SMILES/RDKit basics) and bio-chemoinformatics-molecular-descriptors (Morgan fingerprints).
• Basic PyTorch (nn.Module, optimizers, training loops).

Molecules as Graphs

Goal: Turn a SMILES string into a PyG Data object usable by any GNN layer. Approach: Atoms become nodes with simple physicochemical features; bonds become edges, duplicated in both directions since RDKit bonds are undirected but PyG expects directed edge pairs.

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


def smiles_to_graph(smiles: str, label: float = None) -> Data:
    """Convert a SMILES string into a PyTorch Geometric graph.

    Node features: [atomic_num, degree, formal_charge, is_aromatic, in_ring].
    Edges are bonds, duplicated (i,j)+(j,i) to make the graph undirected.
    """
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"RDKit could not parse SMILES: {smiles}")

    node_features = [
        [a.GetAtomicNum(), a.GetDegree(), a.GetFormalCharge(),
         int(a.GetIsAromatic()), int(a.IsInRing())]
        for a in mol.GetAtoms()
    ]
    x = torch.tensor(node_features, dtype=torch.float)

    edges = [(b.GetBeginAtomIdx(), b.GetEndAtomIdx()) for b in mol.GetBonds()]
    edges += [(j, i) for i, j in edges]  # undirected
    if edges:
        edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
    else:  # single-atom molecule, no bonds
        edge_index = torch.empty((2, 0), dtype=torch.long)

    data = Data(x=x, edge_index=edge_index)
    if label is not None:
        data.y = torch.tensor([label], dtype=torch.float)
    return data

MPNN Architecture and Training

Goal: A 3-layer GCN with global mean pooling and an MLP head, trained for binary classification (e.g. BBBP permeability). Approach: Stack GCNConv layers to pass messages between bonded atoms, pool node embeddings into one graph-level vector with global_mean_pool, then classify with a small MLP. batch.batch tells the pooling op which nodes belong to which molecule in a mini-batch.

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.nn import GCNConv, global_mean_pool
from sklearn.metrics import roc_auc_score


class MolGNN(nn.Module):
    """3-layer GCN + mean pooling + 2-layer MLP head for graph-level prediction."""

    def __init__(self, in_channels: int = 5, hidden_channels: int = 64, out_channels: int = 1):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)
        self.lin1 = nn.Linear(hidden_channels, hidden_channels)
        self.lin2 = nn.Linear(hidden_channels, out_channels)
        self.dropout = nn.Dropout(0.2)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = self.conv3(x, edge_index).relu()
        x = global_mean_pool(x, batch)          # graph-level readout
        x = self.dropout(F.relu(self.lin1(x)))
        return self.lin2(x)                      # raw logits -> use BCEWithLogitsLoss


def evaluate_auc(model, loader, device="cpu") -> float:
    """Compute ROC-AUC of a trained MolGNN over a DataLoader of PyG graphs."""
    model.eval()
    preds, labels = [], []
    with torch.no_grad():
        for batch in loader:
            batch = batch.to(device)
            out = torch.sigmoid(model(batch.x, batch.edge_index, batch.batch).squeeze(-1))
            preds.extend(out.cpu().numpy())
            labels.extend(batch.y.cpu().numpy())
    return roc_auc_score(labels, preds)


def train_gnn(model, train_loader, val_loader, epochs=50, lr=1e-3, device="cpu"):
    """Train MolGNN with Adam + BCEWithLogitsLoss, printing val AUC every 10 epochs."""
    model = model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.BCEWithLogitsLoss()

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0
        for batch in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()
            out = model(batch.x, batch.edge_index, batch.batch).squeeze(-1)
            loss = criterion(out, batch.y.float())
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * batch.num_graphs

        if epoch % 10 == 0:
            val_auc = evaluate_auc(model, val_loader, device)
            print(f"Epoch {epoch:3d} | train_loss={total_loss/len(train_loader.dataset):.4f} | val_auc={val_auc:.3f}")
    return model

Scaffold Split and Fingerprint Baseline

Goal: Split data by Murcko scaffold (not randomly) and check whether the GNN actually beats a Morgan-fingerprint Random Forest — the standard sanity check before claiming a GNN is worthwhile. Approach: Group molecules by scaffold so structurally similar compounds land in the same split, preventing leakage that inflates test AUC.

import numpy as np
from collections import defaultdict
from rdkit.Chem.Scaffolds import MurckoScaffold
from rdkit.Chem import AllChem
from sklearn.ensemble import RandomForestClassifier


def scaffold_split(smiles_list, frac_train: float = 0.8):
    """Split molecule indices by Murcko scaffold so train/test don't share
    near-identical structures (avoids optimistic AUC from random splits)."""
    scaffold_to_idx = defaultdict(list)
    for i, smi in enumerate(smiles_list):
        mol = Chem.MolFromSmiles(smi)
        scaffold = MurckoScaffold.MurckoScaffoldSmiles(mol=mol, includeChirality=False)
        scaffold_to_idx[scaffold].append(i)

    scaffold_sets = sorted(scaffold_to_idx.values(), key=len, reverse=True)
    n_train = int(frac_train * len(smiles_list))
    train_idx, test_idx = [], []
    for group in scaffold_sets:
        if len(train_idx) + len(group) <= n_train:
            train_idx.extend(group)
        else:
            test_idx.extend(group)
    return train_idx, test_idx


def morgan_rf_auc(train_smiles, train_y, test_smiles, test_y, n_bits=2048, radius=2) -> float:
    """Random Forest on Morgan fingerprints — the baseline any GNN must beat to justify its cost."""
    def fp(smi):
        mol = Chem.MolFromSmiles(smi)
        return np.array(AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits))

    X_train = np.array([fp(s) for s in train_smiles])
    X_test = np.array([fp(s) for s in test_smiles])
    clf = RandomForestClassifier(n_estimators=500, random_state=0, n_jobs=-1)
    clf.fit(X_train, train_y)
    rf_probs = clf.predict_proba(X_test)[:, 1]
    return roc_auc_score(test_y, rf_probs)

Pitfalls

• Random splits overestimate performance: MoleculeNet benchmarks (BBBP, HIV, Tox21) use scaffold splits precisely because random splits let near-duplicate scaffolds leak between train/test and inflate AUC by 0.05–0.15.
• Small datasets favor fingerprints: with a few hundred to low-thousands of molecules, a Morgan-FP Random Forest often matches or beats a GNN — GNNs need more data (or pretraining) to learn useful representations from scratch.
• edge_index dtype and shape: must be torch.long with shape [2, num_edges]; forgetting .t().contiguous() or using float silently breaks message passing or throws opaque CUDA errors.
• Forgetting to duplicate edges: RDKit bonds are undirected but PyG message passing is directed per edge — omit the reverse (j, i) edges and half the graph never receives messages.
• batch.batch is required for pooling: global_mean_pool(x, batch) needs the batch vector from the DataLoader; calling the model on a single un-batched Data object requires batch=torch.zeros(x.size(0), dtype=torch.long).
• Class imbalance: BBBP and similar datasets are often skewed (~75% permeable); report AUC/AUPRC, not accuracy, and consider pos_weight in BCEWithLogitsLoss.

See Also

• bio-chemoinformatics-molecular-descriptors — Morgan fingerprints and classical descriptors used as the GNN baseline.
• bio-chemoinformatics-molecular-io — SMILES parsing and RDKit Mol object basics.
• bio-chemoinformatics-virtual-screening — applying trained property predictors to large compound libraries.
• bio-machine-learning-omics-classifiers — general classifier training/evaluation patterns (CV, metrics) applicable beyond molecules.