{ "cells": [ { "cell_type": "markdown", "id": "2993576b", "metadata": {}, "source": [ "# TMol How-To Guide — Advanced Topics\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RosettaCommons/tmol/blob/main/docs/tutorial/08_advanced_topics.ipynb)" ] }, { "cell_type": "markdown", "id": "9e4ac368", "metadata": {}, "source": [ "## 8. Advanced Topics" ] }, { "cell_type": "markdown", "id": "e9396449", "metadata": {}, "source": [ "Two unrelated deep-end topics: customizing which identities the packer offers at a position (beyond the single-target `restrict_absent_name3s` from section 1), and building a `PoseStack` directly from fixed-width atom37 tensors -- the representation most structure-prediction models output -- while keeping the whole thing differentiable." ] }, { "cell_type": "markdown", "id": "ecb35447", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "markdown", "id": "043e6177", "metadata": {}, "source": [ "Wheel install, imports, seeding, and device selection live in [notebook_setup.py](https://github.com/RosettaCommons/tmol/blob/main/docs/tutorial/notebook_setup.py)." ] }, { "cell_type": "code", "execution_count": 1, "id": "8c6870ce", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:49:10.229411Z", "iopub.status.busy": "2026-10-20T00:49:10.229039Z", "iopub.status.idle": "2026-10-20T00:51:18.530593Z", "shell.execute_reply": "2026-10-20T00:51:18.523949Z" }, "tags": [ "collapse-code" ] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Environment variable CCD_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:\n", " (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file\n", " (2) set it in your current shell with 'export VAR_NAME=path/to/variable'\n", " (3) write it to a .env file in the root of the atomworks.io repository\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Environment variable PDB_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:\n", " (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file\n", " (2) set it in your current shell with 'export VAR_NAME=path/to/variable'\n", " (3) write it to a .env file in the root of the atomworks.io repository\n" ] } ], "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 urlretrieve\n", "\n", " urlretrieve(\n", " \"https://raw.githubusercontent.com/RosettaCommons/tmol/\"\n", " \"main/docs/tutorial/notebook_setup.py\",\n", " \"notebook_setup.py\",\n", " )\n", "\n", "from notebook_setup import * # noqa: F401,F403\n", "\n", "device, param_db = setup()\n" ] }, { "cell_type": "markdown", "id": "145a7b46", "metadata": {}, "source": [ "### Extending the packer for design" ] }, { "cell_type": "markdown", "id": "0b3e9313", "metadata": {}, "source": [ "#### PackerPalette" ] }, { "cell_type": "markdown", "id": "3e086a33", "metadata": {}, "source": [ "`PackerPalette` decides, per position, which identities are ever candidates; `restrict_absent_name3s` (section 1) then narrows that set further from outside. Subclass `PackerPalette` and override `block_types_from_original` to change the candidate set itself -- call `super().block_types_from_original(pbt, orig)` for the defaults, then filter or reorder them. Unlike `restrict_absent_name3s`, this applies uniformly everywhere the palette is used, not one position at a time." ] }, { "cell_type": "markdown", "id": "e4ba09ed", "metadata": {}, "source": [ "#### A worked design example" ] }, { "cell_type": "markdown", "id": "6068daf4", "metadata": {}, "source": [ "A palette that only ever offers a small hydrophobic alphabet (plus whatever the position's own original identity is, so a non-matching position can still \"keep its own\"):" ] }, { "cell_type": "code", "execution_count": 2, "id": "035b1d0e", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:51:18.546396Z", "iopub.status.busy": "2026-10-20T00:51:18.545957Z", "iopub.status.idle": "2026-10-20T00:51:18.603502Z", "shell.execute_reply": "2026-10-20T00:51:18.602696Z" } }, "outputs": [], "source": [ "class HydrophobicPalette(PackerPalette):\n", " \"\"\"Allow a small hydrophobic alphabet plus each position's own identity.\"\"\"\n", "\n", " def __init__(self, allowed_name3s=(\"ALA\", \"VAL\", \"LEU\", \"ILE\", \"PHE\")):\n", " super().__init__()\n", " self.allowed_name3s = frozenset(allowed_name3s)\n", "\n", " def block_types_from_original(self, pbt, orig):\n", " n_allowed, allowed, is_original = super().block_types_from_original(pbt, orig)\n", " in_alphabet = torch.tensor(\n", " [block_type.name3 in self.allowed_name3s for block_type in pbt.active_block_types],\n", " dtype=torch.bool, device=pbt.device,\n", " )\n", " real_choice = allowed >= 0\n", " keep = torch.zeros_like(allowed, dtype=torch.bool)\n", " keep[real_choice] = in_alphabet[allowed[real_choice]]\n", " keep = (keep | is_original) & real_choice\n", "\n", " order = torch.argsort((~keep).to(torch.int64), dim=2, stable=True)\n", " allowed = torch.gather(allowed, 2, order)\n", " is_original = torch.gather(is_original, 2, order)\n", " keep = torch.gather(keep, 2, order)\n", " allowed = torch.where(keep, allowed, torch.full_like(allowed, -1))\n", " is_original = is_original & keep\n", " n_allowed = keep.sum(dim=2)\n", " return n_allowed, allowed, is_original\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "670583bf", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:51:18.605331Z", "iopub.status.busy": "2026-10-20T00:51:18.605106Z", "iopub.status.idle": "2026-10-20T00:51:45.903364Z", "shell.execute_reply": "2026-10-20T00:51:45.902323Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n" ] }, { "data": { "application/3dmoljs_load.v0": "
\n

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n
\n", "text/html": [ "
\n", "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", "
\n", "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "block 5: LYS -> LEU\n", "designed identity allowed by the palette: True\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n", "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n", "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n", "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n" ] }, { "data": { "text/html": [ "\n", "
\n", " \n", "
\n", " \n", "
\n", "
block 5: LYS
\n", "
\n", "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", "
\n", "\n", "
\n", "
\n", "
block 5: LEU
\n", "
\n", "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", "
\n", "\n", "
\n", "
\n", "
original vs. designed
\n", "
\n", "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", "
\n", "\n", "
\n", "
\n", "
\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "cif_text = fetch_cif(\"1UBQ\")\n", "pose_stack = pose_stack_from_file(io.StringIO(cif_text), device, param_db=param_db)\n", "\n", "viewer = tmol.view(pose_stack)\n", "viewer.show()\n", "\n", "score_function = beta2016_score_function(device, param_db=param_db)\n", "\n", "target_block = 5\n", "palette = HydrophobicPalette()\n", "design_task = PackerTask(pose_stack, palette)\n", "design_region = torch.zeros_like(design_task.is_real_block)\n", "design_region[:, target_block] = True\n", "design_task.disable_packing_by_block_mask(~design_region)\n", "design_task.add_conformer_sampler(create_dunbrack_sampler_from_database(param_db, device))\n", "design_task.add_conformer_sampler(FixedAAChiSampler())\n", "design_task.add_conformer_sampler(IncludeCurrentSampler())\n", "\n", "torch.manual_seed(SEED)\n", "designed = pack_rotamers(pose_stack, score_function, design_task)\n", "before, after = pose_stack.block_type(0, target_block).name3, designed.block_type(0, target_block).name3\n", "print(f\"block {target_block}: {before} -> {after}\")\n", "print(f\"designed identity allowed by the palette: {after in palette.allowed_name3s or after == before}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"designed\": designed},\n", " notes={\n", " \"original\": f\"block {target_block}: {before}\",\n", " \"designed\": f\"block {target_block}: {after}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "cf4ddd50", "metadata": {}, "source": [ "### Populate a PoseStack from atom37" ] }, { "cell_type": "markdown", "id": "3f368539", "metadata": {}, "source": [ "\"atom37\" is the fixed-width, 37-atom-slot-per-residue layout most structure-prediction models (AlphaFold-style) use. `tmol.io.pose_stack_from_atom37` builds a `PoseStack` directly from such a tensor plus a `PoseBuildContext`, entirely on tensor ops with no `AtomArray` involved. Because of that, gradients flow straight back from the built pose's coordinates to the atom37 input -- a model's output coordinates can be scored by TMol and the loss backpropagated through pose construction with nothing detached along the way.\n", "\n", "We still need an `AtomArray` once here, to resolve chemistry and pull a real structure's canonical-form tensors (`res_types`, `chain_id`, bond records) for re-populating into atom37 slots -- a real caller already has these tensors from its own model and skips this step entirely." ] }, { "cell_type": "code", "execution_count": 4, "id": "c80303cb", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:51:45.913460Z", "iopub.status.busy": "2026-10-20T00:51:45.912538Z", "iopub.status.idle": "2026-10-20T00:51:46.384576Z", "shell.execute_reply": "2026-10-20T00:51:46.375440Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "atom37 tensor shape: (1, 76, 37, 3) (pose, residue, slot, xyz)\n" ] } ], "source": [ "from tmol.io import (\n", " atom_array_from_file, build_context_from_atom_array, canonical_form_from_biotite,\n", " pose_stack_from_atom37, atom37_slot_map_for_ordering,\n", ")\n", "\n", "atom37_cif_text = fetch_cif(\"1UBQ\")\n", "atom37_structure = atom_array_from_file(io.StringIO(atom37_cif_text))\n", "atom37_context = build_context_from_atom_array(atom37_structure, device, param_db=param_db)\n", "canonical_ordering = atom37_context.canonical_ordering\n", "\n", "canonical_form = canonical_form_from_biotite(\n", " atom37_structure, device, co=canonical_ordering, missing_density_distance_threshold=0.0,\n", ")\n", "\n", "\n", "def full_atom_layout(canonical_ordering):\n", " \"\"\"Every canonical atom gets its own atom37 slot, in order -- the simplest possible layout.\"\"\"\n", " return {\n", " name3: list(canonical_ordering.restypes_ordered_atom_names[name3])\n", " for name3 in canonical_ordering.restype_io_equiv_classes\n", " }\n", "\n", "\n", "def scatter_to_slots(coords, res_types, slot_map):\n", " \"\"\"Spread canonical-form coordinates out into atom37 slot order.\"\"\"\n", " n_poses, n_res = res_types.shape\n", " n_slots = slot_map.shape[1]\n", " is_real = res_types >= 0\n", " canonical_atom = slot_map[torch.where(is_real, res_types, torch.zeros_like(res_types)).to(torch.int64)]\n", " routed = (canonical_atom >= 0) & is_real.unsqueeze(-1)\n", " slots = torch.full((n_poses, n_res, n_slots, 3), float(\"nan\"), dtype=torch.float32, device=coords.device)\n", " pose_index = torch.arange(n_poses, device=coords.device).reshape(-1, 1, 1).expand(n_poses, n_res, n_slots)\n", " res_index = torch.arange(n_res, device=coords.device).reshape(1, -1, 1).expand(n_poses, n_res, n_slots)\n", " slots[routed] = coords[pose_index[routed], res_index[routed], canonical_atom[routed]]\n", " return slots\n", "\n", "\n", "slot_map = atom37_slot_map_for_ordering(canonical_ordering, full_atom_layout(canonical_ordering), device)\n", "atom37_coords = scatter_to_slots(canonical_form.coords, canonical_form.res_types, slot_map)\n", "atom37_coords.requires_grad_(True)\n", "print(f\"atom37 tensor shape: {tuple(atom37_coords.shape)} (pose, residue, slot, xyz)\")\n" ] }, { "cell_type": "markdown", "id": "ef00a739", "metadata": {}, "source": [ "`pose_stack_from_atom37` takes the atom37 coordinates, a residue-type tensor, a chain-id tensor, and a `PoseBuildContext` -- the same kind of context object built in notebook 4, here without any ligand." ] }, { "cell_type": "code", "execution_count": 5, "id": "89147319", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:51:46.391338Z", "iopub.status.busy": "2026-10-20T00:51:46.390974Z", "iopub.status.idle": "2026-10-20T00:51:47.671124Z", "shell.execute_reply": "2026-10-20T00:51:47.644754Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "n_poses=1, max_n_blocks=76\n", "pose coordinates require grad: True\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/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.\n", " record = np.char.array(np.where(array.hetero, \"HETATM\", \"ATOM\"))\n" ] }, { "data": { "application/3dmoljs_load.v0": "
\n

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n
\n", "text/html": [ "
\n", "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", "
\n", "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "atom37_pose_stack = pose_stack_from_atom37(\n", " atom37_coords, canonical_form.res_types, canonical_form.chain_id, atom37_context,\n", " slot_map=slot_map, disulfides=canonical_form.disulfides, cyclic_bonds=canonical_form.cyclic_bonds,\n", " covalent_bonds=canonical_form.covalent_bonds, res_not_connected=canonical_form.res_not_connected,\n", ")\n", "print(f\"n_poses={atom37_pose_stack.n_poses}, max_n_blocks={atom37_pose_stack.max_n_blocks}\")\n", "print(f\"pose coordinates require grad: {atom37_pose_stack.coords.requires_grad}\")\n", "\n", "viewer = tmol.view(atom37_pose_stack)\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "d5920401", "metadata": {}, "source": [ "Scoring that pose and calling `.backward()` reaches all the way back to the original atom37 tensor -- the gradient a model's own training loop would need." ] }, { "cell_type": "code", "execution_count": 6, "id": "4375e8e7", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T00:51:47.682744Z", "iopub.status.busy": "2026-10-20T00:51:47.678350Z", "iopub.status.idle": "2026-10-20T00:51:58.520275Z", "shell.execute_reply": "2026-10-20T00:51:58.519208Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "score: 261.03\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "gradient reached the atom37 input: True\n" ] } ], "source": [ "atom37_score_function = beta2016_score_function(device, param_db=atom37_context.parameter_database)\n", "atom37_scorer = atom37_score_function.render_whole_pose_scoring_module(atom37_pose_stack)\n", "atom37_total_score = atom37_scorer(atom37_pose_stack.coords)\n", "print(f\"score: {float(atom37_total_score[0].detach()):.2f}\")\n", "\n", "atom37_total_score.sum().backward()\n", "print(f\"gradient reached the atom37 input: {torch.count_nonzero(atom37_coords.grad).item() > 0}\")\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4" }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.12" }, "tmol": { "gpu_smoke": true } }, "nbformat": 4, "nbformat_minor": 5 }