TorchDrug Documentation

TorchDrug provides benchmarks for property prediction models that utilize pre-training. The benchmarks evaluate two primary pre-training methodologies:

1. Self-supervised pre-training: Focuses on learning graph structural information. In these benchmarks, pre-training is performed on a subset of 2 million molecules from the ZINC15 dataset.
2. Supervised pre-training: Involves pre-training on a large supervised dataset. The benchmarks use 456k molecules and 1,310 tasks from the ChEMBL dataset.

Evaluation Methodology:

• Splitting Strategy: Downstream tasks use scaffold splitting for molecule data.
• Data Split: The train/validation/test sets are split into an 80%:10%:10% ratio.
• Metrics: Performance is evaluated using the AUROC metric. Results are reported as the mean and derivation across 10 random splits.

For generative modeling of molecules, use the following task classes:

• AutoregressiveGeneration: Generates molecules step-by-step in an autoregressive manner.
• GCPNGeneration: Implements Graph Convolutional Policy Network (GCPN) based generation.

TorchDrug provides benchmarks for retrosynthesis models evaluated on the USPTO50k dataset. The dataset is partitioned into training, validation, and test sets using an 80%:10%:10% ratio. Performance is measured using top-k accuracy for two distinct experimental settings:

1. Given Reaction Class: The reaction class is provided to the model.
2. Unknown Reaction Class: The reaction class is not provided to the model.

These benchmarks serve as a reference for evaluating the predictive capabilities of retrosynthesis models like G2Gs.

Use the torchdrug.tasks module to perform various property prediction tasks on molecular or graph data. Available task classes include:

• PropertyPrediction: For predicting continuous or discrete properties of a molecule.
• MultipleBinaryClassification: For multi-label classification tasks.
• NodePropertyPrediction: For predicting properties at the node level.
• InteractionPrediction: For predicting interactions between entities.

TorchDrug provides benchmarks for graph generative models focused on goal-directed property optimization. The goal is to generate novel molecules with optimized chemical properties.

Benchmark Methodology:

1. Pretraining: Models are pretrained on the ZINC250k dataset.
2. Finetuning: Reinforcement learning algorithms are applied to finetune the networks toward specific target properties.

Target Properties Used:

• Penalized logP score: The octanol-water partition coefficient, penalized by the synthetic accessibility score and the number of long cycles.
• QED score: A measure of the drug-likeness of the molecule.

Constraints: The maximum graph size for these benchmarks is set to 38, matching the maximum graph size of molecules in the ZINC250k dataset.

TorchDrug supports tasks related to chemical retrosynthesis and reaction analysis:

• CenterIdentification: Identifying the reaction center in a molecule.
• SynthonCompletion: Completing synthons for reaction prediction.
• Retrosynthesis: The core task for predicting precursors from a target molecule.

TorchDrug provides several task classes designed for self-supervised or pre-training molecular representations. These tasks focus on reconstructing or predicting structural features:

• EdgePrediction: Predicting the existence or type of edges.
• AttributeMasking: Reconstructing masked node or edge attributes.
• ContextPrediction: Predicting local context within a graph.
• DistancePrediction: Predicting inter-atomic distances.
• AnglePrediction: Predicting bond angles.
• DihedralPrediction: Predicting dihedral angles.
• Unsupervised: General unsupervised learning tasks.

Data processing in TorchDrug is handled via the DataLoader and several specialized splitting functions to ensure proper data distribution during training.

Data Loading:

• DataLoader: Used to iterate over datasets in batches.

Dataset Splitting Methods:

• graph_collate: Function for collating graphs into batches.
• key_split: Splits data based on specific keys.
• scaffold_split: Splits data based on molecular scaffolds.
• ordered_scaffold_split: An ordered version of scaffold splitting.
• semisupervised: Splits data for semi-supervised learning scenarios.

TorchDrug includes specialized modules for various drug discovery workflows:

Pretraining Molecular Representations:

• InfoGraph: Unsupervised/Semi-supervised graph-level representation learning.
• Pretraining tasks: EdgePrediction, AttributeMasking, and ContextPrediction.

De Novo Molecule Design:

• GCPNGeneration: Goal-directed molecular graph generation.
• GraphAutoregressiveFlow: Flow-based autoregressive model.
• AutoregressiveGeneration: Task for autoregressive generation.

Retrosynthesis:

• Tasks: CenterIdentification, SynthonCompletion, and Retrosynthesis.

TorchDrug follows a convention of separating representation models from task-specific designs to improve reusability.

1. Representation Models: These are subclasses of nn.Module and core.Configurable. They focus on learning latent features (e.g., node, edge, or graph representations) and should return a dictionary containing these features.
2. Tasks: These are subclasses of tasks.Task and core.Configurable. They define how the model's representations are used to compute losses and metrics (e.g., link prediction, node classification).

By separating these, you can use the same representation model across different tasks or different tasks with the same model architecture.