Tutorial 01 — Working with structures#
Load 1UBQ from mmCIF, inspect its blocks and author labels, select atoms, and export structures. The examples use CPU or CUDA. Prefer mmCIF over PDB when metadata or ligand bonds matter.
Setup#
In Colab, select T4 GPU, then Run all. Locally, follow the installation guide. Setup downloads the checked-in fixture; calculations do not query a live structure database.
[1]:
try:
import google.colab # noqa: F401
except ImportError:
IN_COLAB = False
else:
IN_COLAB = True
if IN_COLAB:
from urllib.request import urlopen
exec(
urlopen(
"https://raw.githubusercontent.com/uw-ipd/tmol/"
"master/docs/tutorial/colab_setup.py"
).read(),
globals(),
)
setup_colab(["tmol/tests/data/cif/1UBQ.cif", "tmol/tests/data/pdb/1ubq.pdb"])
[2]:
from collections import Counter
from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from pathlib import Path
import tempfile
import warnings
import numpy as np
import pandas as pd
import torch
from IPython.display import display
from biotite.structure import AtomArray
from biotite.structure.io import load_structure
from biotite.structure.io.pdb import PDBFile
import tmol
from tmol.database import ParameterDatabase
from tmol.io import biotite_from_pose_stack, pose_stack_from_biotite
from tmol.io import write_pose_stack_pdb
from tmol.pose import PoseStackBuilder
SEED = 20260807
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(SEED)
device = (
torch.device("cuda", torch.cuda.current_device())
if torch.cuda.is_available()
else torch.device("cpu")
)
repo_root = Path.cwd()
if not (repo_root / "tmol/tests/data/cif/1UBQ.cif").exists():
repo_root = Path(tmol.__file__).resolve().parents[1]
cif_path = repo_root / "tmol/tests/data/cif/1UBQ.cif"
atom_array = load_structure(
str(cif_path),
model=1,
include_bonds=True,
extra_fields=["occupancy", "b_factor"],
)
assert isinstance(atom_array, AtomArray)
param_db = ParameterDatabase.get_default()
# Structure I/O may emit diagnostics while rebuilding missing atoms or handling
# unrecognized residues. Keep them out of the tutorial unless conversion fails.
pose_diagnostics = StringIO()
try:
with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):
pose_stack, build_context = pose_stack_from_biotite(
atom_array,
torch_device=device,
param_db=param_db,
no_optH=True,
return_context=True,
)
except Exception:
print(pose_diagnostics.getvalue())
raise
def show_table(frame):
"""Use sortable tables in rendered docs, with a pandas fallback."""
try:
from itables import show
except ImportError:
return display(frame)
return show(frame)
environment_frame = pd.DataFrame(
[
{"component": "TMol", "version": tmol.__version__},
{"component": "PyTorch", "version": torch.__version__},
{"component": "device", "version": str(device)},
{"component": "input", "version": cif_path.name},
]
)
show_table(environment_frame)
print(f"input atoms={atom_array.array_length()}")
Environment variable CCD_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
(1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
(2) set it in your current shell with 'export VAR_NAME=path/to/variable'
(3) write it to a .env file in the root of the atomworks.io repository
Environment variable PDB_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
(1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
(2) set it in your current shell with 'export VAR_NAME=path/to/variable'
(3) write it to a .env file in the root of the atomworks.io repository
input atoms=660
From chemistry to coordinates#
ParameterDatabase holds chemical and scoring definitions. PackedBlockTypes packs the required residue types onto one device. Each chemical unit is a block; PoseStack stores a batch as padded tensors. Keep the pose, packed types, and score function on the same device.
Conversion selects terminal and supported histidine variants, detects geometrically compatible disulfides, and builds supported missing atoms. It leaves the deposited atom_array unchanged. The returned context stores reusable chemistry and ordering; pass context=build_context when importing compatible structures.
Hydrogen preparation#
Default
no_optH=False: optimize supported proton chis/NHQ alternatives before interpreting all-atom scores.no_optH=True: skip optimization; built hydrogens retain ideal kinematic positions. This notebook uses it to test I/O round trips, not hydrogen-bond energies.Ligands still require authoritative bonds, charges, and parameterization.
no_optHdoes not supply these.
Record hydrogen preparation with score comparisons. Tutorials 03 and 04 optimize hydrogens; examples retaining prepared geometry state that choice.
The tables compare deposited and built counts and author labels. Matching labels permit coordinate comparisons, but the public API does not report retained-versus-constructed provenance for every atom.
[3]:
pbt = pose_stack.packed_block_types
real_blocks = pose_stack.block_type_ind64[0] >= 0
block_type_indices = pose_stack.block_type_ind64[0, real_blocks].detach().cpu().tolist()
block_names = [pbt.active_block_types[i].name for i in block_type_indices]
shape_table = pd.DataFrame(
[
("coords", tuple(pose_stack.coords.shape), str(pose_stack.coords.dtype)),
("block_coord_offset", tuple(pose_stack.block_coord_offset.shape), str(pose_stack.block_coord_offset.dtype)),
("block_type_ind", tuple(pose_stack.block_type_ind.shape), str(pose_stack.block_type_ind.dtype)),
("chain_id", tuple(pose_stack.chain_id.shape), str(pose_stack.chain_id.dtype)),
("real_atoms", tuple(pose_stack.real_atoms.shape), str(pose_stack.real_atoms.dtype)),
],
columns=["field", "shape", "dtype"],
)
show_table(shape_table)
print(
f"n_poses={pose_stack.n_poses}, max_n_blocks={pose_stack.max_n_blocks}, "
f"max_n_pose_atoms={pose_stack.max_n_pose_atoms}"
)
print("first five block types:", block_names[:5])
print("PackedBlockTypes device:", pbt.device)
print(
"all resolved atom coordinates finite:",
bool(torch.isfinite(pose_stack.coords[pose_stack.real_atoms]).all()),
)
print("context reuses ParameterDatabase:", build_context.parameter_database is param_db)
n_poses=1, max_n_blocks=76, max_n_pose_atoms=1231
first five block types: ['MET:nterm', 'GLN', 'ILE', 'PHE', 'VAL']
PackedBlockTypes device: cpu
all resolved atom coordinates finite: True
context reuses ParameterDatabase: True
coords has shape [n_poses, max_n_pose_atoms, 3]; block fields have shape [n_poses, max_n_blocks, ...]. Use sentinel block indices and real_atoms to identify padding.
TMol selects among compatible terminal variants. Warnings about unselected candidates can occur during successful construction; this notebook captures them separately. Finite coordinates check construction, while the round-trip section measures coordinate agreement.
pdb_info stores author labels separately from kernel indices. Selected block types may include atoms absent from the deposited structure.
[4]:
residue_frame = pd.DataFrame(
{
"block_index": np.arange(pose_stack.max_n_blocks)[real_blocks.cpu().numpy()],
"chain": pose_stack.pdb_info.chain_labels[0, real_blocks.cpu().numpy()],
"residue_number": pose_stack.pdb_info.residue_labels[
0, real_blocks.cpu().numpy()
],
"block_type": block_names,
"n_atoms": pose_stack.n_ats_per_block[0, real_blocks].detach().cpu().numpy(),
}
)
show_table(residue_frame)
Round-trip structures and export PDB#
biotite_from_pose_stack() returns TMol-built atoms and coordinates directly as a Biotite structure. Pass the build context’s canonical ordering for custom residue or ligand types.
The next cell matches (chain, residue number, insertion code, atom name) labels and reports heavy-atom RMSD, maximum displacement, and net atom-count changes. RMSD is unaligned, in the original coordinate frame; assertions are specific to 1UBQ.
Reimporting with context=build_context reuses chemistry while rebuilding coordinates and labels. A separate PDB write/read measures coordinate rounding. PDB cannot preserve all mmCIF metadata or ligand bond orders.
[5]:
tmol_atom_array = biotite_from_pose_stack(
pose_stack, build_context.canonical_ordering
)
assert isinstance(tmol_atom_array, AtomArray)
def atom_indices_by_author_label(structure):
"""Map practical author labels to indices; labels are not provenance."""
labels = [
(
str(structure.chain_id[i]),
int(structure.res_id[i]),
str(structure.ins_code[i]).strip(),
str(structure.atom_name[i]).strip(),
)
for i in range(structure.array_length())
]
duplicates = [label for label, count in Counter(labels).items() if count > 1]
if duplicates:
raise ValueError(f"author atom labels are not unique: {duplicates[:3]}")
return dict(zip(labels, range(len(labels))))
def common_heavy_atom_comparison(reference, comparison):
reference_indices = atom_indices_by_author_label(reference)
comparison_indices = atom_indices_by_author_label(comparison)
common_labels = sorted(reference_indices.keys() & comparison_indices.keys())
heavy_labels = [
label
for label in common_labels
if reference.element[reference_indices[label]].upper() != "H"
and comparison.element[comparison_indices[label]].upper() != "H"
]
if not heavy_labels:
raise ValueError("no common heavy-atom author labels")
reference_xyz = np.array(
[reference.coord[reference_indices[label]] for label in heavy_labels]
)
comparison_xyz = np.array(
[comparison.coord[comparison_indices[label]] for label in heavy_labels]
)
displacements = np.linalg.norm(comparison_xyz - reference_xyz, axis=1)
displacement_frame = pd.DataFrame(
{
"author_atom_label": [
f"{chain}/{resid}{ins_code}/{atom_name}"
for chain, resid, ins_code, atom_name in heavy_labels
],
"displacement_A": displacements,
}
).sort_values("displacement_A", ascending=False)
metrics = {
"common_heavy_atom_labels": len(heavy_labels),
"reference_only_atom_labels": len(reference_indices.keys() - comparison_indices.keys()),
"comparison_only_atom_labels": len(comparison_indices.keys() - reference_indices.keys()),
"heavy_atom_RMSD_A": float(np.sqrt(np.mean(displacements**2))),
"max_heavy_atom_displacement_A": float(displacements.max()),
}
return metrics, displacement_frame
deposited_indices = atom_indices_by_author_label(atom_array)
built_indices = atom_indices_by_author_label(tmol_atom_array)
deposited_only_labels = deposited_indices.keys() - built_indices.keys()
built_only_labels = built_indices.keys() - deposited_indices.keys()
known_residue_names = set(build_context.canonical_ordering.restype_io_equiv_classes)
deposited_only_reasons = Counter()
for label in deposited_only_labels:
atom_index = deposited_indices[label]
residue_name = str(atom_array.res_name[atom_index])
atom_name = str(atom_array.atom_name[atom_index]).strip()
if residue_name == "HOH":
reason = "water is excluded by the current conversion path"
elif residue_name not in known_residue_names:
reason = "residue name is not recognized by the canonical ordering"
elif atom_name not in build_context.canonical_ordering.restypes_atom_index_mapping.get(
residue_name, {}
):
reason = "atom name is absent from the residue's canonical mapping"
else:
reason = "label absent after block/variant selection; no finer public reason is exposed"
deposited_only_reasons[reason] += 1
common_label_count = len(deposited_indices.keys() & built_indices.keys())
audit_rows = [
{
"category": "deposited model 1",
"atoms": atom_array.array_length(),
"interpretation": "atoms read from the checked-in mmCIF model",
},
{
"category": "common author labels",
"atoms": common_label_count,
"interpretation": "same chain/residue/insertion-code/atom-name; not provenance",
},
]
audit_rows.extend(
{
"category": "deposited-only labels",
"atoms": count,
"interpretation": reason,
}
for reason, count in deposited_only_reasons.items()
)
audit_rows.append(
{
"category": "TMol-built-only labels",
"atoms": len(built_only_labels),
"interpretation": (
"labels present only after chemical-model selection/building; "
"exact per-atom construction provenance is not exposed"
),
}
)
audit_rows.append(
{
"category": "TMol-built total",
"atoms": tmol_atom_array.array_length(),
"interpretation": "atoms exported from the selected TMol block types",
}
)
show_table(pd.DataFrame(audit_rows))
deposited_metrics, deposited_displacements = common_heavy_atom_comparison(
atom_array, tmol_atom_array
)
assert deposited_metrics["heavy_atom_RMSD_A"] < 0.05
assert deposited_metrics["max_heavy_atom_displacement_A"] < 0.10
reuse_diagnostics = StringIO()
try:
with redirect_stdout(reuse_diagnostics), redirect_stderr(reuse_diagnostics):
reused_pose_stack = pose_stack_from_biotite(
tmol_atom_array,
torch_device=device,
context=build_context,
no_optH=True,
)
except Exception:
print(reuse_diagnostics.getvalue())
raise
reused_atom_array = biotite_from_pose_stack(
reused_pose_stack, build_context.canonical_ordering
)
reuse_metrics, reuse_displacements = common_heavy_atom_comparison(
tmol_atom_array, reused_atom_array
)
assert reuse_metrics["heavy_atom_RMSD_A"] < 0.01
assert reuse_metrics["max_heavy_atom_displacement_A"] < 0.02
roundtrip_path = Path(tempfile.gettempdir()) / "1ubq_tmol_roundtrip.pdb"
# PDB cannot encode every TMol/Biotite annotation; that limitation is already
# explained above, so suppress Biotite's duplicate compatibility warning.
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
write_pose_stack_pdb(pose_stack, str(roundtrip_path))
roundtrip_array = PDBFile.read(str(roundtrip_path)).get_structure(
model=1,
include_bonds=True,
extra_fields=["occupancy", "b_factor"],
)
pdb_metrics, pdb_displacements = common_heavy_atom_comparison(
tmol_atom_array, roundtrip_array
)
assert pdb_metrics["heavy_atom_RMSD_A"] < 0.01
assert pdb_metrics["max_heavy_atom_displacement_A"] < 0.02
comparison_frame = pd.DataFrame(
[
{"comparison": "deposited model 1 → TMol-built", **deposited_metrics},
{"comparison": "TMol-built → context-reused TMol-built", **reuse_metrics},
{"comparison": "TMol-built → PDB compatibility read", **pdb_metrics},
]
)
show_table(comparison_frame)
print("largest deposited-to-built common-heavy-atom displacements:")
show_table(deposited_displacements.head(5))
print(
"context PackedBlockTypes reused:",
reused_pose_stack.packed_block_types is build_context.packed_block_types,
)
print("wrote:", roundtrip_path.resolve())
| Loading ITables v2.9.1 from the internet... (need help?) |
| ⓘcategory | atoms | interpretation |
|---|---|---|
| deposited model 1 | 660 | atoms read from the checked-in mmCIF model |
| common author labels | 602 | same chain/residue/insertion-code/atom-name; not provenance |
| deposited-only labels | 58 | water is excluded by the current conversion path |
| TMol-built-only labels | 629 | labels present only after chemical-model selection/building; exact per-atom construction provenance is not exposed |
| TMol-built total | 1231 | atoms exported from the selected TMol block types |
| Loading ITables v2.9.1 from the internet... (need help?) |
| ⓘcomparison | common_heavy_atom_labels | reference_only_atom_labels | comparison_only_atom_labels | heavy_atom_RMSD_A | max_heavy_atom_displacement_A |
|---|---|---|---|---|---|
| deposited model 1 → TMol-built | 602 | 58 | 629 | 0.0 | 0.0 |
| TMol-built → context-reused TMol-built | 602 | 0 | 0 | 0.0 | 0.0 |
| TMol-built → PDB compatibility read | 602 | 0 | 0 | 0.0 | 0.0 |
largest deposited-to-built common-heavy-atom displacements:
context PackedBlockTypes reused: True
wrote: /tmp/1ubq_tmol_roundtrip.pdb
Direct PDB API, residue slices, and batched output#
tmol.pose_stack_from_pdb() is the concise compatibility path for a PDB filename or PDB lines. residue_start and residue_end select a zero-based, half-open range in parsed residue order; they are not PDB author residue numbers. A slice is normally treated as a new chain segment with termini. To represent an internal unresolved cut instead, pass res_not_connected[p, i, 0] = True at a missing upstream connection or [..., 1] = True at a missing downstream connection. At a
selected range boundary, those flags preserve a nonterminal block with an incomplete connection rather than inventing terminal chemistry.
write_pose_stack_pdb() writes every pose in one PoseStack as a PDB MODEL. Split the batch first when downstream software requires one file per model. Both exports inherit PDB’s metadata and ligand-chemistry limitations.
Prediction tensors#
Use tmol.io.pose_stack_from_canonical_aa_atom37() for canonical proteins in AtomWorks’ unified Atom37 encoding. Other atom37/atom14 layouts map named slots and token IDs into CanonicalForm; the executable model-input tutorial demonstrates OpenFold and RF2 mappings, including an explicit policy for RF2 hydrogen slots. Missing supported atoms are built differentiably, so scores can backpropagate to supplied prediction coordinates. Repeated Atom37 guidance with annotated topology can reuse
prepare_atom37_pose_builder() and a compatible scoring module.
See the Task index and structure I/O and integrations for these input contracts and the executable tutorial.
[6]:
pdb_path = repo_root / "tmol/tests/data/pdb/1ubq.pdb"
full_pdb_pose = tmol.pose_stack_from_pdb(str(pdb_path), device=device)
# Parsed residue positions 19:25 correspond to six residues in this fixture.
# Mark both outer connections incomplete so this is an internal fragment, not
# newly capped N/C termini.
internal_cut_flags = torch.zeros((1, 6, 2), dtype=torch.bool, device=device)
internal_cut_flags[0, 0, 0] = True
internal_cut_flags[0, -1, 1] = True
internal_slice = tmol.pose_stack_from_pdb(
str(pdb_path),
device=device,
residue_start=19,
residue_end=25,
res_not_connected=internal_cut_flags,
)
assert internal_slice.max_n_blocks == 6
assert int(internal_slice.inter_residue_connections[0, 0, 0, 0]) == -1
assert int(internal_slice.inter_residue_connections[0, -1, 1, 0]) == -1
pdb_batch = PoseStackBuilder.from_poses([full_pdb_pose] * 3, device=device)
with tempfile.TemporaryDirectory(prefix="tmol-pdb-output-") as temp_dir:
temp_dir = Path(temp_dir)
multi_model_path = temp_dir / "ubiquitin_batch.pdb"
write_pose_stack_pdb(pdb_batch, str(multi_model_path))
model_count = sum(
line.startswith("MODEL ") for line in multi_model_path.read_text().splitlines()
)
separate_paths = []
for pose_index in range(pdb_batch.n_poses):
output_path = temp_dir / f"ubiquitin_{pose_index:02d}.pdb"
write_pose_stack_pdb(pdb_batch.split(pose_index), str(output_path))
separate_paths.append(output_path)
assert model_count == pdb_batch.n_poses
assert all(path.is_file() and path.stat().st_size > 0 for path in separate_paths)
print(
f"direct PDB blocks={full_pdb_pose.max_n_blocks}; "
f"internal slice blocks={internal_slice.max_n_blocks}"
)
print(
f"batch poses={pdb_batch.n_poses}; multi-model MODEL records={model_count}; "
f"separate files={len(separate_paths)}"
)
direct PDB blocks=76; internal slice blocks=6
batch poses=3; multi-model MODEL records=3; separate files=3
Select and visualize deposited atoms#
Combine Biotite’s NumPy chain, residue, atom-name, and element annotations into Boolean masks. These masks index the deposited atom_array; recompute them for the differently sized TMol-built structure.
tmol.selection_gallery() displays an AtomArray with named selection buttons. Use pose_stack or tmol_atom_array when inspecting built atoms or resolved chemistry.
[7]:
selection_mask = (
(atom_array.chain_id == "A")
& (atom_array.res_id >= 1)
& (atom_array.res_id <= 10)
& np.isin(atom_array.atom_name, ["N", "CA", "C", "O"])
)
sidechain_mask = (
(atom_array.chain_id == "A")
& (atom_array.res_id >= 1)
& (atom_array.res_id <= 10)
& ~np.isin(atom_array.atom_name, ["N", "CA", "C", "O", "OXT"])
)
selected_atoms = atom_array[selection_mask]
print("selected backbone atoms:", selected_atoms.array_length())
print("selected residues:", np.unique(selected_atoms.res_id))
try:
display(
tmol.selection_gallery(
atom_array,
{
"Backbone, residues 1–10": selection_mask,
"Side chains, residues 1–10": sidechain_mask,
"All Cα atoms": atom_array.atom_name == "CA",
},
width=720,
height=420,
)
)
except ImportError as exc:
print("Selection viewer unavailable in this environment:", exc)
selected backbone atoms: 40
selected residues: [ 1 2 3 4 5 6 7 8 9 10]
/home/runner/work/tmol/tmol/.venv/lib/python3.12/site-packages/biotite/structure/io/pdb/file.py:629: DeprecationWarning: The chararray class is deprecated and will be removed in a future release. Use an ndarray with a string or bytes dtype instead.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
Next#
Continue with GPU batching and scoring.
Exercises#
Change the device selection to force CPU and confirm all tensor devices agree.
Select residues 20–30 and heavy side-chain atoms with a Biotite/NumPy Boolean mask; assert the selected author labels.
Build the same PDB residue range once as a new terminal fragment and once with internal-cut flags; compare selected block-type names and connectivity.
Reuse
build_contextto import the PDB-readroundtrip_array; compare its block types with the directreused_pose_stackand explain any differences.Write a two-pose batch as both one multi-model PDB and separate files; verify model count and author residue labels rather than raw text equality.