Predicting Drug Toxicity with Tox21 and Chemprop

Introduction

Before a new drug can reach clinical trials, it must be extensively tested for toxicity. Computational prediction of toxicity can save millions of dollars and years of research by filtering out dangerous compounds early. In this tutorial, we will use the Tox21 dataset (a public dataset of 10k compounds and their toxicity profiles across 12 different assays) and Chemprop (a Message Passing Neural Network designed specifically for molecular property prediction) to build a multi-task toxicity classifier.

Architecture Overview

graph TD
    SMILES(SMILES String: CC(=O)OC1=CC=CC=C1C(=O)O) --> Graph[2D Molecular Graph]
    
    subgraph Chemprop D-MPNN
        Graph --> MP1[Message Passing Round 1: Edge Updates]
        MP1 --> MP2[Message Passing Round 2: Edge Updates]
        MP2 --> MP3[Message Passing Round 3: Edge Updates]
        MP3 --> Agg[Readout / Aggregation: Sum of hidden states]
        Agg --> FFN[Feed Forward Neural Network]
    end
    
    FFN --> Out1[p53 Activation Prob]
    FFN --> Out2[Estrogen Receptor Prob]
    FFN --> Out3[... 10 other assays]
        

Prerequisites

• Python 3.8+
• Basic understanding of SMILES strings (Simplified Molecular-Input Line-Entry System)
• A GPU is recommended for faster training, but not strictly required for this dataset size

Step 1: Install Dependencies

Chemprop relies on RDKit for chemical cheminformatics and PyTorch for deep learning.

# Install RDKit via conda (recommended) or pip
conda install -c conda-forge rdkit
pip install chemprop torch

Step 2: Download the Tox21 Dataset

The Tox21 dataset is a standard benchmark in computational chemistry. We'll download the CSV file. It contains SMILES strings representing the molecules, and 12 columns representing different biological toxicity targets (e.g., p53 activation, estrogen receptor binding). The targets are binary (1 for toxic/active, 0 for non-toxic/inactive, and empty for missing data).

wget https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz
gunzip tox21.csv.gz

Step 3: Hyperparameter Optimization and Cross-Validation

Chemprop is designed to be incredibly easy to use via the command line. Before running the final training, it's highly recommended to optimize hyperparameters (like depth, hidden size, and dropout) using Bayesian optimization.

# Run hyperparameter optimization
chemprop_hyperopt \
    --data_path tox21.csv \
    --dataset_type classification \
    --num_iters 20 \
    --config_save_path hyperopts.json

Once we have optimal hyperparameters, we train the model using 5-fold cross-validation to ensure robustness. Chemprop automatically parses the CSV, identifies the SMILES column, recognizes the multi-task targets, and handles masking for missing labels.

chemprop_train \
    --data_path tox21.csv \
    --dataset_type classification \
    --save_dir tox21_checkpoints \
    --epochs 30 \
    --split_type scaffold_balanced \
    --num_folds 5 \
    --ensemble_size 5 \
    --config_path hyperopts.json \
    --metric auc

Note: We use scaffold_balanced splitting instead of random splitting. This ensures that molecules with similar core structures (scaffolds) are kept in the same split, forcing the model to generalize to new, unseen chemical spaces during validation.

Step 4: Making Predictions on New Molecules

Once the ensemble of models is trained across the 5 folds, we can use it to predict the toxicity of new, unseen molecules. Let's create a file called new_drugs.csv with some SMILES strings.

smiles
CC(=O)OC1=CC=CC=C1C(=O)O  # Aspirin
CN1C=NC2=C1C(=O)N(C(=O)N2C)C # Caffeine

Now, run the prediction command. This will output a new CSV file containing the probabilities of toxicity for all 12 assays for each molecule. Because we trained an ensemble of 5 models, Chemprop automatically averages their predictions to give a more reliable result.

chemprop_predict \
    --test_path new_drugs.csv \
    --checkpoint_dir tox21_checkpoints \
    --preds_path toxicity_predictions.csv

Conclusion

Graph Neural Networks, specifically Message Passing Neural Networks like Chemprop, have revolutionized cheminformatics. By treating atoms as nodes and bonds as edges, the network natively understands the 2D topology of the molecule, vastly outperforming traditional methods that rely on hand-crafted Morgan fingerprints. Coupling this with Scaffold splitting ensures our models are truly learning generalizable toxicity rules rather than just memorizing structural datasets.