Model and structure integrations#

TMol is designed to sit inside PyTorch-based structural-biology workflows. It can score structures loaded from standard files and convert outputs from structure prediction systems into PoseStack objects.

RoseTTAFold2#

Install TMol into the RoseTTAFold2 environment:

cd <tmol repo root>
pip install -e .

Convert one prediction by passing one-dimensional residue-type indices, three-dimensional atom coordinates, and chain lengths:

from tmol.io import pose_stack_from_rosettafold2

pose_stack = pose_stack_from_rosettafold2(
    seq=sequence_indices,
    xyz=atom_coordinates,
    chainlens=chain_lengths,
)

The adapter supports canonical amino acids and canonical termini. It returns a single-pose stack on the same device as the input tensors. RoseTTAFold2 inference code often disables gradients globally; enable them around any differentiable TMol scoring or minimization:

import torch

from tmol.score import beta2016_score_function

sfxn = beta2016_score_function(pose_stack.device)
scorer = sfxn.render_whole_pose_scoring_module(pose_stack)

with torch.enable_grad():
    coords = pose_stack.coords.detach().clone().requires_grad_(True)
    score = scorer(coords).sum()
    score.backward()

OpenFold#

The OpenFold adapter consumes a result dictionary containing aatype, positions, and chain_index tensors:

from tmol.io import pose_stack_from_openfold

pose_stack = pose_stack_from_openfold(openfold_output)

It supports batched canonical-protein predictions, uses the final entry in positions, and preserves the input device. Additional keys are ignored by the adapter.

Biotite and AtomArray#

The preferred path for rich structure IO is Biotite AtomArray:

import biotite.structure as struc
from biotite.structure.io import load_structure
import torch

from tmol.io import pose_stack_from_biotite
from tmol.score import beta2016_score_function

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
structure = load_structure(
    "complex.cif",
    model=1,
    include_bonds=True,
)
if isinstance(structure, struc.AtomArrayStack):
    structure = structure[0]

pose_stack = pose_stack_from_biotite(structure, device)
sfxn = beta2016_score_function(device)

The score function and PoseStack must use the same device. Batching unlike structures requires compatible chemistry and PoseStackBuilder.from_poses(); TMol does not schedule multi-GPU work internally.

Biotite is especially important for ligands because TMol uses explicit bond tables from the AtomArray during ligand preparation. Follow 07 — Ligands and Parameter Files when the structure contains non-standard residues; the resulting score function must use the ligand-extended parameter database.

Differentiable AtomWorks Atom37 coordinates#

For a model that predicts AtomWorks unified Atom37 coordinates, keep chemical identity and connectivity in its Biotite AtomArray and route only coordinates from the model tensor. Annotate each supported atom with its model token_id and atom37_slot, build the chemistry context once, and reuse it across model steps:

import torch

from tmol.io import (
    build_context_from_biotite,
    prepare_pose_stack_from_atom37,
)
from tmol.score import beta2016_score_function

# atom_array is the AtomWorks-produced topology. atom37_slots is a per-atom
# integer array from the same unified encoding; use -1 for an unmapped atom.
atom_array.set_annotation("atom37_slot", atom37_slots)

context = build_context_from_biotite(
    atom_array,
    atom37_coords.device,
    prepare_ligands=True,
)
pose_builder = prepare_pose_stack_from_atom37(atom_array, context)

# One call accepts one sample or a same-topology batch. The builder can be
# reused at every diffusion, guidance, or search step.
pose_stack = pose_builder(atom37_coords)  # [sample, token, 37, xyz], float32
sfxn = beta2016_score_function(
    pose_stack.device,
    param_db=context.parameter_database,
)
scorer = sfxn.render_whole_pose_scoring_module(pose_stack)
score = scorer(pose_stack.coords).sum()
score.backward()

A single call handles one sample or a whole same-topology batch; hydrogen optimization is enabled by default. For high-throughput diffusion or search, reuse both pose_builder and the scoring module for every compatible topology and batch shape. Pass opt_h=False to the builder only when ideal hydrogen placement is an intentional speed/accuracy tradeoff. For one-off conversion, pose_stack_from_atom37_and_biotite(atom37_coords, atom_array, context) remains available.

A single AtomArray may provide topology for a batch of coordinate tensors. An AtomArrayStack must either have the same number of models as the tensor batch or one model that can be broadcast. Finite mapped tensor coordinates replace the reference coordinates; negative indices, non-finite tensor entries, and unmapped atoms retain their reference coordinates. TMol-generated atoms, such as hydrogens, are left in their built or optimized positions. Gradients from TMol coordinates route back to the mapped Atom37 entries even when hydrogen optimization is enabled. Missing, out-of-range, or ambiguous routing annotations raise Atom37MappingError, so callers can handle mapping failures without catching unrelated pose-construction errors.

This path supports the intersection of the two chemistry systems: canonical protein, DNA, RNA, ordinary prepared ligands, and fragmented ligands on current TMol. Metal-containing ligands and covalently linked modified components require corresponding TMol parameterization support. The adapter has no residue or element allowlist, so those chemistries use this same interface once a context can represent them; strict ligand preparation does not silently drop them in the meantime.