{ "cells": [ { "cell_type": "markdown", "id": "8424f6fb", "metadata": {}, "source": [ "# Tutorial 01 — Working with TMol\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/uw-ipd/tmol/blob/master/docs/tutorial/01_working_with_tmol.ipynb)\n", "\n", "This tutorial introduces `PoseStack`, TMol's tensor-backed molecular representation. You will load 1UBQ from mmCIF, inspect its blocks and labels, select atoms, and write structures.\n", "\n", "## Learning objectives\n", "\n", "- Choose a PyTorch device and load a structure.\n", "- Inspect blocks, coordinates, and author labels.\n", "- Write one pose or a batch of poses.\n", "\n", "## Before you begin\n", "\n", "- **Prerequisites:** Basic Python, NumPy, and PyTorch tensor familiarity; no prior TMol tutorial is required.\n", "- **Curriculum:** Start here, then continue through [02 — GPU batching](02_gpu_batching.ipynb) and [03 — Scoring and analysis](03_scoring_and_analysis.ipynb). After 03, complete the parallel [04 — Packing](04_packing_and_mutation_scan.ipynb) and [05 — Minimization](05_minimization_constraints_kinematics.ipynb) branches before [06 — FastRelax](06_fast_relax.ipynb); 07 and 08 are specialized paths.\n", "- **Related:** [Structure I/O and integrations](../user_guide/integrations.md) · [PoseStack API](../api/pose.rst)\n", "\n", "Use mmCIF when possible because it preserves richer metadata and chemical bonds than PDB. The examples run on CPU; CUDA is used automatically when available." ] }, { "cell_type": "markdown", "id": "0d5206d7", "metadata": {}, "source": [ "## Setup\n", "\n", "Imports, reproducibility, device selection, and fixture discovery live under this heading so the documentation can collapse setup details. The input is checked into the repository; no network access is used." ] }, { "cell_type": "code", "execution_count": 1, "id": "7e06483e", "metadata": {}, "outputs": [], "source": [ "try:\n", " import google.colab # noqa: F401\n", "except ImportError:\n", " IN_COLAB = False\n", "else:\n", " IN_COLAB = True\n", "\n", "if IN_COLAB:\n", " from urllib.request import urlopen\n", "\n", " exec(\n", " urlopen(\n", " \"https://raw.githubusercontent.com/uw-ipd/tmol/\"\n", " \"master/docs/tutorial/colab_setup.py\"\n", " ).read(),\n", " globals(),\n", " )\n", " setup_colab([\"tmol/tests/data/cif/1UBQ.cif\"])" ] }, { "cell_type": "code", "execution_count": 2, "id": "1c41103a", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", "\n", "\n", "\n", "\n", "
ⓘcomponentversion
TMol0.1.54
PyTorch2.14.0+cpu
devicecpu
input1UBQ.cif
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "input atoms=660\n" ] } ], "source": [ "from collections import Counter\n", "from contextlib import redirect_stderr, redirect_stdout\n", "from io import StringIO\n", "from pathlib import Path\n", "import tempfile\n", "import warnings\n", "\n", "import numpy as np\n", "import pandas as pd\n", "import torch\n", "from IPython.display import display\n", "from biotite.structure import AtomArray\n", "from biotite.structure.io import load_structure\n", "from biotite.structure.io.pdb import PDBFile\n", "\n", "import tmol\n", "from tmol.database import ParameterDatabase\n", "from tmol.io import biotite_from_pose_stack, pose_stack_from_biotite\n", "from tmol.io import write_pose_stack_pdb\n", "from tmol.pose import PoseStackBuilder\n", "\n", "SEED = 20260807\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)\n", "if torch.cuda.is_available():\n", " torch.cuda.manual_seed_all(SEED)\n", "\n", "device = (\n", " torch.device(\"cuda\", torch.cuda.current_device())\n", " if torch.cuda.is_available()\n", " else torch.device(\"cpu\")\n", ")\n", "repo_root = Path.cwd()\n", "if not (repo_root / \"tmol/tests/data/cif/1UBQ.cif\").exists():\n", " repo_root = Path(tmol.__file__).resolve().parents[1]\n", "cif_path = repo_root / \"tmol/tests/data/cif/1UBQ.cif\"\n", "\n", "atom_array = load_structure(\n", " str(cif_path),\n", " model=1,\n", " include_bonds=True,\n", " extra_fields=[\"occupancy\", \"b_factor\"],\n", ")\n", "assert isinstance(atom_array, AtomArray)\n", "\n", "param_db = ParameterDatabase.get_default()\n", "# Structure I/O may emit diagnostics while rebuilding missing atoms or handling\n", "# unrecognized residues. Keep them out of the tutorial unless conversion fails.\n", "pose_diagnostics = StringIO()\n", "try:\n", " with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):\n", " pose_stack, build_context = pose_stack_from_biotite(\n", " atom_array,\n", " torch_device=device,\n", " param_db=param_db,\n", " no_optH=True,\n", " return_context=True,\n", " )\n", "except Exception:\n", " print(pose_diagnostics.getvalue())\n", " raise\n", "\n", "\n", "def show_table(frame):\n", " \"\"\"Use sortable tables in rendered docs, with a pandas fallback.\"\"\"\n", " try:\n", " from itables import show\n", " except ImportError:\n", " return display(frame)\n", " return show(frame)\n", "\n", "\n", "environment_frame = pd.DataFrame(\n", " [\n", " {\"component\": \"TMol\", \"version\": tmol.__version__},\n", " {\"component\": \"PyTorch\", \"version\": torch.__version__},\n", " {\"component\": \"device\", \"version\": str(device)},\n", " {\"component\": \"input\", \"version\": cif_path.name},\n", " ]\n", ")\n", "show_table(environment_frame)\n", "print(f\"input atoms={atom_array.array_length()}\")" ] }, { "cell_type": "markdown", "id": "037808c2", "metadata": {}, "source": [ "## From chemistry to coordinates\n", "\n", "`ParameterDatabase` is the immutable source of chemical and scoring parameters. A `PackedBlockTypes` object packs the residue types needed by a system onto one device. Each residue-like unit is a **block**, and `PoseStack` stores one or more poses as padded, contiguous tensors over those blocks.\n", "\n", "A `PoseStack` lives on one `torch.device`; build its `PackedBlockTypes` and score function on that same device. A batch of N structures is one `PoseStack` with `n_poses=N`, not N separate Rosetta `Pose` objects.\n", "\n", "Biotite conversion maps the deposited atoms into TMol's chemical model. During that process TMol detects geometrically compatible disulfides, resolves supported HIS protonation/tautomer variants from the supplied atoms, chooses terminal variants, and builds supported missing atoms. These are chemical-model decisions in the new `PoseStack`; they do not mutate the deposited `atom_array`.\n", "\n", "The returned build context contains the database, residue-type set, canonical ordering, and packed block types used for conversion. Passing it back as `context=build_context` reuses that structure-independent chemistry for another compatible input while still rebuilding per-structure coordinates and labels. The round-trip section below demonstrates this explicitly.\n", "\n", "### Preparation and `no_optH` decision guide\n", "\n", "- Use the default `no_optH=False` before interpreting all-atom scores, especially hydrogen bonds: TMol builds supported missing atoms, samples supported proton chis/NHQ alternatives, and optimizes hydrogen placement as part of conversion.\n", "- Use `no_optH=True` for I/O round trips, geometry plumbing, or controlled benchmarks where hydrogen optimization is intentionally outside the experiment. Hydrogens are then built in ideal kinematic positions but are not optimized; do not silently treat resulting hbond values as hydrogen-prepared scores.\n", "- For deposited ligands or custom chemistry, preserve authoritative CIF/MOL2 bonds and charges and use a validated preparation path. `no_optH` is not a substitute for ligand parameterization.\n", "- Record the choice with every score comparison. Changing hydrogen preparation changes the molecular model, not merely runtime.\n", "\n", "This notebook uses `no_optH=True` because it evaluates coordinate and label round trips rather than hbond energies. Tutorials 03 and 04 use optimized hydrogen coordinates for score interpretation; tutorials that deliberately retain deposited or prepared hydrogens state that limitation.\n", "\n", "The comparison below reports deposited and TMol-built counts, author-label matches, and reasons exposed by the current APIs, such as excluded water or an unsupported atom label. A matching label supports coordinate comparison but does not prove retained-versus-constructed atom provenance, which the public conversion API does not expose." ] }, { "cell_type": "code", "execution_count": 3, "id": "3ae683f1", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", "\n", "\n", "
ⓘfieldshapedtype
coords(1, 1231, 3)torch.float32
block_coord_offset(1, 76)torch.int32
block_type_ind(1, 76)torch.int32
chain_id(1, 76)torch.int32
real_atoms(1, 1231)torch.bool
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "n_poses=1, max_n_blocks=76, max_n_pose_atoms=1231\n", "first five block types: ['MET:nterm', 'GLN', 'ILE', 'PHE', 'VAL']\n", "PackedBlockTypes device: cpu\n", "all resolved atom coordinates finite: True\n", "context reuses ParameterDatabase: True\n" ] } ], "source": [ "pbt = pose_stack.packed_block_types\n", "real_blocks = pose_stack.block_type_ind64[0] >= 0\n", "block_type_indices = pose_stack.block_type_ind64[0, real_blocks].detach().cpu().tolist()\n", "block_names = [pbt.active_block_types[i].name for i in block_type_indices]\n", "\n", "shape_table = pd.DataFrame(\n", " [\n", " (\"coords\", tuple(pose_stack.coords.shape), str(pose_stack.coords.dtype)),\n", " (\"block_coord_offset\", tuple(pose_stack.block_coord_offset.shape), str(pose_stack.block_coord_offset.dtype)),\n", " (\"block_type_ind\", tuple(pose_stack.block_type_ind.shape), str(pose_stack.block_type_ind.dtype)),\n", " (\"chain_id\", tuple(pose_stack.chain_id.shape), str(pose_stack.chain_id.dtype)),\n", " (\"real_atoms\", tuple(pose_stack.real_atoms.shape), str(pose_stack.real_atoms.dtype)),\n", " ],\n", " columns=[\"field\", \"shape\", \"dtype\"],\n", ")\n", "show_table(shape_table)\n", "print(\n", " f\"n_poses={pose_stack.n_poses}, max_n_blocks={pose_stack.max_n_blocks}, \"\n", " f\"max_n_pose_atoms={pose_stack.max_n_pose_atoms}\"\n", ")\n", "print(\"first five block types:\", block_names[:5])\n", "print(\"PackedBlockTypes device:\", pbt.device)\n", "print(\n", " \"all resolved atom coordinates finite:\",\n", " bool(torch.isfinite(pose_stack.coords[pose_stack.real_atoms]).all()),\n", ")\n", "print(\"context reuses ParameterDatabase:\", build_context.parameter_database is param_db)" ] }, { "cell_type": "markdown", "id": "3bdd7aef", "metadata": {}, "source": [ "**Expected observations.** `coords` has shape `[n_poses, max_n_pose_atoms, 3]`; block-indexed fields have shape `[n_poses, max_n_blocks, ...]`. Padding is represented by sentinel block indices, while `real_atoms` identifies coordinate rows belonging to actual atoms. During construction TMol evaluates compatible terminal block-type alternatives and selects the best match. Unselected alternatives are normal internal candidates rather than malformed atoms, although the selector may emit a diagnostic warning while still constructing the correct terminal variant; this notebook captures those diagnostics separately. The finite-coordinate check is only a construction sanity check; the round-trip section performs the quantitative structural validation.\n", "\n", "`pdb_info` preserves author-facing labels separately from the integer chain and block indices used by kernels. The block atom counts below describe the selected TMol residue types, which may include built atoms absent from the deposited structure." ] }, { "cell_type": "code", "execution_count": 4, "id": "f05ea65b", "metadata": { "tags": [ "collapse-code" ] }, "outputs": [ { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
ⓘblock_indexchainresidue_numberblock_typen_atoms
0A1MET:nterm19
1A2GLN17
2A3ILE19
3A4PHE20
4A5VAL16
5A6LYS22
6A7THR14
7A8LEU19
8A9THR14
9A10GLY7
(66 more rows not shown)
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "residue_frame = pd.DataFrame(\n", " {\n", " \"block_index\": np.arange(pose_stack.max_n_blocks)[real_blocks.cpu().numpy()],\n", " \"chain\": pose_stack.pdb_info.chain_labels[0, real_blocks.cpu().numpy()],\n", " \"residue_number\": pose_stack.pdb_info.residue_labels[\n", " 0, real_blocks.cpu().numpy()\n", " ],\n", " \"block_type\": block_names,\n", " \"n_atoms\": pose_stack.n_ats_per_block[0, real_blocks].detach().cpu().numpy(),\n", " }\n", ")\n", "show_table(residue_frame)" ] }, { "cell_type": "markdown", "id": "2bcbd4c1", "metadata": {}, "source": [ "## Scientific round-trip and PDB compatibility export\n", "\n", "`biotite_from_pose_stack()` is the direct scientific inverse of the import path: it returns a Biotite structure from TMol's blocks and coordinates without first serializing through PDB. Supplying the build context's canonical ordering is important when custom residue or ligand types are present. The result is a **TMol-built** `AtomArray`, not a copy of the deposited `atom_array`: it reflects the selected chemical model.\n", "\n", "The next cell reports net atom-count changes without claiming unavailable atom-level provenance. It then compares coordinates by the practical author label `(chain, residue number, insertion code, atom name)`. For common heavy-atom labels it reports an unaligned RMSD and maximum displacement in the original coordinate frame, plus the largest-displacement labels. The clean 1UBQ fixture has fixture-specific assertions on those quantitative values.\n", "\n", "The built structure is also converted back to a `PoseStack` with `context=build_context`. This reuses the parameter database, canonical ordering, residue-type set, and device-packed block types; only structure-dependent canonicalization and pose construction are repeated.\n", "\n", "PDB remains useful for compatibility, so the cell also writes and reads a PDB file and quantifies the coordinate rounding on common heavy atoms. Its path is visible for download. PDB cannot carry all mmCIF metadata or reliable ligand bond orders, so compare intended chemistry rather than treating that file export as lossless." ] }, { "cell_type": "code", "execution_count": 5, "id": "b13c826c", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", "\n", "\n", "
ⓘcategoryatomsinterpretation
deposited model 1660atoms read from the checked-in mmCIF model
common author labels602same chain/residue/insertion-code/atom-name; not provenance
deposited-only labels58water is excluded by the current conversion path
TMol-built-only labels629labels present only after chemical-model selection/building; exact per-atom construction provenance is not exposed
TMol-built total1231atoms exported from the selected TMol block types
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", "
ⓘcomparisoncommon_heavy_atom_labelsreference_only_atom_labelscomparison_only_atom_labelsheavy_atom_RMSD_Amax_heavy_atom_displacement_A
deposited model 1 → TMol-built602586290.00.0
TMol-built → context-reused TMol-built602000.00.0
TMol-built → PDB compatibility read602000.00.0
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "largest deposited-to-built common-heavy-atom displacements:\n" ] }, { "data": { "text/html": [ "\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "
\n", " \n", " \n", " \n", "\n", "\n", "\n", "\n", "\n", "
ⓘauthor_atom_labeldisplacement_A
A/1/C0.0
A/1/CA0.0
A/1/CB0.0
A/1/CE0.0
A/1/CG0.0
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "context PackedBlockTypes reused: True\n", "wrote: /tmp/1ubq_tmol_roundtrip.pdb\n" ] } ], "source": [ "tmol_atom_array = biotite_from_pose_stack(\n", " pose_stack, build_context.canonical_ordering\n", ")\n", "assert isinstance(tmol_atom_array, AtomArray)\n", "\n", "\n", "def atom_indices_by_author_label(structure):\n", " \"\"\"Map practical author labels to indices; labels are not provenance.\"\"\"\n", " labels = [\n", " (\n", " str(structure.chain_id[i]),\n", " int(structure.res_id[i]),\n", " str(structure.ins_code[i]).strip(),\n", " str(structure.atom_name[i]).strip(),\n", " )\n", " for i in range(structure.array_length())\n", " ]\n", " duplicates = [label for label, count in Counter(labels).items() if count > 1]\n", " if duplicates:\n", " raise ValueError(f\"author atom labels are not unique: {duplicates[:3]}\")\n", " return dict(zip(labels, range(len(labels))))\n", "\n", "\n", "def common_heavy_atom_comparison(reference, comparison):\n", " reference_indices = atom_indices_by_author_label(reference)\n", " comparison_indices = atom_indices_by_author_label(comparison)\n", " common_labels = sorted(reference_indices.keys() & comparison_indices.keys())\n", " heavy_labels = [\n", " label\n", " for label in common_labels\n", " if reference.element[reference_indices[label]].upper() != \"H\"\n", " and comparison.element[comparison_indices[label]].upper() != \"H\"\n", " ]\n", " if not heavy_labels:\n", " raise ValueError(\"no common heavy-atom author labels\")\n", "\n", " reference_xyz = np.array(\n", " [reference.coord[reference_indices[label]] for label in heavy_labels]\n", " )\n", " comparison_xyz = np.array(\n", " [comparison.coord[comparison_indices[label]] for label in heavy_labels]\n", " )\n", " displacements = np.linalg.norm(comparison_xyz - reference_xyz, axis=1)\n", " displacement_frame = pd.DataFrame(\n", " {\n", " \"author_atom_label\": [\n", " f\"{chain}/{resid}{ins_code}/{atom_name}\"\n", " for chain, resid, ins_code, atom_name in heavy_labels\n", " ],\n", " \"displacement_A\": displacements,\n", " }\n", " ).sort_values(\"displacement_A\", ascending=False)\n", " metrics = {\n", " \"common_heavy_atom_labels\": len(heavy_labels),\n", " \"reference_only_atom_labels\": len(reference_indices.keys() - comparison_indices.keys()),\n", " \"comparison_only_atom_labels\": len(comparison_indices.keys() - reference_indices.keys()),\n", " \"heavy_atom_RMSD_A\": float(np.sqrt(np.mean(displacements**2))),\n", " \"max_heavy_atom_displacement_A\": float(displacements.max()),\n", " }\n", " return metrics, displacement_frame\n", "\n", "\n", "deposited_indices = atom_indices_by_author_label(atom_array)\n", "built_indices = atom_indices_by_author_label(tmol_atom_array)\n", "deposited_only_labels = deposited_indices.keys() - built_indices.keys()\n", "built_only_labels = built_indices.keys() - deposited_indices.keys()\n", "known_residue_names = set(build_context.canonical_ordering.restype_io_equiv_classes)\n", "\n", "deposited_only_reasons = Counter()\n", "for label in deposited_only_labels:\n", " atom_index = deposited_indices[label]\n", " residue_name = str(atom_array.res_name[atom_index])\n", " atom_name = str(atom_array.atom_name[atom_index]).strip()\n", " if residue_name == \"HOH\":\n", " reason = \"water is excluded by the current conversion path\"\n", " elif residue_name not in known_residue_names:\n", " reason = \"residue name is not recognized by the canonical ordering\"\n", " elif atom_name not in build_context.canonical_ordering.restypes_atom_index_mapping.get(\n", " residue_name, {}\n", " ):\n", " reason = \"atom name is absent from the residue's canonical mapping\"\n", " else:\n", " reason = \"label absent after block/variant selection; no finer public reason is exposed\"\n", " deposited_only_reasons[reason] += 1\n", "\n", "common_label_count = len(deposited_indices.keys() & built_indices.keys())\n", "audit_rows = [\n", " {\n", " \"category\": \"deposited model 1\",\n", " \"atoms\": atom_array.array_length(),\n", " \"interpretation\": \"atoms read from the checked-in mmCIF model\",\n", " },\n", " {\n", " \"category\": \"common author labels\",\n", " \"atoms\": common_label_count,\n", " \"interpretation\": \"same chain/residue/insertion-code/atom-name; not provenance\",\n", " },\n", "]\n", "audit_rows.extend(\n", " {\n", " \"category\": \"deposited-only labels\",\n", " \"atoms\": count,\n", " \"interpretation\": reason,\n", " }\n", " for reason, count in deposited_only_reasons.items()\n", ")\n", "audit_rows.append(\n", " {\n", " \"category\": \"TMol-built-only labels\",\n", " \"atoms\": len(built_only_labels),\n", " \"interpretation\": (\n", " \"labels present only after chemical-model selection/building; \"\n", " \"exact per-atom construction provenance is not exposed\"\n", " ),\n", " }\n", ")\n", "audit_rows.append(\n", " {\n", " \"category\": \"TMol-built total\",\n", " \"atoms\": tmol_atom_array.array_length(),\n", " \"interpretation\": \"atoms exported from the selected TMol block types\",\n", " }\n", ")\n", "show_table(pd.DataFrame(audit_rows))\n", "\n", "deposited_metrics, deposited_displacements = common_heavy_atom_comparison(\n", " atom_array, tmol_atom_array\n", ")\n", "assert deposited_metrics[\"heavy_atom_RMSD_A\"] < 0.05\n", "assert deposited_metrics[\"max_heavy_atom_displacement_A\"] < 0.10\n", "\n", "reuse_diagnostics = StringIO()\n", "try:\n", " with redirect_stdout(reuse_diagnostics), redirect_stderr(reuse_diagnostics):\n", " reused_pose_stack = pose_stack_from_biotite(\n", " tmol_atom_array,\n", " torch_device=device,\n", " context=build_context,\n", " no_optH=True,\n", " )\n", "except Exception:\n", " print(reuse_diagnostics.getvalue())\n", " raise\n", "reused_atom_array = biotite_from_pose_stack(\n", " reused_pose_stack, build_context.canonical_ordering\n", ")\n", "reuse_metrics, reuse_displacements = common_heavy_atom_comparison(\n", " tmol_atom_array, reused_atom_array\n", ")\n", "assert reuse_metrics[\"heavy_atom_RMSD_A\"] < 0.01\n", "assert reuse_metrics[\"max_heavy_atom_displacement_A\"] < 0.02\n", "\n", "roundtrip_path = Path(tempfile.gettempdir()) / \"1ubq_tmol_roundtrip.pdb\"\n", "# PDB cannot encode every TMol/Biotite annotation; that limitation is already\n", "# explained above, so suppress Biotite's duplicate compatibility warning.\n", "with warnings.catch_warnings():\n", " warnings.simplefilter(\"ignore\", UserWarning)\n", " write_pose_stack_pdb(pose_stack, str(roundtrip_path))\n", " roundtrip_array = PDBFile.read(str(roundtrip_path)).get_structure(\n", " model=1,\n", " include_bonds=True,\n", " extra_fields=[\"occupancy\", \"b_factor\"],\n", " )\n", "pdb_metrics, pdb_displacements = common_heavy_atom_comparison(\n", " tmol_atom_array, roundtrip_array\n", ")\n", "assert pdb_metrics[\"heavy_atom_RMSD_A\"] < 0.01\n", "assert pdb_metrics[\"max_heavy_atom_displacement_A\"] < 0.02\n", "\n", "comparison_frame = pd.DataFrame(\n", " [\n", " {\"comparison\": \"deposited model 1 → TMol-built\", **deposited_metrics},\n", " {\"comparison\": \"TMol-built → context-reused TMol-built\", **reuse_metrics},\n", " {\"comparison\": \"TMol-built → PDB compatibility read\", **pdb_metrics},\n", " ]\n", ")\n", "show_table(comparison_frame)\n", "print(\"largest deposited-to-built common-heavy-atom displacements:\")\n", "show_table(deposited_displacements.head(5))\n", "print(\n", " \"context PackedBlockTypes reused:\",\n", " reused_pose_stack.packed_block_types is build_context.packed_block_types,\n", ")\n", "print(\"wrote:\", roundtrip_path.resolve())" ] }, { "cell_type": "markdown", "id": "bae713cf", "metadata": {}, "source": [ "## Direct PDB API, residue slices, and batched output\n", "\n", "`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.\n", "\n", "`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.\n", "\n", "### OpenFold-style prediction tensors\n", "\n", "For canonical proteins, `tmol.pose_stack_from_openfold(result)` consumes `aatype` with shape `[batch, residues]`, final `positions` with shape `[batch, residues, atom14, 3]` (stored by OpenFold under the final recycle), and `chain_index` with shape `[batch, residues]`. Missing supported atoms, including hydrogens, are built differentiably: if the input position tensor requires gradients, a TMol score can backpropagate to the supplied prediction coordinates. Combine this adapter with the per-residue-root forest in Tutorial 05 for NN-like frames.\n", "\n", "See the [Task index](recipe_index.md) and [structure I/O and integrations](../user_guide/integrations.md) for the stable prediction-adapter entry points." ] }, { "cell_type": "code", "execution_count": 6, "id": "9c99d195", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "direct PDB blocks=76; internal slice blocks=6\n", "batch poses=3; multi-model MODEL records=3; separate files=3\n" ] } ], "source": [ "pdb_path = repo_root / \"tmol/tests/data/pdb/1ubq.pdb\"\n", "full_pdb_pose = tmol.pose_stack_from_pdb(str(pdb_path), device=device)\n", "\n", "# Parsed residue positions 19:25 correspond to six residues in this fixture.\n", "# Mark both outer connections incomplete so this is an internal fragment, not\n", "# newly capped N/C termini.\n", "internal_cut_flags = torch.zeros((1, 6, 2), dtype=torch.bool, device=device)\n", "internal_cut_flags[0, 0, 0] = True\n", "internal_cut_flags[0, -1, 1] = True\n", "internal_slice = tmol.pose_stack_from_pdb(\n", " str(pdb_path),\n", " device=device,\n", " residue_start=19,\n", " residue_end=25,\n", " res_not_connected=internal_cut_flags,\n", ")\n", "assert internal_slice.max_n_blocks == 6\n", "assert int(internal_slice.inter_residue_connections[0, 0, 0, 0]) == -1\n", "assert int(internal_slice.inter_residue_connections[0, -1, 1, 0]) == -1\n", "\n", "pdb_batch = PoseStackBuilder.from_poses([full_pdb_pose] * 3, device=device)\n", "with tempfile.TemporaryDirectory(prefix=\"tmol-pdb-output-\") as temp_dir:\n", " temp_dir = Path(temp_dir)\n", " multi_model_path = temp_dir / \"ubiquitin_batch.pdb\"\n", " write_pose_stack_pdb(pdb_batch, str(multi_model_path))\n", " model_count = sum(\n", " line.startswith(\"MODEL \") for line in multi_model_path.read_text().splitlines()\n", " )\n", " separate_paths = []\n", " for pose_index in range(pdb_batch.n_poses):\n", " output_path = temp_dir / f\"ubiquitin_{pose_index:02d}.pdb\"\n", " write_pose_stack_pdb(pdb_batch.split(pose_index), str(output_path))\n", " separate_paths.append(output_path)\n", " assert model_count == pdb_batch.n_poses\n", " assert all(path.is_file() and path.stat().st_size > 0 for path in separate_paths)\n", "\n", "print(\n", " f\"direct PDB blocks={full_pdb_pose.max_n_blocks}; \"\n", " f\"internal slice blocks={internal_slice.max_n_blocks}\"\n", ")\n", "print(\n", " f\"batch poses={pdb_batch.n_poses}; multi-model MODEL records={model_count}; \"\n", " f\"separate files={len(separate_paths)}\"\n", ")" ] }, { "cell_type": "markdown", "id": "f24f7550", "metadata": {}, "source": [ "## Select and visualize deposited atoms\n", "\n", "Biotite stores chain, residue, atom-name, and element annotations as NumPy arrays. Compose Boolean masks directly so the selection is explicit, dependency-free, and easy to test. Here the masks describe `atom_array`, the deposited experimental model; do not apply them to the differently sized `tmol_atom_array` built by TMol.\n", "\n", "Pass `pose_stack` or `tmol_atom_array` to the viewer instead when the scientific question concerns built atoms or resolved chemistry, and recompute highlights in that object's indexing. `tmol.selection_gallery()` accepts an `AtomArray` plus named Boolean masks; click a button to restyle and center the same model." ] }, { "cell_type": "code", "execution_count": 7, "id": "58f6ae03", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "selected backbone atoms: 40\n", "selected residues: [ 1 2 3 4 5 6 7 8 9 10]\n" ] }, { "data": { "text/html": [ "\n", "\n", "\n", "
\n", "
\n", "
\n", "
\n", "
\n", " pick a selection · drag to rotate · scroll to zoom · click a highlighted atom to label it\n", "
\n", "
\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "selection_mask = (\n", " (atom_array.chain_id == \"A\")\n", " & (atom_array.res_id >= 1)\n", " & (atom_array.res_id <= 10)\n", " & np.isin(atom_array.atom_name, [\"N\", \"CA\", \"C\", \"O\"])\n", ")\n", "sidechain_mask = (\n", " (atom_array.chain_id == \"A\")\n", " & (atom_array.res_id >= 1)\n", " & (atom_array.res_id <= 10)\n", " & ~np.isin(atom_array.atom_name, [\"N\", \"CA\", \"C\", \"O\", \"OXT\"])\n", ")\n", "selected_atoms = atom_array[selection_mask]\n", "print(\"selected backbone atoms:\", selected_atoms.array_length())\n", "print(\"selected residues:\", np.unique(selected_atoms.res_id))\n", "try:\n", " display(\n", " tmol.selection_gallery(\n", " atom_array,\n", " {\n", " \"Backbone, residues 1–10\": selection_mask,\n", " \"Side chains, residues 1–10\": sidechain_mask,\n", " \"All Cα atoms\": atom_array.atom_name == \"CA\",\n", " },\n", " width=720,\n", " height=420,\n", " )\n", " )\n", "except ImportError as exc:\n", " print(\"Selection viewer unavailable in this environment:\", exc)" ] }, { "cell_type": "markdown", "id": "3130299a", "metadata": {}, "source": [ "## Rosetta comparison\n", "\n", "A Rosetta `Pose` is a single rich object with residues, conformation, energies, and attached metadata. A TMol `PoseStack` is deliberately batch-first: chemistry/connectivity metadata is block-indexed, while Cartesian coordinates are contiguous PyTorch tensors. One TMol block is closest to a residue-like chemical unit, but not every block must be a canonical amino acid.\n", "\n", "Rosetta applications often obtain behavior through a process-wide flags/options system. TMol has no global Rosetta-style flags layer: device, parameter database, score-function options, I/O choices, and protocol settings are explicit Python arguments or object configuration. This is more verbose but makes notebook state and batching assumptions visible.\n", "\n", "Keep the [Rosetta-to-TMol crosswalk](rosetta_crosswalk.md) open when translating a workflow; it separates genuine API parallels from protocol layers TMol does not implement. See the official full tutorials for [Working with Rosetta](https://docs.rosettacommons.org/demos/latest/tutorials/Working_With_Rosetta/working_with_rosetta), [input and output](https://docs.rosettacommons.org/demos/latest/tutorials/input_and_output/input_and_output), [core concepts](https://docs.rosettacommons.org/demos/latest/tutorials/Core_Concepts/Core_Concepts), [commonly used options](https://docs.rosettacommons.org/demos/latest/tutorials/commonly_used_options/commonly_used_options), and [PyRosetta Pose basics](https://nbviewer.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/02.01-Pose-Basics.ipynb)." ] }, { "cell_type": "markdown", "id": "2f773b88", "metadata": {}, "source": [ "## Next: batch scoring\n", "\n", "This notebook followed one structure from deposited atoms to a TMol-built `PoseStack` and back. [GPU Batching with TMol](02_gpu_batching.ipynb) next keeps those I/O and device choices explicit while assembling many compatible poses for one scoring call; [Scoring and Analysis](03_scoring_and_analysis.ipynb) then interprets the score terms." ] }, { "cell_type": "markdown", "id": "1c5b08c2", "metadata": {}, "source": [ "## Exercises\n", "\n", "1. Change the device selection to force CPU and confirm all tensor devices agree.\n", "2. Select residues 20–30 and heavy side-chain atoms with a Biotite/NumPy Boolean mask; assert the selected author labels.\n", "3. 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.\n", "4. Reuse `build_context` to import the PDB-read `roundtrip_array`; compare its block types with the direct `reused_pose_stack` and explain any differences.\n", "5. 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." ] }, { "cell_type": "markdown", "id": "189319c7", "metadata": {}, "source": [ "## References\n", "\n", "- [TMol Task index](recipe_index.md)\n", "- [Rosetta-to-TMol crosswalk](rosetta_crosswalk.md)\n", "- [TMol repository](https://github.com/uw-ipd/tmol)\n", "- [TMol input/output API](../api/io.rst)\n", "- [Biotite structure documentation](https://www.biotite-python.org/latest/apidoc/biotite.structure.html)\n", "- [Rosetta: Working with Rosetta](https://docs.rosettacommons.org/demos/latest/tutorials/Working_With_Rosetta/working_with_rosetta)\n", "- [Rosetta input/output tutorial](https://docs.rosettacommons.org/demos/latest/tutorials/input_and_output/input_and_output)\n", "- [Rosetta Basics: IO and Scripting workshop slides](https://meilerlab.org/wp-content/uploads/2025/11/Rosetta_IO_and_Scripting.pdf)\n", "- [PyRosetta Pose basics](https://nbviewer.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/02.01-Pose-Basics.ipynb)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }