ContextQMD
Back to Claude Scientific Skills

Protein Modeling

skills/torchdrug/references/protein_modeling.md

2.54.06.1 KB
Original Source

Protein Modeling

TorchDrug 0.2.1 documents protein data structures, datasets, sequence encoders, and geometry-aware graph models in its data, dataset, and model APIs. The primary tutorial index focuses on molecular and knowledge-graph workflows, so avoid inventing a protein tutorial API that upstream does not provide.

Build protein objects

From sequence

python
from torchdrug import data

protein = data.Protein.from_sequence(
    "MKTAYIAKQRQISFVKSHFSRQ",
    atom_feature=None,
    bond_feature=None,
    residue_feature="default",
)
print(protein.to_sequence())

For sequence-only work, setting atom and bond features to None avoids the cost of constructing a full atom-level representation.

From PDB

python
protein = data.Protein.from_pdb(
    "protein.pdb",
    atom_feature="default",
    bond_feature="default",
    residue_feature="default",
)

Use trusted local PDB files and validate chain selection, missing residues, alternate locations, and nonstandard residues before training.

Documented conversion methods include:

  • Protein.from_sequence
  • Protein.from_pdb
  • Protein.from_molecule
  • Protein.to_sequence
  • Protein.to_pdb
  • Protein.to_molecule

Packed equivalents operate on lists:

  • PackedProtein.from_sequence(sequences)
  • PackedProtein.from_pdb(pdb_files)
  • PackedProtein.from_molecule(mols)

Protein datasets

Documented dataset families include:

  • Property / sequence: BetaLactamase, BinaryLocalization, SubcellularLocalization
  • Function / structure: EnzymeCommission, GeneOntology, AlphaFoldDB
  • Structure labels: Fold, SecondaryStructure
  • Protein-protein: HumanPPI, YeastPPI, PPIAffinity
  • Protein-ligand: BindingDB, PDBBind

Example:

python
from torchdrug import datasets

dataset = datasets.EnzymeCommission(
    "~/protein-datasets/",
    atom_feature=None,
    bond_feature=None,
    residue_feature="default",
)
train_set, valid_set, test_set = dataset.split()

Class signatures differ. Options such as branch, test_cutoff, lazy, or species/split IDs are dataset-specific; check the API before using them.

Sequence encoders

ESM

models.ESM is the alias for EvolutionaryScaleModeling. The constructor takes a directory for downloaded weights, not a checkpoint filename:

python
from torchdrug import models

model = models.ESM(
    path="~/model-weights/esm/",
    model="ESM-2-150M",
    readout="mean",
)

TorchDrug 0.2.1 supports these ESM-2 names:

  • ESM-2-8M
  • ESM-2-35M
  • ESM-2-150M
  • ESM-2-650M
  • ESM-2-3B
  • ESM-2-15B

It also supports ESM-1b and ESM-1v. Maximum sequence input is 1022 residues before special tokens. Large checkpoints require substantial memory; start with ESM-2-8M or ESM-2-35M for pipeline validation.

Other sequence models

Documented classes include:

  • models.ProteinCNN
  • models.ProteinResNet
  • models.ProteinLSTM
  • models.ProteinBERT

These models require explicit input/hidden dimensions. Derive input dimensions from the dataset's residue feature configuration.

Structure encoders

Documented structure-aware models include:

  • models.GearNet
  • models.SchNet
  • general graph models such as GCN, GAT, GIN, and RGCN

SchNet requires node_position. GearNet requires a graph whose relation and geometric feature configuration matches its constructor.

Use TorchDrug graph-construction and geometry layers to create sequential, radius, and nearest-neighbor relations. Do not use a nonexistent protein.residue_graph(...) method.

Before training a structure model, inspect:

python
print(protein.num_node)
print(protein.num_residue)
print(protein.node_position.shape)
print(protein.residue_feature.shape)

Confirm whether nodes represent atoms or residues and ensure the model input matches that choice.

Property-prediction task

Protein-level classification or regression can use the same task abstraction as molecules:

python
from torchdrug import tasks

task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=("auprc", "auroc"),
)

Choose criterion and metrics from the actual dataset target:

  • binary or multi-label classification: BCE, AUPRC/AUROC
  • multiclass classification: CE and the documented compatible metrics
  • regression: MSE, MAE/RMSE

For large multi-label ontology tasks, inspect tasks.MultipleBinaryClassification rather than treating labels as one multiclass target.

Workflow checks

  1. Decide sequence-only versus structure-aware modeling.
  2. Configure protein features to match that representation.
  3. Verify dataset splits and sequence identity cutoffs.
  4. Check maximum sequence length before selecting ESM.
  5. Build graph relations explicitly for structure models.
  6. Derive dimensions from the loaded dataset.
  7. Smoke-test one batch before long training.
  8. Record checkpoint name, feature settings, split, and TorchDrug version.

Common failures

ESM constructor error

Use models.ESM(path=<directory>, model=<supported-name>). Do not pass a downloaded .pt filename as path.

Out-of-memory error

Choose a smaller ESM model, reduce batch size, crop or filter long sequences, or freeze the encoder and precompute embeddings.

Missing coordinates

Sequence-created proteins do not acquire experimental 3D coordinates. Load a PDB or another validated structure source before using coordinate-dependent models.

Relation mismatch

Build the same relation types expected by the structure model and set num_relation accordingly.