Build your own model input interface#
Map OpenFold-style atom14 and RoseTTAFold2 prediction tensors to TMol, score them, and differentiate back to the original coordinates. The adapter functions below are editable application code using TMol’s canonical construction and scoring APIs.
Run cells in order. In Colab, select a GPU runtime; the shared setup installs published wheels matching PyTorch’s CUDA major and downloads example predictions. No compilation is needed. Locally, follow the installation guide, then run from a source checkout for the fixtures. CPU and CUDA are supported. Set TMOL_EXAMPLE_SOURCE to another fixture checkout or TMOL_EXAMPLE_DEVICE to select a device explicitly.
[ ]:
import os
from pathlib import Path
try:
import google.colab # noqa: F401
except ImportError:
IN_COLAB = False
else:
IN_COLAB = True
if IN_COLAB:
from urllib.request import urlretrieve
urlretrieve(
"https://raw.githubusercontent.com/RosettaCommons/tmol/"
"main/docs/tutorial/notebook_setup.py",
"notebook_setup.py",
)
import notebook_setup
notebook_setup.setup_colab(
[
"tmol/tests/data/openfold/openfold_ubq_and_sumo.pt",
"tmol/tests/data/rosettafold2/ubiquitin.pt",
"tmol/tests/data/rosettafold2/input_layout.json",
]
)
repo = Path(os.environ.get("TMOL_EXAMPLE_SOURCE", ".")).resolve()
os.environ.setdefault("OMP_NUM_THREADS", "2")
1. Describe a coordinate layout once#
A slot count does not specify atom order or residue-token numbering. Supply ordered residue names and named slots for each residue. AtomWorks supplies common atom14/atom37 definitions; check the token order against your predictor version.
This builder takes [batch, residues, atoms, 3] float32 tensors. chain_id=-1 means a padded residue; NaN triplets or an observation mask mean missing atoms of a real residue. Torch indexed assignment preserves gradients from supplied coordinates. The example caches the mappings, while canonical pose construction still runs for each call. It expects the heavy atoms required for construction; it does not implement missing-sidechain packing.
[ ]:
import torch
import tmol
from tmol.io import (
CanonicalForm,
default_canonical_ordering,
default_packed_block_types,
pose_stack_from_canonical_form,
)
from tmol.score import beta2016_score_function
device = torch.device(
os.environ.get(
"TMOL_EXAMPLE_DEVICE",
"cuda" if tmol.has_cuda() and torch.cuda.is_available() else "cpu",
)
)
print("PyTorch:", torch.__version__, "device:", device)
def prepare_named_layout(sequence, chain_id, residue_names, atom_names):
"""Prepare a model's fixed sequence/layout once; bind coordinates in Torch.
This example uses the default protein database. General chemistry callers
should use the ordering and packed types from their PoseBuildContext.
Padding is explicitly identified by chain_id == -1.
"""
if sequence.ndim != 2 or chain_id.shape != sequence.shape:
raise ValueError("sequence and chain_id must have shape [batch, residues]")
if sequence.dtype not in (torch.int32, torch.int64):
raise TypeError("sequence must contain integer token IDs")
if chain_id.dtype not in (torch.int32, torch.int64):
raise TypeError("chain_id must contain integer chain IDs")
if sequence.device != chain_id.device or torch.any(chain_id < -1):
raise ValueError(
"chain IDs must share the sequence device and use -1 for padding"
)
present = chain_id >= 0
if torch.any(present & ((sequence < 0) | (sequence >= len(residue_names)))):
raise ValueError("An observed residue has an unknown token ID")
sequence = sequence.clone().long()
sequence[~present] = 0
chain_id = chain_id.clone().to(torch.int32)
ordering = default_canonical_ordering()
packed_types = default_packed_block_types(sequence.device)
restype_map, atom_map, real_atoms = ordering.create_src_2_tmol_mappings(
residue_names, atom_names, sequence.device
)
restypes = restype_map[sequence].to(torch.int32)
if torch.any(present & (restypes < 0)):
raise ValueError("A residue needs a chemical definition before construction")
restypes[~present] = -1
real = real_atoms[sequence] & present.unsqueeze(-1)
pose, residue, slot = torch.nonzero(real, as_tuple=True)
target_atom = atom_map[sequence][real]
shape = (*sequence.shape, real.shape[-1], 3)
def build(coords, *, observed=None):
if tuple(coords.shape) != shape or coords.dtype != torch.float32:
raise ValueError(f"Expected float32 coordinates shaped {shape}")
if coords.device != sequence.device:
raise ValueError("Coordinates and prepared mapping must share a device")
source = coords[pose, residue, slot]
if observed is not None:
if observed.shape != coords.shape[:-1] or observed.dtype != torch.bool:
raise ValueError(
"observed must be a boolean mask over coordinate slots"
)
source = torch.where(
observed[pose, residue, slot, None], source, float("nan")
)
valid = torch.isfinite(source).all(-1) | torch.isnan(source).all(-1)
if not bool(valid.all()):
raise ValueError(
"Supply finite coordinate triplets or all-NaN missing atoms"
)
canonical = coords.new_full(
(*sequence.shape, ordering.max_n_canonical_atoms, 3), float("nan")
)
canonical[pose, residue, target_atom] = source
form = CanonicalForm(
chain_id=chain_id,
res_types=restypes,
coords=canonical,
chain_labels=None,
res_labels=None,
residue_insertion_codes=None,
atom_occupancy=None,
atom_b_factor=None,
disulfides=None,
res_not_connected=None,
)
return pose_stack_from_canonical_form(ordering, packed_types, *form)
return build
2. OpenFold-style atom14 output#
Read the final structure-module coordinates, sequence tokens, chain IDs and atom-existence mask from a saved prediction. These dictionary keys describe this example’s output format. Adapt them to your predictor; no OpenFold package or tmol-specific OpenFold interface is needed.
[ ]:
def openfold_example(prediction):
"""Adapt this predictor's token IDs, final atom14 coordinates and chain mask."""
from atomworks.ml.encoding_definitions import AF2_ATOM14_ENCODING
# This token order is part of the predictor's contract, independent of shape.
names = list(AF2_ATOM14_ENCODING.token_atoms)[:20]
builder = prepare_named_layout(
prediction["aatype"],
prediction["chain_index"],
names,
AF2_ATOM14_ENCODING.token_atoms,
)
return builder(
prediction["positions"][-1], observed=prediction["atom14_atom_exists"].bool()
)
[ ]:
from tmol.io import canonical_form_from_pose_stack
def check_and_score(pose, input_coords, expected_ca):
"""Check identity and differentiate the score into supplied coordinates."""
ordering = default_canonical_ordering()
canonical = canonical_form_from_pose_stack(ordering, pose)
ca_slot = torch.tensor(
[
ordering.restypes_atom_index_mapping[name].get("CA", -1)
for name in ordering.restype_io_equiv_classes
],
device=device,
)
batch, residue = torch.nonzero(canonical.res_types >= 0, as_tuple=True)
observed_ca = canonical.coords[
batch, residue, ca_slot[canonical.res_types[batch, residue]]
]
torch.testing.assert_close(observed_ca, expected_ca[batch, residue], rtol=0, atol=0)
assert torch.isfinite(pose.coords[pose.real_atoms]).all()
module = beta2016_score_function(device).render_whole_pose_scoring_module(pose)
energy = module(pose.coords).sum()
(gradient,) = torch.autograd.grad(energy, input_coords)
assert torch.isfinite(energy) and torch.isfinite(gradient).all()
assert torch.any(gradient != 0)
return float(energy.detach()), gradient
prediction = torch.load(
repo / "tmol/tests/data/openfold/openfold_ubq_and_sumo.pt",
map_location=device,
weights_only=True,
)
openfold_coords = prediction["positions"].detach().clone().requires_grad_()
prediction["positions"] = openfold_coords
openfold_pose = openfold_example(prediction)
assert openfold_pose.n_poses == 2
openfold_energy, openfold_gradient = check_and_score(
openfold_pose, openfold_coords, openfold_coords[-1, :, :, 1]
)
print("OpenFold-style batch energy:", openfold_energy)
print("Coordinate-gradient norm:", float(openfold_gradient.norm()))
3. RF2-style output and hydrogen policy#
Use num2aa and aa2long from the RF2 version that produced the tensor. The legacy prediction below has 27 slots; other RF2-family versions can use different layouts. The saved layout is example data, not a library interface.
Choose hydrogens="preserve" to keep supplied nonterminal H coordinates, or "rebuild" to mark all supplied H missing. Both rebuild a generic N-terminal amide H using the appropriate terminal model. Rebuilt slots have zero input gradient. Supplied coordinates keep their gradient connection.
[ ]:
def rf2_example(prediction, residue_names, atom_names, *, hydrogens):
"""Use num2aa/aa2long from the RF2 version that produced this prediction.
``hydrogens='preserve'`` retains supplied nonterminal H coordinates;
``'rebuild'`` marks all supplied H as missing. The generic amide H at each
chain's N terminus is always rebuilt as the appropriate terminal hydrogens.
"""
if hydrogens not in ("preserve", "rebuild"):
raise ValueError("Choose hydrogens='preserve' or 'rebuild'")
sequence = prediction["seq"].unsqueeze(0)
coords = prediction["xyz"].unsqueeze(0)
lengths = prediction["chainlens"]
if any(length <= 0 for length in lengths) or sum(lengths) != sequence.shape[1]:
raise ValueError("Positive chain lengths must cover the entire sequence")
chain_id = torch.repeat_interleave(
torch.arange(len(lengths), device=coords.device),
torch.tensor(lengths, device=coords.device),
).unsqueeze(0)
names = {
name: [atom.strip() if atom else "" for atom in row]
for name, row in zip(residue_names, atom_names)
}
builder = prepare_named_layout(sequence, chain_id, residue_names, names)
hydrogen_slots = torch.tensor(
[
[atom.lstrip("123").startswith("H") for atom in names[name]]
for name in residue_names
],
device=coords.device,
)
amide_h = torch.tensor(
[[atom == "H" for atom in names[name]] for name in residue_names],
device=coords.device,
)
nterm = torch.ones_like(chain_id, dtype=torch.bool)
nterm[:, 1:] = chain_id[:, 1:] != chain_id[:, :-1]
suppressed = amide_h[sequence] & nterm.unsqueeze(-1)
if hydrogens == "rebuild":
suppressed |= hydrogen_slots[sequence]
return builder(coords, observed=~suppressed)
[ ]:
import json
rf2_prediction = torch.load(
repo / "tmol/tests/data/rosettafold2/ubiquitin.pt",
map_location=device,
weights_only=True,
)
layout = json.loads(
(repo / "tmol/tests/data/rosettafold2/input_layout.json").read_text()
)
rf2_results = {}
for policy in ("preserve", "rebuild"):
coords = rf2_prediction["xyz"].detach().clone().requires_grad_()
rf2_prediction["xyz"] = coords
pose = rf2_example(
rf2_prediction, layout["num2aa"], layout["aa2long"], hydrogens=policy
)
assert pose.n_poses == 1
energy, gradient = check_and_score(pose, coords, coords[None, :, 1])
if policy == "rebuild":
# These saved legacy tables place hydrogens at slots 14 onward.
assert torch.count_nonzero(gradient[:, 14:]) == 0
rf2_results[policy] = {"energy": energy, "gradient_norm": float(gradient.norm())}
print(rf2_results)
4. Scoring and repeated guidance#
For a fixed sequence and topology, reuse the named mapping and score module. Bind new coordinate tensors entirely in Torch; do not detach or convert them to NumPy on the differentiable path. AtomWorks file parsing and chemical preparation belong before this loop. NaN completion preserves atom identity and does not itself impute finite coordinates.
Discrete protonation choices, chemical topology changes and rotamer selection are outside this coordinate-gradient contract. A prepared-topology builder would also cache final topology; this small example only caches slot mappings.
[ ]:
from atomworks.ml.encoding_definitions import AF2_ATOM14_ENCODING
names = list(AF2_ATOM14_ENCODING.token_atoms)[:20]
bind = prepare_named_layout(
prediction["aatype"],
prediction["chain_index"],
names,
AF2_ATOM14_ENCODING.token_atoms,
)
mask = prediction["atom14_atom_exists"].bool()
coordinates = openfold_coords[-1].detach().clone()
score_module = None
guidance_energies = []
for step in range(3):
coordinates.requires_grad_()
pose = bind(coordinates, observed=mask)
if score_module is None:
score_module = beta2016_score_function(device).render_whole_pose_scoring_module(
pose
)
topology = pose.block_type_ind.clone()
torch.testing.assert_close(pose.block_type_ind, topology, rtol=0, atol=0)
energy = score_module(pose.coords).sum()
(gradient,) = torch.autograd.grad(energy, coordinates)
assert torch.isfinite(energy) and torch.isfinite(gradient).all()
guidance_energies.append(float(energy.detach()))
# Stop the tape between iterations, after obtaining the current gradient.
# An application can instead use this gradient in its diffusion/search update.
with torch.no_grad():
coordinates = coordinates - 0.001 * gradient / gradient.norm().clamp_min(1.0)
print("Three guidance evaluations:", guidance_energies)