alterlab-torchdrug

Files (10)

• SKILL.md
• evals.json · evals/
• core_concepts.md · references/
• datasets.md · references/
• knowledge_graphs.md · references/
• models_architectures.md · references/
• molecular_generation.md · references/
• molecular_property_prediction.md · references/
• protein_modeling.md · references/
• retrosynthesis.md · references/

SKILL.md (first lines)

# TorchDrug

## Overview

TorchDrug is a comprehensive PyTorch-based machine learning toolbox for drug discovery and molecular science. Apply graph neural networks, pre-trained models, and task definitions to molecules, proteins, and biological knowledge graphs, including molecular property prediction, protein modeling, knowledge graph reasoning, molecular generation, retrosynthesis planning, with 40+ curated datasets and 20+ model architectures.

## When to Use This Skill

This skill should be used when working with:

**Data Types:**
- SMILES strings or molecular structures
- Protein sequences or 3D structures (PDB files)
- Chemical reactions and retrosynthesis
- Biomedical knowledge graphs
- Drug discovery datasets

**Tasks:**
- Predicting molecular properties (solubility, toxicity, activity)
- Protein function or structure prediction
- Drug-target binding prediction
- Generating new molecular structures
- Planning chemical synthesis routes
- Link prediction in biomedical knowledge bases
- Training graph neural networks on scientific data

**Libraries and Integration:**
- TorchDrug is the primary library
- Often used with RDKit for cheminformatics
- Compatible with PyTorch and PyTorch Lightning
- Integrates with AlphaFold and ESM for proteins

### Does NOT Trigger

| Scenario | Use Instead |
|----------|-------------|
| Building a general custom GNN today (active library, current Python/torch) | `alterlab-torch-geometric` |
| Ready-made featurizers, MoleculeNet benchmarks, and pretrained models with less code | `alterlab-deepchem` |
| Only turning molecules into feature matrices for your own models | `alterlab-molfeat` |
| Only sourcing a labeled benchmark dataset (ADMET/DTI) with official splits | `alterlab-pytdc` |
| Querying an existing biomedical knowledge graph rather than training a KG model | `alterlab-primekg` |

## Getting Started

### Installation

```bash
# TorchDrug 0.2.1 (last release, Jul 2023) requires Python >=3.7,<3.11 and
# torch >=1.8. It will NOT solve on Python 3.11+ — pin an older interpreter:
uv venv --python 3.10
uv pip install torch                        # install torch first (see below)
# torch-scatter / torch-cluster ship only as sdists, so they compile against the
# installed torch — either install from the PyG wheel index or ensure a C++ toolchain:
uv pip install torch-scatter torch-cluster \
  --find-links "https://data.pyg.org/whl/torch-${TORCH_VERSION}+${CUDA}.html"  # e.g. cpu, cu121
uv pip install torchdrug==0.2.1
```

Gotchas:
- Install `torch` first; TorchDrug and `torch-scatter`/`torch-cluster` build graph ops against the installed torch, and mismatched builds fail at import.
- TorchDrug vendors its own `data.DataLoader`, `data.Graph`, and `core.Engine` — reach for those, not the bare PyTorch equivalents (see the loop below).
- TorchDrug depends on `rdkit-pypi`, which has no wheels past 2022.09 (Python ≤ 3.11) — another reason the stack is pinned to an older interpreter.
- It is unmaintained (no release since 0.2.1, July 2023); for new Python/torch stacks prefer `alterlab-torch-geometric` or `alterlab-deepchem`. Use this skill when you specifically need TorchDrug's task/dataset abstractions (e.g. GearNet, its retrosynthesis pipeline).

### Quick Example

```python
import torch
from torchdrug import data, datasets, models, tasks

# Load molecular dataset. MoleculeNet-style datasets such as BBBP have no .split()
# method (only datasets with predefined splits do) — split explicitly:
dataset = datasets.BBBP("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set = data.ordered_scaffold_split(dataset, lengths)
# (or torch.utils.data.random_split(dataset, lengths) for a random split)

# Define GNN model
model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean"
)

# Create property prediction task
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=["auroc", "auprc"]
)

# Train with a native PyTorch loop.
# preprocess() builds the task's MLP prediction head (and label statistics) — core.Engine
# calls it for you; in a manual loop call it BEFORE creating the optimizer, or the head
# is missing and its parameters are never optimized.
task.preprocess(train_set, valid_set, test_set)
# NOTE: use torchdrug.data.DataLoader (NOT torch.utils.data.DataLoader) — its
# default graph_collate packs data.Graph objects; the stock PyTorch collate cannot.
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = data.DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        # task.forward returns (loss, metric), not a bare tensor.
        loss, metric = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
```

For the standard (non-custom) path, `core.Engine(task, train_set, valid_set, test_set, optimizer, batch_size=...)` wraps this loop and handles batching/devices. Use the manual loop above when you need full control over the training step.

## Core Capabilities

### 1. Molecular Property Prediction

From AlterLab-IEU/AlterLab-Academic-Skills, Apache-2.0.

More from AlterLab-IEU/AlterLab-Academic-Skills