torchdrug

TorchDrug

概要

TorchDrug は、創薬と分子科学のための包括的な PyTorch ベースの機械学習ツールボックスです。グラフニューラルネットワーク、事前学習済みモデル、タスク定義を分子、タンパク質、生物学的ナレッジグラフに適用できます。分子特性予測、タンパク質モデリング、ナレッジグラフ推論、分子生成、逆合成計画を含み、40 を超えるキュレーション済みデータセットと 20 を超えるモデルアーキテクチャを備えています。

このスキルを使用するタイミング

このスキルは、次の作業を行う場合に使用します。

データの種類:

• SMILES 文字列または分子構造
• タンパク質配列または 3D 構造（PDB ファイル）
• 化学反応と逆合成
• バイオメディカルナレッジグラフ
• 創薬データセット

タスク:

• 分子特性（溶解度、毒性、活性）の予測
• タンパク質の機能または構造の予測
• 薬物と標的の結合予測
• 新規分子構造の生成
• 化学合成経路の計画
• バイオメディカルナレッジベースにおけるリンク予測
• 科学データを用いたグラフニューラルネットワークの学習

ライブラリと統合:

• TorchDrug は主要なライブラリです
• ケモインフォマティクスでは RDKit と併用されることが多いです
• PyTorch および PyTorch Lightning と互換性があります
• タンパク質については AlphaFold および ESM と統合されています

はじめに

インストール

uv pip install torchdrug
# Or with optional dependencies
uv pip install torchdrug[full]

クイック例

from torchdrug import datasets, models, tasks
from torch.utils.data import DataLoader

# Load molecular dataset
dataset = datasets.BBBP("~/molecule-datasets/")
train_set, valid_set, test_set = dataset.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 PyTorch
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        loss = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

主な機能

1. 分子特性予測

構造から分子の化学的、物理的、生物学的特性を予測します。

ユースケース:

• 薬らしさ（drug-likeness）および ADMET 特性
• 毒性スクリーニング
• 量子化学的特性
• 結合親和性予測

主要コンポーネント:

• 20 を超える分子データセット（BBBP、HIV、Tox21、QM9 など）
• GNN モデル（GIN、GAT、SchNet）
• PropertyPrediction および MultipleBinaryClassification タスク

リファレンス: 以下の内容については references/molecular_property_prediction.md を参照してください。

• データセットの完全なカタログ
• モデル選択ガイド
• 学習ワークフローとベストプラクティス
• 特徴量エンジニアリングの詳細

2. タンパク質モデリング

タンパク質の配列、構造、特性を扱います。

ユースケース:

• 酵素機能予測
• タンパク質の安定性と溶解度
• 細胞内局在
• タンパク質間相互作用
• 構造予測

主要コンポーネント:

• 15 を超えるタンパク質データセット（EnzymeCommission、GeneOntology、PDBBind など）
• 配列モデル（ESM、ProteinBERT、ProteinLSTM）
• 構造モデル（GearNet、SchNet）
• 予測レベルの異なる複数のタスクタイプ

リファレンス: 以下の内容については references/protein_modeling.md を参照してください。

• タンパク質固有のデータセット
• 配列モデルと構造モデルの比較
• 事前学習戦略
• AlphaFold および ESM との統合

3. ナレッジグラフ推論

生物学的ナレッジグラフにおける欠落したリンクと関係を予測します。

ユースケース:

• 薬剤の再利用（ドラッグリポジショニング）
• 疾患メカニズムの発見
• 遺伝子と疾患の関連性
• 多段階のバイオメディカル推論

主要コンポーネント:

• 一般的な KG（FB15k、WN18）およびバイオメディカル（Hetionet）
• 埋め込みモデル（TransE、RotatE、ComplEx）
• KnowledgeGraphCompletion タスク

リファレンス: 以下の内容については references/knowledge_graphs.md を参照してください。

• ナレッジグラフのデータセット（4.5 万のバイオメディカルエンティティを含む Hetionet を含む）
• 埋め込みモデルの比較
• 評価指標とプロトコル
• バイオメディカルへの応用

4. 分子生成

目的の特性を持つ新規分子構造を生成します。

ユースケース:

• デノボ創薬
• リード最適化
• 化学空間の探索
• 特性に基づく生成

主要コンポーネント:

• 自己回帰生成
• GCPN（方策ベースの生成）
• GraphAutoregressiveFlow
• 特性最適化ワークフロー

リファレンス: 以下の内容については references/molecular_generation.md を参照してください。

• 生成戦略（無条件、条件付き、スキャフォールドベース）
• 多目的最適化
• 検証とフィルタリング
• 特性予測との統合

5. レトロシンセシス

標的分子から出発原料までの合成経路を予測します。

ユースケース:

• 合成計画
• 経路の最適化
• 合成容易性の評価
• 多段階計画

主要コンポーネント:

• USPTO-50k 反応データセット
• CenterIdentification（反応中心の予測）
• SynthonCompletion（反応物の予測）
• エンドツーエンドのレトロシンセシスパイプライン

リファレンス: 以下の内容については references/retrosynthesis.md を参照してください:

• タスクの分解（中心の特定 → シントン補完）
• 多段階合成計画
• 市販入手可能性の確認
• 他のレトロシンセシスツールとの統合

6. グラフニューラルネットワークモデル

データの種類とタスクに応じた GNN アーキテクチャの包括的なカタログです。

利用可能なモデル:

• 汎用 GNN: GCN、GAT、GIN、RGCN、MPNN
• 3D 対応: SchNet、GearNet
• タンパク質特化: ESM、ProteinBERT、GearNet
• ナレッジグラフ: TransE、RotatE、ComplEx、SimplE
• 生成系: GraphAutoregressiveFlow

リファレンス: 以下の内容については references/models_architectures.md を参照してください:

• モデルの詳細な説明
• タスクとデータセットに基づくモデル選択ガイド
• アーキテクチャの比較
• 実装のヒント

7. データセット

化学、生物学、ナレッジグラフにまたがる 40 種類以上のキュレーション済みデータセットです。

カテゴリ:

• 分子特性（創薬、量子化学）
• タンパク質特性（機能、構造、相互作用）
• ナレッジグラフ（汎用および生物医学）
• レトロシンセシス反応

リファレンス: 以下の内容については references/datasets.md を参照してください:

• サイズとタスクを含む完全なデータセットカタログ
• データセット選択ガイド
• 読み込みと前処理
• 分割戦略（ランダム、スカフォールド）

一般的なワークフロー

ワークフロー 1: 分子特性予測

シナリオ: 創薬候補の血液脳関門透過性を予測します。

手順:

1. データセットを読み込む: datasets.BBBP()
2. モデルを選択する: 分子グラフ向けの GIN
3. タスクを定義する: 二値分類の PropertyPrediction
4. 現実的な評価のためにスカフォールド分割で学習する
5. AUROC と AUPRC を用いて評価する

ナビゲーション: references/molecular_property_prediction.md → データセットの選択 → モデルの選択 → 学習

ワークフロー 2: タンパク質機能予測

シナリオ: 配列から酵素機能を予測します。

手順:

1. データセットを読み込む: datasets.EnzymeCommission()
2. モデルを選択する: ESM（事前学習済み）または GearNet（構造情報を使用）
3. タスクを定義する: 多クラス分類の PropertyPrediction
4. 事前学習済みモデルをファインチューニングするか、ゼロから学習する
5. 正解率とクラスごとの指標を用いて評価する

ナビゲーション: references/protein_modeling.md → モデルの選択（配列 vs 構造）→ 事前学習戦略

ワークフロー 3: ナレッジグラフによる薬剤再利用

シナリオ: Hetionet で新たな疾患の治療法を探索します。

手順:

1. データセットを読み込む: datasets.Hetionet()
2. モデルを選択する: RotatE または ComplEx
3. タスクを定義する: KnowledgeGraphCompletion
4. 負例サンプリングを用いて学習する
5. "Compound-treats-Disease" の予測をクエリする
6. 妥当性と作用機序でフィルタリングする

ナビゲーション: references/knowledge_graphs.md → Hetionet データセット → モデルの選択 → 生物医学への応用

ワークフロー 4: デノボ分子生成

シナリオ: 標的への結合に最適化された創薬様分子を生成します。

手順:

1. 活性データで特性予測器を学習する
2. 生成アプローチを選択する: 強化学習ベースの最適化には GCPN
3. 親和性、創薬様性、合成容易性を組み合わせた報酬関数を定義する
4. 特性の制約を付けて候補を生成する
5. 化学的妥当性を検証し、創薬様性でフィルタリングする
6. 多目的スコアリングでランク付けする

ナビゲーション: references/molecular_generation.md → 条件付き生成 → 多目的最適化

ワークフロー 5: レトロシンセシス計画

シナリオ: 対象分子の合成ルートを計画する。

手順:

1. データセットを読み込む: datasets.USPTO50k()
2. 中心原子特定モデル（RGCN）を学習する
3. シントン補完モデル（GIN）を学習する
4. エンドツーエンドのレトロシンセシスパイプラインに統合する
5. 多段階計画のために再帰的に適用する
6. 構成単位の市販入手可能性を確認する

ナビゲーション: references/retrosynthesis.md → Task types → Multi-step planning

統合パターン

RDKit との連携

TorchDrug の分子と RDKit の間で変換する:

from torchdrug import data
from rdkit import Chem

# SMILES → TorchDrug molecule
smiles = "CCO"
mol = data.Molecule.from_smiles(smiles)

# TorchDrug → RDKit
rdkit_mol = mol.to_molecule()

# RDKit → TorchDrug
rdkit_mol = Chem.MolFromSmiles(smiles)
mol = data.Molecule.from_molecule(rdkit_mol)

AlphaFold/ESM との連携

予測された構造を使用する:

from torchdrug import data

# Load AlphaFold predicted structure
protein = data.Protein.from_pdb("AF-P12345-F1-model_v4.pdb")

# Build graph with spatial edges
graph = protein.residue_graph(
    node_position="ca",
    edge_types=["sequential", "radius"],
    radius_cutoff=10.0
)

PyTorch Lightning との連携

タスクを Lightning での学習用にラップする:

import pytorch_lightning as pl

class LightningTask(pl.LightningModule):
    def __init__(self, torchdrug_task):
        super().__init__()
        self.task = torchdrug_task

    def training_step(self, batch, batch_idx):
        return self.task(batch)

    def validation_step(self, batch, batch_idx):
        pred = self.task.predict(batch)
        target = self.task.target(batch)
        return {"pred": pred, "target": target}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

技術的詳細

TorchDrug のアーキテクチャを詳しく知るには:

コアコンセプト: 以下については references/core_concepts.md を参照してください:

• アーキテクチャの設計思想（モジュール式、設定可能）
• データ構造（Graph、Molecule、Protein、PackedGraph）
• モデルのインターフェースと forward 関数のシグネチャ
• タスクのインターフェース（predict、target、forward、evaluate）
• 学習ワークフローとベストプラクティス
• 損失関数と評価指標
• よくある落とし穴とデバッグ

クイックリファレンス チートシート

データセットの選択:

• 分子特性 → references/datasets.md → Molecular section
• タンパク質タスク → references/datasets.md → Protein section
• ナレッジグラフ → references/datasets.md → Knowledge graph section

モデルの選択:

• 分子 → references/models_architectures.md → GNN section → GIN/GAT/SchNet
• タンパク質（配列） → references/models_architectures.md → Protein section → ESM
• タンパク質（構造） → references/models_architectures.md → Protein section → GearNet
• ナレッジグラフ → references/models_architectures.md → KG section → RotatE/ComplEx

一般的なタスク:

• 特性予測 → references/molecular_property_prediction.md または references/protein_modeling.md
• 生成 → references/molecular_generation.md
• レトロシンセシス → references/retrosynthesis.md
• ナレッジグラフ推論 → references/knowledge_graphs.md

アーキテクチャの理解:

• データ構造 → references/core_concepts.md → Data Structures
• モデル設計 → references/core_concepts.md → Model Interface
• タスク設計 → references/core_concepts.md → Task Interface

よくある問題のトラブルシューティング

問題: 次元不一致エラー → model.input_dim が dataset.node_feature_dim と一致しているか確認する → references/core_concepts.md → Essential Attributes を参照

問題: 分子タスクで性能が低い → ランダム分割ではなく scaffold 分割を使用する → GCN ではなく GIN を試す → references/molecular_property_prediction.md → Best Practices を参照

問題: タンパク質モデルが学習しない → 配列タスクには事前学習済みの ESM を使用する → 構造モデルの場合はエッジの構築を確認する → references/protein_modeling.md → Training Workflows を参照

問題: 大規模グラフでメモリエラーが発生する → バッチサイズを小さくする → 勾配累積を使用する → references/core_concepts.md → Memory Efficiency を参照

問題: 生成された分子が無効である → 妥当性の制約を追加する → RDKit による妥当性検証で後処理する → references/molecular_generation.md → Validation and Filtering を参照

リソース

公式ドキュメント: https://torchdrug.ai/docs/ GitHub: https://github.com/DeepGraphLearning/torchdrug 論文: TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery

まとめ

タスクに応じて、適切なリファレンスファイルに移動してください:

1. 分子特性予測 → molecular_property_prediction.md
2. タンパク質モデリング → protein_modeling.md
3. ナレッジグラフ → knowledge_graphs.md
4. 分子生成 → molecular_generation.md
5. レトロシンセシス → retrosynthesis.md
6. モデルの選択 → models_architectures.md
7. データセットの選択 → datasets.md
8. 技術的詳細 → core_concepts.md

各リファレンスでは、例、ベストプラクティス、一般的なユースケースとともに、その分野を包括的に解説しています。