### Install TorchDrug from Source

Source: https://torchdrug.ai/docs/installation.html

Clone the repository and install the package from the source code.

```bash
git clone https://github.com/DeepGraphLearning/torchdrug
cd torchdrug
pip install -r requirements.txt
python setup.py install
```

--------------------------------

### Install PyTorch via Pip

Source: https://torchdrug.ai/docs/installation.html

Initial step to install the PyTorch dependency.

```bash
pip install torch
```

--------------------------------

### Registry Usage Examples

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/core.html

Demonstrates how to search for registered objects and register custom hooks.

```python
>>> gcn = R.search("GCN")(128, [128])

>>> @R.register("features.atom.my_feature")
>>> def my_featurizer(atom):
>>>     ...
>>>
>>> data.Molecule.from_smiles("C1=CC=CC=C1", atom_feature="my_feature")
```

--------------------------------

### Install TorchDrug on Apple Silicon

Source: https://torchdrug.ai/docs/installation.html

Install TorchDrug and compile dependencies from source for M1/M2 chips.

```bash
pip install torch==1.13.0
pip install git+https://github.com/rusty1s/pytorch_scatter.git
pip install git+https://github.com/rusty1s/pytorch_cluster.git
pip install torchdrug
```

--------------------------------

### Load and Save Configuration with Configurable

Source: https://torchdrug.ai/docs/api/core.html

Use `config_dict()` to get the configuration of an instance and `load_config_dict()` to create an instance from a configuration. This is convenient for building models from configuration files.

```python
config = models.GCN(128, [128]).config_dict()
gcn = Configurable.load_config_dict(config)

```

--------------------------------

### Sequential Layer Usage Examples

Source: https://torchdrug.ai/docs/api/layers.html

Demonstrates the usage of the improved Sequential container, supporting multiple inputs/outputs, global arguments, and dictionary outputs.

```python
>>> # layer1 signature: (...) -> (a, b)
>>> # layer2 signature: (a, b) -> (...)
>>> layer = layers.Sequential(layer1, layer2)
```

```python
>>> # layer1 signature: (graph, a) -> b
>>> # layer2 signature: (graph, b) -> c
>>> layer = layers.Sequential(layer1, layer2, global_args=("graph",))
```

```python
>>> # layer1 signature: (graph, a) -> b
>>> # layer2 signature: b -> c
>>> # layer3 signature: (graph, c) -> d
>>> layer = layers.Sequential(layer1, layer2, global_args=("graph",))
```

```python
>>> # layer1 signature: a -> {"b": b, "c": c}
>>> # layer2 signature: b -> d
>>> layer = layers.Sequential(layer1, layer2, allow_unused=True)
```

```python
>>> # layer1 signature: (graph, a) -> {"graph": graph, "b": b}
>>> # layer2 signature: (graph, b) -> c
>>> # layer2 takes in the graph output by layer1
>>> layer = layers.Sequential(layer1, layer2, global_args=("graph",))
```

--------------------------------

### Example evaluation output

Source: https://torchdrug.ai/docs/tutorials/property_prediction.html

Sample output format for model evaluation metrics.

```text
auprc [CT_TOX]: 0.455744
auprc [FDA_APPROVED]: 0.985126
auroc [CT_TOX]: 0.861976
auroc [FDA_APPROVED]: 0.816788
```

--------------------------------

### Install PyTorch Geometric dependencies

Source: https://torchdrug.ai/docs/installation.html

Install torch-scatter and torch-cluster by specifying the correct PyTorch and CUDA versions in the URL.

```bash
pip install torch-scatter torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html
```

--------------------------------

### Engine Initialization

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/engine.html

Initializes the Engine class with task, datasets, optimizer, and other training configurations. Supports distributed training setup.

```APIDOC
## Engine Initialization

### Description
Initializes the Engine class with task, datasets, optimizer, and other training configurations. Supports distributed training setup.

### Method
__init__

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **task** (nn.Module) - The task to be trained or tested.
- **train_set** (data.Dataset) - The training dataset.
- **valid_set** (data.Dataset) - The validation dataset.
- **test_set** (data.Dataset) - The test dataset.
- **optimizer** (optim.Optimizer) - The optimizer for training.
- **scheduler** (lr_scheduler._LRScheduler, optional) - The learning rate scheduler.
- **gpus** (list of int, optional) - GPU ids to use. Defaults to CPU.
- **batch_size** (int, optional) - Batch size per GPU/CPU. Defaults to 1.
- **gradient_interval** (int, optional) - Number of batches before performing a gradient update. Defaults to 1.
- **num_worker** (int, optional) - Number of CPU workers per GPU. Defaults to 0.
- **logger** (str or core.LoggerBase, optional) - Logger type or instance. Available types are 'logging' and 'wandb'. Defaults to 'logging'.
- **log_interval** (int, optional) - Log every n gradient updates. Defaults to 100.

### Request Example
```python
engine = Engine(task=my_task, train_set=train_data, valid_set=valid_data, test_set=test_data, optimizer=my_optimizer)
```

### Response
#### Success Response (200)
None (Initialization does not return a value)

#### Response Example
None
```

--------------------------------

### Get World Size

Source: https://torchdrug.ai/docs/api/utils.html

Returns the total number of processes in the distributed group. Returns 1 for a single-process setup.

```python
comm.get_world_size()
```

--------------------------------

### Get Current Process Rank

Source: https://torchdrug.ai/docs/api/utils.html

Returns the rank of the current process within the distributed group. Returns 0 for a single-process setup.

```python
comm.get_rank()
```

--------------------------------

### Initialize Graph and Register Node Attributes

Source: https://torchdrug.ai/docs/api/data.html

Demonstrates creating a graph from an edge list and using a context manager to register dynamic node attributes.

```python
>>> graph = data.Graph(torch.randint(10, (30, 2)))
>>> with graph.node():
>>>     graph.my_node_attr = torch.rand(10, 5, 5)
```

--------------------------------

### Install TorchDrug via Pip

Source: https://torchdrug.ai/docs/installation.html

Final command to install the TorchDrug package.

```bash
pip install torchdrug
```

--------------------------------

### Launch Distributed Training (Multi-Node)

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/engine.html

Use this command for multi-node, multi-process distributed training. Specify the total number of nodes, the rank of the current node, and the number of GPUs per node. Replace placeholders with your specific values.

```bash
python -m torch.distributed.launch --nnodes={number_of_nodes} --node_rank={rank_of_this_node} --nproc_per_node={number_of_gpus} {your_script.py} {your_arguments...}
```

--------------------------------

### Install PowerShell Extensions

Source: https://torchdrug.ai/docs/installation.html

Install required PowerShell modules for JIT compilation on Windows.

```powershell
Install-Module Pscx -AllowClobber
Install-Module VSSetup
```

--------------------------------

### Initialize and Train the Engine

Source: https://torchdrug.ai/docs/quick_start.html

Configures the optimizer and engine to begin the training process.

```python
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(task, train_set, valid_set, test_set, optimizer,
                     batch_size=1024)
solver.train(num_epoch=100)
```

--------------------------------

### Install TorchDrug via Conda

Source: https://torchdrug.ai/docs/installation.html

Use this command to install TorchDrug and its dependencies from the specified channels.

```bash
conda install torchdrug -c milagraph -c conda-forge -c pytorch -c pyg
```

--------------------------------

### Initialize Physicochemical Model

Source: https://torchdrug.ai/docs/_modules/torchdrug/models/physicochemical.html

Initializes the Physicochemical model, downloading necessary property data and setting up the MLP. Ensure the path exists for storing feature files.

```python
import os

import torch
from torch import nn

from torch_scatter import scatter_mean, scatter_add

from torchdrug import core, layers, utils, data
from torchdrug.layers import functional
from torchdrug.core import Registry as R



@R.register("models.Physicochemical")
class Physicochemical(nn.Module, core.Configurable):
    """
    The physicochemical feature engineering for protein sequence proposed in
    `Prediction of Membrane Protein Types based on the Hydrophobic Index of Amino Acids`_.

    .. _Prediction of Membrane Protein Types based on the Hydrophobic Index of Amino Acids:
        https://link.springer.com/article/10.1023/A:1007091128394

    Parameters:
        path (str): path to store feature file
        type (str, optional): physicochemical feature. Available features are ``moran``, ``geary`` and ``nmbroto``.
        nlag (int, optional): maximum position interval to compute features
        hidden_dims (list of int, optional): hidden dimensions
    """

    url = "https://miladeepgraphlearningproteindata.s3.us-east-2.amazonaws.com/documents/AAidx.txt"
    md5 = "ec612f4df41b93ae03c31ae376c23ce0"

    property_key = ["CIDH920105", "BHAR880101", "CHAM820101", "CHAM820102",
                    "CHOC760101", "BIGC670101", "CHAM810101", "DAYM780201"]
    num_residue_type = len(data.Protein.id2residue_symbol)

    def __init__(self, path, type="moran", nlag=30, hidden_dims=(512,)):
        super(Physicochemical, self).__init__()
        self.type = type
        path = os.path.expanduser(path)
        if not os.path.exists(path):
            os.makedirs(path)
        self.path = path
        index_file = utils.download(self.url, path, md5=self.md5)
        property = self.read_property(index_file)
        self.register_buffer("property", property)

        self.nlag = nlag
        self.input_dim = len(self.property_key) * nlag
        self.output_dim = hidden_dims[-1]
        
        self.mlp = layers.Sequential(
            layers.MultiLayerPerceptron(self.input_dim, hidden_dims),
            nn.ReLU()
        )
```

--------------------------------

### Launch Distributed Training (Single-Node)

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/engine.html

Use this command to launch single-node, multi-process distributed training. Replace `{number_of_gpus}` with the actual number of GPUs available and `{your_script.py}` with your training script.

```bash
python -m torch.distributed.launch --nproc_per_node={number_of_gpus} {your_script.py} {your_arguments...}
```

--------------------------------

### Get Protein Sequence

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/protein.html

Retrieves the amino acid sequence of the protein.

```APIDOC
## GET /api/proteins/{protein_id}/to_sequence

### Description
Returns the amino acid sequence of the protein.

### Method
GET

### Endpoint
/api/proteins/{protein_id}/to_sequence

### Parameters
#### Path Parameters
- **protein_id** (str) - Required - The identifier of the protein.

#### Query Parameters
None

### Request Body
None

### Response
#### Success Response (200)
- **str** - The amino acid sequence of the protein.

#### Response Example
```json
{
  "sequence": "MGLSDGEWQLVLNVWGKVE"
}
```
```

--------------------------------

### GET /molecule/scaffold

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/molecule.html

Generates a scaffold SMILES string for the current molecule.

```APIDOC
## GET /molecule/scaffold

### Description
Returns a scaffold SMILES string of the molecule using Murcko Scaffold decomposition.

### Parameters
#### Query Parameters
- **chirality** (bool) - Optional - Whether to consider chirality in the scaffold generation.

### Response
#### Success Response (200)
- **scaffold** (str) - The scaffold SMILES string.
```

--------------------------------

### Load and Split Datasets

Source: https://torchdrug.ai/docs/quick_start.html

Download and prepare the ClinTox dataset, splitting it into training, validation, and test sets.

```python
import torch
from torchdrug import datasets

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)
```

--------------------------------

### Initialize Visual Studio for JIT

Source: https://torchdrug.ai/docs/installation.html

Configure the environment variables for JIT compilation in PowerShell.

```powershell
Import-VisualStudioVars -Architecture x64
$env:LIB += ";C:\Program Files\Python37\libs"
```

--------------------------------

### Get Item by Residue

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/protein.html

Retrieves a masked protein based on residue indexing.

```python
    def __getitem__(self, index):
        # why do we check tuple?
        # case 1: x[0, 1] is parsed as (0, 1)
        # case 2: x[[0, 1]] is parsed as [0, 1]
        if not isinstance(index, tuple):
            index = (index,)

        if len(index) > 1:
            raise ValueError("Protein has only 1 axis, but %d axis is indexed" % len(index))

        return self.residue_mask(index[0], compact=True)
```

--------------------------------

### Get CPU Count

Source: https://torchdrug.ai/docs/_modules/torchdrug/utils/comm.html

Returns the number of available CPUs on the current node.

```python
def get_cpu_count():
    """
    Get the number of CPUs on this node.
    """
    return multiprocessing.cpu_count()
```

--------------------------------

### Create and Query a Tensor Dictionary

Source: https://torchdrug.ai/docs/api/data.html

Demonstrates creating a Dictionary with tensor keys and values, and performing lookups and key checks. Ensure keys and values are PyTorch tensors.

```python
>>> keys = torch.tensor([[0, 0], [1, 1], [2, 2]])
>>> values = torch.tensor([[0, 1], [1, 2], [2, 3]])
>>> d = data.Dictionary(keys, values)
>>> assert (d[[[0, 0], [2, 2]]] == values[[0, 2]]).all()
>>> assert (d.has_key([[0, 1], [1, 2]]) == torch.tensor([False, False])).all()
```

--------------------------------

### Display Evaluation Accuracy

Source: https://torchdrug.ai/docs/tutorials/retrosynthesis.html

Example output format for validation set accuracy.

```text
accuracy: 0.836367
```

--------------------------------

### Initialize Distributed Process Group

Source: https://torchdrug.ai/docs/_modules/torchdrug/utils/comm.html

Sets up CPU and GPU process groups. Use 'nccl' for GPU backends and 'gloo' for CPU.

```python
def init_process_group(backend, init_method=None, **kwargs):
    """
    Initialize CPU and/or GPU process groups.

    Parameters:
        backend (str): Communication backend. Use ``nccl`` for GPUs and ``gloo`` for CPUs.
        init_method (str, optional): URL specifying how to initialize the process group
    """
    global cpu_group
    global gpu_group

    dist.init_process_group(backend, init_method, **kwargs)
    gpu_group = dist.group.WORLD
    if backend == "nccl":
        cpu_group = dist.new_group(backend="gloo")
    else:
        cpu_group = gpu_group
```

--------------------------------

### Sample Actions

Source: https://torchdrug.ai/docs/_modules/torchdrug/tasks/generation.html

Samples actions from the model, supporting both on-policy and off-policy modes.

```python
    @torch.no_grad()
    def _sample_action(self, graph, off_policy):
        if off_policy:
            model = self.agent_model
            new_atom_embeddings = self.agent_new_atom_embeddings
            mlp_stop = self.agent_mlp_stop
            mlp_node1 = self.agent_mlp_node1
            mlp_node2 = self.agent_mlp_node2
            mlp_edge = self.agent_mlp_edge
        else:
            model = self.model
            new_atom_embeddings = self.new_atom_embeddings
            mlp_stop = self.mlp_stop
            mlp_node1 = self.mlp_node1
            mlp_node2 = self.mlp_node2
            mlp_edge = self.mlp_edge

        # step1: get feature
        output = model(graph, graph.node_feature.float())

        extended_node2graph = torch.arange(len(graph), device=self.device).repeat_interleave(len(self.id2atom)) # (num_graph * 16)
        extended_node2graph = torch.cat((graph.node2graph, extended_node2graph)) # (num_node + 16 * num_graph)

        graph_feature_per_node = output["graph_feature"][extended_node2graph]
```

--------------------------------

### GCPN Training Logs

Source: https://torchdrug.ai/docs/tutorials/generation.html

Example output logs generated during the pretraining procedure.

```text
edge acc: 0.896366
edge loss: 0.234644
node1 acc: 0.596209
node1 loss: 1.04997
node2 acc: 0.747235
node2 loss: 0.723717
stop acc: 0.849681
stop bce loss: 0.247942
total loss: 2.25627
```

--------------------------------

### Prepare the ClinTox dataset

Source: https://torchdrug.ai/docs/tutorials/property_prediction.html

Downloads the ClinTox dataset and splits it into training, validation, and test sets.

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

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)
```

--------------------------------

### GET /PackedGraph/match

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/graph.html

Returns all matched indexes for a given pattern, supporting -1 as a wildcard.

```APIDOC
## GET /PackedGraph/match

### Description
Return all matched indexes for each pattern.

### Parameters
#### Query Parameters
- **pattern** (array_like) - Required - Index of shape (N, 2) or (N, 3).

### Response
#### Success Response (200)
- **matched_indexes** (LongTensor) - The indexes of matched edges.
- **num_matches** (LongTensor) - Number of matches per edge.
```

--------------------------------

### USPTO50k Dataset Loading and Initialization

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/uspto50k.html

This snippet details the initialization of the USPTO50k dataset, including downloading the data, loading it from CSV, and processing it either for reaction centers or as synthons.

```APIDOC
## USPTO50k Dataset Initialization

### Description
Initializes the USPTO50k dataset, downloading the data if necessary and processing it to identify reaction centers or decompose reactions into synthons.

### Method
`__init__`

### Parameters
- **path** (str) - Required - Path to store the dataset.
- **as_synthon** (bool, optional) - Whether to decompose (reactant, product) pairs into (reactant, synthon) pairs. Defaults to False.
- **verbose** (int, optional) - Output verbosity level. Defaults to 1.
- **kwargs - Additional keyword arguments passed to `data.ReactionDataset.load_csv`.

### Request Example
```python
dataset = USPTO50k(path="/path/to/data", as_synthon=True)
```

### Response
- **USPTO50k object**: An instance of the USPTO50k dataset.

### Notes
The dataset is downloaded from a specified URL and processed using RDKit and NetworkX. The `target_fields` are aliased for convenience.
```

--------------------------------

### Distributed Training Launch Commands

Source: https://torchdrug.ai/docs/api/core.html

Commands for launching distributed training using `torch.distributed.launch` for single-node and multi-node scenarios.

```bash
python -m torch.distributed.launch --nproc_per_node={number_of_gpus} {your_script.py} {your_arguments...}

```

```bash
python -m torch.distributed.launch --nnodes={number_of_nodes} --node_rank={rank_of_this_node}
--nproc_per_node={number_of_gpus} {your_script.py} {your_arguments...}

```

--------------------------------

### GET /PackedGraph/get_edge

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/graph.html

Retrieves the weight of a specific edge based on the provided edge index.

```APIDOC
## GET /PackedGraph/get_edge

### Description
Get the weight of an edge.

### Parameters
#### Query Parameters
- **edge** (array_like) - Required - Index of shape (2,) or (3,).

### Response
#### Success Response (200)
- **weight** (Tensor) - The weight of the edge.
```

--------------------------------

### Get Process Group by Device

Source: https://torchdrug.ai/docs/api/utils.html

Retrieves the process group associated with a specific device.

```python
comm.get_group(_device_)
```

--------------------------------

### Load Configuration Dictionary

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/engine.html

Constructs an instance of a class from a configuration dictionary. It handles nested configurations and sets up the optimizer with task parameters.

```python
if getattr(cls, "_registry_key", cls.__name__) != config["class"]:
            raise ValueError("Expect config class to be `%s`, but found `%s`" % (cls.__name__, config["class"]))

        optimizer_config = config.pop("optimizer")
        new_config = {}
        for k, v in config.items():
            if isinstance(v, dict) and "class" in v:
                v = core.Configurable.load_config_dict(v)
            if k != "class":
                new_config[k] = v
        optimizer_config["params"] = new_config["task"].parameters()
        new_config["optimizer"] = core.Configurable.load_config_dict(optimizer_config)

        return cls(**new_config)
```

--------------------------------

### CPU Count

Source: https://torchdrug.ai/docs/_modules/torchdrug/utils/comm.html

Utility to get the number of available CPU cores on the current node.

```APIDOC
## GET /comm/cpu_count

### Description
Get the number of CPUs available on the current node.

### Method
GET

### Endpoint
/comm/cpu_count

### Parameters
None

### Response
#### Success Response (200)
- **cpu_count** (int) - The number of CPU cores.

#### Response Example
```json
{
  "cpu_count": 8
}
```
```

--------------------------------

### Fluorescence Dataset Initialization

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/fluorescence.html

This snippet shows how to initialize the Fluorescence dataset, which involves downloading and extracting data, and then loading it into a dataset object.

```APIDOC
## Fluorescence Dataset

### Description
Provides access to the fitness values of green fluorescent protein mutants. It downloads and processes data from a remote URL.

### Method
__init__

### Parameters
#### Path Parameters
- **path** (str) - Required - The path to store the dataset.
- **verbose** (int, optional) - Output verbose level.

#### Request Body
None

### Request Example
```python
from torchdrug.datasets import Fluorescence

# Initialize the dataset
fluorescence_dataset = Fluorescence(path="/path/to/store/data")

# Access dataset splits
train_split, valid_split, test_split = fluorescence_dataset.split()
```

### Response
#### Success Response (200)
- **None** - The `__init__` method does not return a value, but initializes the dataset object.

#### Response Example
None
```

--------------------------------

### Load YeastPPI Dataset

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/yeast_ppi.html

Initializes and loads the YeastPPI dataset. Requires a path for storage and optionally accepts a verbose level. Downloads and extracts data if not present.

```python
import os

from torch.utils import data as torch_data

from torchdrug import data, utils
from torchdrug.core import Registry as R



@R.register("datasets.YeastPPI")
@utils.copy_args(data.ProteinPairDataset.load_lmdbs, ignore=("sequence_field", "target_fields"))
class YeastPPI(data.ProteinPairDataset):
    """
    Binary labels indicating whether two yeast proteins interact or not.

    Statistics:
        - #Train: 1,668
        - #Valid: 131
        - #Test: 373

    Parameters:
        path (str): the path to store the dataset
        verbose (int, optional): output verbose level
        **kwargs
    """

    url = "https://miladeepgraphlearningproteindata.s3.us-east-2.amazonaws.com/ppidata/yeast_ppi.zip"
    md5 = "3993b02c3080d74996cddf6fe798b1e8"
    splits = ["train", "valid", "test", "cross_species_test"]
    target_fields = ["interaction"]

    def __init__(self, path, verbose=1, **kwargs):
        path = os.path.expanduser(path)
        if not os.path.exists(path):
            os.makedirs(path)
        self.path = path

        zip_file = utils.download(self.url, path, md5=self.md5)
        data_path = utils.extract(zip_file)
        lmdb_files = [os.path.join(data_path, "yeast_ppi/yeast_ppi_%s.lmdb" % split)
                      for split in self.splits]

        self.load_lmdbs(lmdb_files, sequence_field=["primary_1", "primary_2"],
                        target_fields=self.target_fields, verbose=verbose, **kwargs)
```

--------------------------------

### GET /to_molecule

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/molecule.html

Converts the internal molecular representation into a list of RDKit molecule objects.

```APIDOC
## GET /to_molecule

### Description
Constructs RDKit molecule objects from internal atom and bond data.

### Parameters
#### Query Parameters
- **ignore_error** (bool) - Optional - If true, return None for illegal molecules instead of raising an exception.

### Response
#### Success Response (200)
- **list of rdchem.Mol** - A list of RDKit molecule objects.
```

--------------------------------

### Initialize WN18 Dataset

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/wn18.html

Initializes the WN18 dataset by downloading and loading the training, validation, and test sets. Ensure the path exists or will be created.

```python
import os

from torch.utils import data as torch_data

from torchdrug import data, utils
from torchdrug.core import Registry as R



@R.register("datasets.WN18")
class WN18(data.KnowledgeGraphDataset):
    """
    WordNet knowledge base.

    Statistics:
        - #Entity: 40,943
        - #Relation: 18
        - #Triplet: 151,442

    Parameters:
        path (str): path to store the dataset
        verbose (int, optional): output verbose level
    """

    urls = [
        "https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding/raw/master/data/wn18/train.txt",
        "https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding/raw/master/data/wn18/valid.txt",
        "https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding/raw/master/data/wn18/test.txt",
    ]
    md5s = [
        "7d68324d293837ac165c3441a6c8b0eb",
        "f4f66fec0ca83b5ebe7ad7003404e61d",
        "b035247a8916c7ec3443fa949e1ff02c"
    ]

    def __init__(self, path, verbose=1):
        path = os.path.expanduser(path)
        if not os.path.exists(path):
            os.makedirs(path)
        self.path = path

        txt_files = []
        for url, md5 in zip(self.urls, self.md5s):
            save_file = "wn18_%s" % os.path.basename(url)
            txt_file = utils.download(url, self.path, save_file=save_file, md5=md5)
            txt_files.append(txt_file)

        self.load_tsvs(txt_files, verbose=verbose)

    def split(self):
        offset = 0
        splits = []
        for num_sample in self.num_samples:
            split = torch_data.Subset(self, range(offset, offset + num_sample))
            splits.append(split)
            offset += num_sample
        return splits

```

--------------------------------

### GET /evaluate

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/engine.html

Evaluates the model on a specified data split and returns the calculated metrics.

```APIDOC
## GET /evaluate

### Description
Evaluate the model on a specific data split (train, valid, or test).

### Parameters
#### Query Parameters
- **split** (str) - Required - Split to evaluate (e.g., 'train', 'valid', 'test').
- **log** (bool) - Optional - Whether to log metrics.

### Response
#### Success Response (200)
- **metrics** (dict) - The calculated evaluation metrics.
```

--------------------------------

### YeastPPI Dataset Initialization

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/yeast_ppi.html

This snippet shows how to initialize the YeastPPI dataset, which involves downloading and extracting the data, and then loading it into a ProteinPairDataset.

```APIDOC
## YeastPPI Dataset

### Description
Provides binary labels indicating whether two yeast proteins interact or not. It includes statistics for training, validation, and testing sets.

### Method
`__init__`

### Endpoint
N/A (This is a dataset class, not an API endpoint)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters
- **path** (str) - Required - The path to store the dataset.
- **verbose** (int, optional) - Output verbose level. Defaults to 1.
- **kwargs** - Additional keyword arguments passed to `data.ProteinPairDataset.load_lmdbs`.

### Request Example
```python
from torchdrug.datasets import YeastPPI

dataset = YeastPPI(path='/path/to/store/dataset')
```

### Response
#### Success Response (200)
N/A (Initialization does not return a response in the typical API sense)

#### Response Example
N/A
```

--------------------------------

### Prepare a batch of graphs

Source: https://torchdrug.ai/docs/notes/variadic.html

Initialize two graphs and pack them into a single batch for processing.

```python
edge_list = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]]
graph1 = data.Graph(edge_list, num_node=6)
edge_list = [[0, 1], [1, 2], [2, 3], [3, 0], [0, 2], [1, 3]]
graph2 = data.Graph(edge_list, num_node=4)
graph = data.Graph.pack([graph1, graph2])
with graph.graph():
    graph.id = torch.arange(2)
```

--------------------------------

### GET /graph/subgraph

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/graph.html

Extracts a subgraph from the current graph based on a provided list of node indices.

```APIDOC
## GET /graph/subgraph

### Description
Returns a subgraph based on the specified nodes. This is equivalent to calling node_mask with compact=True.

### Method
GET

### Endpoint
/graph/subgraph

### Parameters
#### Query Parameters
- **index** (array_like) - Required - The list of node indices to include in the subgraph.

### Response
#### Success Response (200)
- **subgraph** (Graph) - A new Graph object containing only the specified nodes and their associated edges.
```

--------------------------------

### Initialize Forward Pass

Source: https://torchdrug.ai/docs/_modules/torchdrug/tasks/retrosynthesis.html

Sets up the initial loss and metric containers for the forward pass.

```python
def forward(self, batch):
        """"""
        all_loss = torch.tensor(0, dtype=torch.float32, device=self.device)
        metric = {}
```

--------------------------------

### Get Dictionary Device

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/dictionary.html

Returns the device (CPU or CUDA) where the dictionary's keys are located.

```python
return self.keys.device
```

--------------------------------

### Initialize PerfectHash with Keys

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/dictionary.html

Constructs a perfect hash table from a given set of keys. If weight, bias, sub_weights, or sub_biases are not provided, they are randomly generated. The keys can be 1D or 2D tensors.

```python
def __init__(self, keys, weight=None, bias=None, sub_weights=None, sub_biases=None):
        if keys.ndim == 1:
            keys = keys.unsqueeze(-1)
        num_input, input_dim = keys.shape
        if weight is None:
            weight = torch.randint(0, self.prime, (1, input_dim), device=keys.device)
        if bias is None:
            bias = torch.randint(0, self.prime, (1,), device=keys.device)
        if sub_weights is None:
            sub_weights = torch.randint(0, self.prime, (num_input, input_dim), device=keys.device)
        if sub_biases is None:
            sub_biases = torch.randint(0, self.prime, (num_input,), device=keys.device)

        self.keys = keys
        self.weight = weight
        self.bias = bias
        self.sub_weights = sub_weights
        self.sub_biases = sub_biases
        self.num_input = num_input
        self.num_output = num_input
        self.input_dim = input_dim

        self._construct_hash_table()
```

--------------------------------

### InfoGraph Training Output

Source: https://torchdrug.ai/docs/tutorials/pretrain.html

Example output showing the average graph-node mutual information after training.

```text
average graph-node mutual information: 1.30658
```

--------------------------------

### Load and Preprocess ZINC250k Dataset

Source: https://torchdrug.ai/docs/tutorials/generation.html

Downloads and initializes the ZINC250k dataset for molecular graph generation tasks.

```python
import torch
from torchdrug import datasets

dataset = datasets.ZINC250k("~/molecule-datasets/", kekulize=True,
                            atom_feature="symbol")
# with open("path_to_dump/zinc250k.pkl", "wb") as fout:
#     pickle.dump(dataset, fout)
# with open("path_to_dump/zinc250k.pkl", "rb") as fin:
#     dataset = pickle.load(fin)
```

--------------------------------

### GET /get_item

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/protein.html

Retrieves a specific item from the dataset by index, including node, edge, and residue information.

```APIDOC
## GET /get_item

### Description
Retrieves a specific item from the dataset based on the provided index. It masks node, edge, and residue indices to return an unpacked object.

### Method
GET

### Parameters
#### Query Parameters
- **index** (int) - Required - The index of the item to retrieve.

### Response
#### Success Response (200)
- **unpacked_type** (object) - The unpacked molecule or protein object.
```

--------------------------------

### Initialize SecondaryStructure Dataset

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/secondary_structure.html

Initializes the SecondaryStructure dataset by downloading and extracting data, then loading LMDB files. Ensure the path exists or will be created.

```python
import os

import torch
from torch.utils import data as torch_data

from torchdrug import data, utils
from torchdrug.core import Registry as R



@R.register("datasets.SecondaryStructure")
@utils.copy_args(data.ProteinDataset.load_lmdbs, ignore=("target_fields",))
class SecondaryStructure(data.ProteinDataset):
    """
    Secondary structure labels for a set of proteins determined by the local structures 
    of protein residues in their natural state

    Statistics:
        - #Train: 8,678
        - #Valid: 2,170
        - #Test: 513

    Parameters:
        path (str): the path to store the dataset
        verbose (int, optional): output verbose level
        **kwargs
    """

    url = "http://s3.amazonaws.com/songlabdata/proteindata/data_pytorch/secondary_structure.tar.gz"
    md5 = "2f61e8e09c215c032ef5bc8b910c8e97"
    splits = ["train", "valid", "casp12", "ts115", "cb513"]
    target_fields = ["ss3", "valid_mask"]

    def __init__(self, path, verbose=1, **kwargs):
        path = os.path.expanduser(path)
        if not os.path.exists(path):
            os.makedirs(path)
        self.path = path

        zip_file = utils.download(self.url, path, md5=self.md5)
        data_path = utils.extract(zip_file)
        lmdb_files = [os.path.join(data_path, "secondary_structure/secondary_structure_%s.lmdb" % split) 
                                    for split in self.splits]

        self.load_lmdbs(lmdb_files, target_fields=self.target_fields, verbose=verbose, **kwargs)
```

--------------------------------

### Get Distributed World Size

Source: https://torchdrug.ai/docs/_modules/torchdrug/utils/comm.html

Retrieves the total number of processes. Returns 1 if not in a distributed environment.

```python
def get_world_size():
    """
    Get the total number of distributed processes.

    Return 1 for single process case.
    """
    if dist.is_initialized():
        return dist.get_world_size()
    if "WORLD_SIZE" in os.environ:
        return int(os.environ["WORLD_SIZE"])
    return 1
```

--------------------------------

### Get Distributed Process Rank

Source: https://torchdrug.ai/docs/_modules/torchdrug/utils/comm.html

Retrieves the current process rank. Returns 0 if not in a distributed environment.

```python
def get_rank():
    """
    Get the rank of this process in distributed processes.

    Return 0 for single process case.
    """
    if dist.is_initialized():
        return dist.get_rank()
    if "RANK" in os.environ:
        return int(os.environ["RANK"])
    return 0
```

--------------------------------

### Class: FreeSolv

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/freesolv.html

Initializes the FreeSolv dataset, which contains hydration free energy data for 642 molecules.

```APIDOC
## Class: FreeSolv

### Description
Experimental and calculated hydration free energy of small molecules in water.

### Parameters
- **path** (str) - Required - Path to store the dataset.
- **verbose** (int) - Optional - Output verbose level.
- **kwargs** (dict) - Optional - Additional arguments passed to load_csv.

### Statistics
- #Molecule: 642
- #Regression task: 1

### Source
- **URL**: https://s3-us-west-1.amazonaws.com/deepchem.io/datasets/molnet_publish/FreeSolv.zip
- **MD5**: 8d681babd239b15e2f8b2d29f025577a
```

--------------------------------

### Attribute Masking Training Output

Source: https://torchdrug.ai/docs/tutorials/pretrain.html

Example output showing the average accuracy and cross entropy after training.

```text
average accuracy: 0.920366
average cross entropy: 0.22998
```

--------------------------------

### Get Number of Reaction Types

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/uspto50k.html

This property returns the total number of unique reaction types available in the dataset.

```python
@property
    def num_reaction_type(self):
        return len(self.reaction_types)
```

--------------------------------

### QM9 Dataset Initialization

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/qm9.html

Initializes the QM9 dataset, which contains geometric, energetic, electronic, and thermodynamic properties of DFT-modeled small molecules. It can optionally load node positions.

```APIDOC
## QM9 Dataset

### Description
Loads the QM9 dataset, which includes geometric, energetic, electronic, and thermodynamic properties of DFT-modeled small molecules. The dataset contains 133,885 molecules and 12 regression tasks.

### Method
`__init__`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Initialization Parameters
- **path** (str) - Required - The path to store the dataset.
- **node_position** (bool, optional) - If True, loads node positions as a node attribute. Defaults to False.
- **verbose** (int, optional) - The verbosity level for dataset loading. Defaults to 1.

### Request Example
```python
from torchdrug.datasets import QM9

dataset = QM9(path='./data', node_position=True)
```

### Response
#### Success Response (200)
- **data** (list) - A list of Molecule objects.
- **targets** (defaultdict) - A dictionary containing target values for each property (e.g., 'mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'u0', 'u298', 'h298', 'g298').

#### Response Example
```json
{
  "data": [
    // List of Molecule objects
  ],
  "targets": {
    "mu": [value1, value2, ...],
    "alpha": [value1, value2, ...],
    // ... other target fields
  }
}
```
```

--------------------------------

### GET /datasets/opv/download

Source: https://torchdrug.ai/docs/_modules/torchdrug/datasets/opv.html

Endpoints used by the TorchDrug OPV class to retrieve dataset files from the NREL data repository.

```APIDOC
## GET https://cscdata.nrel.gov/api/datasets/ad5d2c9a-af0a-4d72-b943-1e433d5750d6/download/{file_id}

### Description
Retrieves the compressed CSV files containing molecule data for training, validation, or testing.

### Method
GET

### Endpoint
https://cscdata.nrel.gov/api/datasets/ad5d2c9a-af0a-4d72-b943-1e433d5750d6/download/{file_id}

### Parameters
#### Path Parameters
- **file_id** (string) - Required - The unique identifier for the specific dataset split (e.g., b69cf9a5-e7e0-405b-88cb-40df8007242e for training).
```

--------------------------------

### Visualize Graphs in a Figure

Source: https://torchdrug.ai/docs/_modules/torchdrug/data/graph.html

Visualizes a batch of graphs within a matplotlib figure. Ensure matplotlib is installed and configured for plotting.

```python
num_row = math.ceil(self.batch_size / num_col)

        figure_size = (num_col * figure_size[0], num_row * figure_size[1])
        fig = plt.figure(figsize=figure_size)

        for i in range(self.batch_size):
            graph = self.get_item(i)
            ax = fig.add_subplot(num_row, num_col, i + 1)
            graph.visualize(title=titles[i], ax=ax, layout=layout)
        # remove the space of axis labels
        fig.tight_layout()

        if save_file:
            fig.savefig(save_file)
        else:
            fig.show()
```

--------------------------------

### Configurable Initialization Wrapper

Source: https://torchdrug.ai/docs/_modules/torchdrug/core/core.html

Intercepts object initialization to capture configuration arguments.

```python
    def __new__(typ, *args, **kwargs):

        cls = type.__new__(typ, *args, **kwargs)

        @decorator
        def wrapper(init, self, *args, **kwargs):
            sig = inspect.signature(init)
            func = sig.bind(self, *args, **kwargs)
            func.apply_defaults()
            config = {}
            keys = list(sig.parameters.keys())
            for k, v in zip(keys[1:], func.args[1:]): # exclude self
                config[k] = v
            config.update(func.kwargs)
            for k in getattr(self, "_ignore_args", {}):
                config.pop(k)
            self._config = dict(config)
            return init(self, *args, **kwargs)

        def get_function(method):
```