Datasets Reference

Datasets Reference

Overview

TorchDrug provides 40+ curated datasets across multiple domains: molecular property prediction, protein modeling, knowledge graph reasoning, and retrosynthesis. All datasets support lazy loading, automatic downloading, and customizable feature extraction.

Molecular Property Prediction Datasets

Drug Discovery Classification

Key Features:

• All use scaffold splits for realistic evaluation
• Binary classification metrics: AUROC, AUPRC
• Multi-label handles missing values

Use Cases:

• Drug safety prediction
• Virtual screening
• ADMET property prediction

Drug Discovery Regression

Metrics: MAE, RMSE, R² Use Cases: ADME optimization, lead optimization

Quantum Chemistry

Properties (QM9):

• Dipole moment
• Isotropic polarizability
• HOMO/LUMO energies
• Internal energy, enthalpy, free energy
• Heat capacity
• Electronic spatial extent

Use Cases:

• Quantum property prediction
• Method development benchmarking
• Pre-training molecular models

Large Molecule Databases

Protein Datasets

Function Prediction

Features:

• Sequence and/or structure input
• Evolutionary information available
• Multiple train/test splits

Use Cases:

• Protein engineering
• Function annotation
• Enzyme design

Localization and Solubility

Use Cases:

• Protein expression optimization
• Target identification
• Cell biology

Structure Prediction

Use Cases:

• Structure prediction pipelines
• Fold recognition
• Contact map generation

Protein Interactions

Use Cases:

• PPI prediction
• Network biology
• Drug target identification

Protein-Ligand Binding

Use Cases:

• Binding affinity prediction
• Structure-based drug design
• Scoring function development

Large Protein Databases

Knowledge Graph Datasets

General Knowledge

Relation Types (FB15k-237):

• /people/person/nationality
• /film/film/genre
• /location/location/contains
• /business/company/founders
• Many more...

Use Cases:

• Link prediction
• Relation extraction
• Knowledge base completion

Biomedical Knowledge

Entity Types in Hetionet:

• Genes (20,945)
• Compounds (1,552)
• Diseases (137)
• Anatomy (400)
• Pathways (1,822)
• Pharmacologic classes
• Side effects
• Symptoms
• Molecular functions
• Biological processes
• Cellular components

Relation Types:

• Compound-binds-Gene
• Gene-associates-Disease
• Disease-presents-Symptom
• Compound-treats-Disease
• Compound-causes-Side effect
• Gene-participates-Pathway
• And 18 more...

Use Cases:

• Drug repurposing
• Disease mechanism discovery
• Target identification
• Multi-hop reasoning in biomedicine

Citation Network Datasets

Use Cases:

• Node classification
• GNN baseline comparisons
• Method development

Retrosynthesis Datasets

Features:

• Product → Reactants mapping
• Atom mapping for reaction centers
• Canonicalized SMILES
• Balanced across reaction types

Splits:

• Train: ~40,000
• Validation: ~5,000
• Test: ~5,000

Use Cases:

• Retrosynthesis prediction
• Reaction type classification
• Synthetic route planning

Dataset Usage Patterns

Loading Datasets

from torchdrug import datasets

# Basic loading
dataset = datasets.BBBP("~/molecule-datasets/")

# With transforms
from torchdrug import transforms
transform = transforms.VirtualNode()
dataset = datasets.BBBP("~/molecule-datasets/", transform=transform)

# Protein dataset
dataset = datasets.EnzymeCommission("~/protein-datasets/")

# Knowledge graph
dataset = datasets.FB15k237("~/kg-datasets/")

Data Splitting

# Random split
train, valid, test = dataset.split([0.8, 0.1, 0.1])

# Scaffold split (for molecules)
from torchdrug import utils
train, valid, test = dataset.split(
    utils.scaffold_split(dataset, [0.8, 0.1, 0.1])
)

# Predefined splits (some datasets)
train, valid, test = dataset.split()

Feature Extraction

Node Features (Molecules):

• Atom type (one-hot or embedding)
• Formal charge
• Hybridization
• Aromaticity
• Number of hydrogens
• Chirality

Edge Features (Molecules):

• Bond type (single, double, triple, aromatic)
• Stereochemistry
• Conjugation
• Ring membership

Node Features (Proteins):

• Amino acid type (one-hot)
• Physicochemical properties
• Position in sequence
• Secondary structure
• Solvent accessibility

Edge Features (Proteins):

• Edge type (sequential, spatial, contact)
• Distance
• Angles and dihedrals

Choosing Datasets

By Task

Molecular Property Prediction:

• Start with BBBP or HIV (medium size, clear task)
• Use QM9 for quantum properties
• ESOL/FreeSolv for regression

Protein Function:

• EnzymeCommission (well-defined classes)
• GeneOntology (comprehensive annotations)

Drug Safety:

• Tox21 (standard benchmark)
• ClinTox (clinical relevance)

Structure-Based:

• PDBBind (protein-ligand)
• ProteinNet (structure prediction)

Knowledge Graph:

• FB15k-237 (standard benchmark)
• Hetionet (biomedical applications)

Generation:

• ZINC250k (training)
• QM9 (with properties)

Retrosynthesis:

• USPTO-50k (only choice)

By Size and Resources

Small (<5k, for testing):

• BACE, FreeSolv, ClinTox
• Core set of PDBBind

Medium (5k-100k):

• BBBP, HIV, ESOL, Tox21
• EnzymeCommission, Fold
• FB15k-237, WN18RR

Large (>100k):

• QM9, MUV, PCQM4M
• GeneOntology, AlphaFoldDB
• ZINC2M, BindingDB

By Domain

Drug Discovery: BBBP, HIV, Tox21, ESOL, ZINC Quantum Chemistry: QM7, QM8, QM9, PCQM4M Protein Engineering: Fluorescence, Stability, Solubility Structural Biology: Fold, PDBBind, ProteinNet, AlphaFoldDB Biomedical: Hetionet, GeneOntology, EnzymeCommission Retrosynthesis: USPTO-50k

Best Practices

1. Start Small: Test on small datasets before scaling
2. Scaffold Split: Use for realistic drug discovery evaluation
3. Balanced Metrics: Use AUROC + AUPRC for imbalanced data
4. Multiple Runs: Report mean ± std over multiple random seeds
5. Data Leakage: Be careful with pre-trained models
6. Domain Knowledge: Understand what you're predicting
7. Validation: Always use held-out test set
8. Preprocessing: Standardize features, handle missing values