{ "cells": [ { "cell_type": "markdown", "id": "a3434f0d", "metadata": {}, "source": [ "# TMol How-To Guide — Scoring and Optimization\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/02_scoring_and_optimization.ipynb)" ] }, { "cell_type": "markdown", "id": "8aae5b6b", "metadata": {}, "source": [ "## 2. Scoring and Optimization" ] }, { "cell_type": "markdown", "id": "6e358900", "metadata": {}, "source": [ "Score a structure, minimize it, pack its side chains, and relax it -- the core scoring/optimization loop. This notebook assumes the loading/querying basics from [tmol_how_to_guide.ipynb](01_tmol_how_to_guide.ipynb) section 1." ] }, { "cell_type": "markdown", "id": "471ceff0", "metadata": {}, "source": [ "### Setup (TODO - these extra setup steps will go away when the notebooks are combined)" ] }, { "cell_type": "markdown", "id": "ee57450f", "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": "48adfbb3", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:11.000094Z", "iopub.status.busy": "2026-10-20T01:12:10.999923Z", "iopub.status.idle": "2026-10-20T01:12:17.812709Z", "shell.execute_reply": "2026-10-20T01:12:17.811988Z" }, "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": "3d482b3a", "metadata": {}, "source": [ "We use ubiquitin (1UBQ) as the running example, loaded the same way as in section 1, plus a `ScoreFunction`." ] }, { "cell_type": "code", "execution_count": 2, "id": "2c29efb4", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:17.815071Z", "iopub.status.busy": "2026-10-20T01:12:17.814643Z", "iopub.status.idle": "2026-10-20T01:12:21.971485Z", "shell.execute_reply": "2026-10-20T01:12:21.970899Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "n_poses=1, max_n_blocks=76\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": [ "cif_text = fetch_cif(\"1UBQ\")\n", "pose_stack = pose_stack_from_file(io.StringIO(cif_text), device)\n", "score_function = beta2016_score_function(device, param_db=param_db)\n", "print(f\"n_poses={pose_stack.n_poses}, max_n_blocks={pose_stack.max_n_blocks}\")\n", "\n", "viewer = tmol.view(pose_stack)\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "209316c8", "metadata": {}, "source": [ "### Scoring" ] }, { "cell_type": "markdown", "id": "6bdb22af", "metadata": {}, "source": [ "The tmol score function is based on the Rosetta energy function. The score function is composed of many terms and weights for those terms. The default score function terms and weights are set to match the beta2016_cart score function from Rosetta3.\n", "\n", "Generally, you want to optimize a structure before interpreting its score closely -- a freshly-loaded crystal structure often has small clashes an optimizer would remove." ] }, { "cell_type": "markdown", "id": "f86e5caf", "metadata": {}, "source": [ "#### Scoring a PoseStack" ] }, { "cell_type": "markdown", "id": "622e34e8", "metadata": {}, "source": [ "The simplest way to score a PoseStack is with the whole_pose_scoring_module, which will return a tensor with one score for each structure in the stack `[nposes]`." ] }, { "cell_type": "code", "execution_count": 3, "id": "63500ed2", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:21.973964Z", "iopub.status.busy": "2026-10-20T01:12:21.973807Z", "iopub.status.idle": "2026-10-20T01:12:23.567520Z", "shell.execute_reply": "2026-10-20T01:12:23.566878Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "total score: 201.51\n" ] } ], "source": [ "# Create our score function.\n", "score_function = beta2016_score_function(device, param_db=param_db)\n", "\n", "# Create our scoring module.\n", "whole_scorer = score_function.render_whole_pose_scoring_module(pose_stack)\n", "# The scoring module is what does the actual score evaluation.\n", "# This is separate from the ScoreFunction because it also needs details\n", "# about the PoseStack being scored - mainly the Residue Types being used.\n", "\n", "# Feed in the coordinates to the scorer and print the result\n", "total_score = whole_scorer(pose_stack.coords)\n", "print(f\"total score: {float(total_score[0]):.2f}\")\n" ] }, { "cell_type": "markdown", "id": "8e633787", "metadata": {}, "source": [ "#### Score a PoseStack, get per-term breakdowns" ] }, { "cell_type": "markdown", "id": "accaaa52", "metadata": {}, "source": [ "Pass `sum_terms=False` to get one value per score term instead of the weighted total. This changes the output tensor shape to `[nterms x nposes]`.\n", "\n", "`score_function.all_score_types()` gives the matching term order." ] }, { "cell_type": "code", "execution_count": 4, "id": "b6887721", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:23.569124Z", "iopub.status.busy": "2026-10-20T01:12:23.568955Z", "iopub.status.idle": "2026-10-20T01:12:23.576022Z", "shell.execute_reply": "2026-10-20T01:12:23.575487Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " fa_ljatr: -415.45\n", " fa_lk: 300.99\n", " dunbrack_rotdev: 165.81\n", " lk_ball_iso: -159.97\n", " lk_ball: 158.11\n", " fa_elec: -146.66\n", " fa_ljrep: 126.12\n", " cart_angles: 91.90\n" ] } ], "source": [ "weighted_terms = whole_scorer(pose_stack.coords, sum_terms=False, apply_weights=True)\n", "\n", "# Get the score types so that we can label them with the names\n", "score_types = score_function.all_score_types()\n", "\n", "# Combine the scores and names and then print\n", "term_breakdown = {st.name: float(weighted_terms[i, 0]) for i, st in enumerate(score_types)}\n", "for name, value in sorted(term_breakdown.items(), key=lambda kv: -abs(kv[1]))[:8]:\n", " print(f\"{name:>16}: {value:9.2f}\")\n" ] }, { "cell_type": "markdown", "id": "c819becc", "metadata": {}, "source": [ "#### Setting score function weights" ] }, { "cell_type": "markdown", "id": "78df6d34", "metadata": {}, "source": [ "`beta2016_score_function` is just a preset `ScoreFunction`; terms are added or dropped lazily by setting a weight on their `ScoreType`. Setting every weight that covers a term back to zero drops it." ] }, { "cell_type": "code", "execution_count": 5, "id": "a1ca09b5", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:23.577429Z", "iopub.status.busy": "2026-10-20T01:12:23.577276Z", "iopub.status.idle": "2026-10-20T01:12:23.595905Z", "shell.execute_reply": "2026-10-20T01:12:23.595373Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "fa_ljrep + hbond only: 66.10\n" ] } ], "source": [ "focused = ScoreFunction(param_db, device)\n", "focused.set_weight(ScoreType.fa_ljrep, 0.55)\n", "focused.set_weight(ScoreType.hbond, 1.0)\n", "focused_scorer = focused.render_whole_pose_scoring_module(pose_stack)\n", "print(f\"fa_ljrep + hbond only: {float(focused_scorer(pose_stack.coords)[0]):.2f}\")\n" ] }, { "cell_type": "markdown", "id": "4f39941e", "metadata": {}, "source": [ "#### Block-pair scoring" ] }, { "cell_type": "markdown", "id": "81e18b39", "metadata": {}, "source": [ "The 'whole_pose' scoring module above only returns a single energy for each pose. \n", "\n", "You can alternatively use 'block-pair' scoring to get the energies between any pair of residues (or a residue and itself for 1-body energies).\n", "\n", "Block-pair scoring will return a tensor of shape `[nposes x nres x nres]`, or `[nterms x nposes x nres x nres]` if you ask for individual term scores. " ] }, { "cell_type": "code", "execution_count": 6, "id": "7ebf921b", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:23.600415Z", "iopub.status.busy": "2026-10-20T01:12:23.600106Z", "iopub.status.idle": "2026-10-20T01:12:23.640474Z", "shell.execute_reply": "2026-10-20T01:12:23.639825Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Score: -0.44379860162734985\n" ] } ], "source": [ "block_pair_scorer = score_function.render_block_pair_scoring_module(pose_stack)\n", "block_pair_scores = block_pair_scorer(pose_stack.coords)\n", "\n", "# Print the energy between residues 4 and 5 on pose 0\n", "print(f\"Score: {block_pair_scores[0, 4, 5]}\")" ] }, { "cell_type": "markdown", "id": "caede24b", "metadata": {}, "source": [ "#### Interface energies" ] }, { "cell_type": "markdown", "id": "f41139c3", "metadata": {}, "source": [ "We briefly switch to a two-chain example, [1BL8](https://www.rcsb.org/structure/1BL8), to score the interface between its chains A and B. Chain masks like the ones below are the general boolean-selection mechanism that the Selections notebook covers in depth." ] }, { "cell_type": "code", "execution_count": 7, "id": "82f2eb43", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:23.642176Z", "iopub.status.busy": "2026-10-20T01:12:23.642007Z", "iopub.status.idle": "2026-10-20T01:12:27.178287Z", "shell.execute_reply": "2026-10-20T01:12:27.177395Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "chain A: 100 blocks, chain B: 97 blocks\n", "A-B interface score: 83.06\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": [ "interface_cif_text = fetch_cif(\"1BL8\")\n", "interface_pose_stack = pose_stack_from_file(io.StringIO(interface_cif_text), device)\n", "\n", "chain_a = get_chain_mask(interface_pose_stack, \"A\")\n", "chain_b = get_chain_mask(interface_pose_stack, \"B\")\n", "\n", "interface_score = calculate_block_pair_ddg(\n", " interface_pose_stack, chain_a, chain_b,\n", " sfxn=score_function, minimize=False, pack=False,\n", ")\n", "print(f\"chain A: {int(chain_a.sum())} blocks, chain B: {int(chain_b.sum())} blocks\")\n", "print(f\"A-B interface score: {float(interface_score[0]):.2f}\")\n", "\n", "chain_a_atom_mask = res_mask_to_coord_mask(interface_pose_stack, chain_a)\n", "viewer = tmol.view(interface_pose_stack, highlighted=chain_a_atom_mask[interface_pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "3d9d41fc", "metadata": {}, "source": [ "#### Per-residue energies" ] }, { "cell_type": "markdown", "id": "7ab92434", "metadata": {}, "source": [ "`tmol.ops.per_residue_block_pair_energies` splits each block-pair score evenly between its two partner residues, so the per-residue sum reconciles exactly with the whole-pose total. It only reduces over the block dimensions, so a per-term (`sum_terms=False`) input comes back with its term dimension still unsummed." ] }, { "cell_type": "code", "execution_count": 8, "id": "27b78822", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:27.187546Z", "iopub.status.busy": "2026-10-20T01:12:27.187305Z", "iopub.status.idle": "2026-10-20T01:12:27.190811Z", "shell.execute_reply": "2026-10-20T01:12:27.190251Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "residue score for residue 10 3.28\n" ] } ], "source": [ "per_residue_scores = per_residue_block_pair_energies(block_pair_scores)\n", "\n", "print(f\"residue score for residue 10 {float(per_residue_scores[0, 10]):7.2f}\")\n" ] }, { "cell_type": "markdown", "id": "7a4e6ef4", "metadata": {}, "source": [ "### Minimization" ] }, { "cell_type": "markdown", "id": "977d11e7", "metadata": {}, "source": [ "tmol provides 2 different methods of minimizing energies - cartesian and torsion-space minimzation." ] }, { "cell_type": "markdown", "id": "9f22be95", "metadata": {}, "source": [ "#### Cartesian-space minimization" ] }, { "cell_type": "markdown", "id": "72fda8fc", "metadata": {}, "source": [ "Cartesian space minimization can be run with `tmol.run_cart_min(pose_stack, score_function)`. This will compute gradients for the XYZ coordinates of each atom directly, and will move these XYZ coordinates based on those gradients." ] }, { "cell_type": "code", "execution_count": 9, "id": "868d1fe1", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:27.193146Z", "iopub.status.busy": "2026-10-20T01:12:27.192941Z", "iopub.status.idle": "2026-10-20T01:12:30.917726Z", "shell.execute_reply": "2026-10-20T01:12:30.917038Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after: -235.66\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: -235.66
\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. cartesian minimized
\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": [ "# Minimize the pose_stack\n", "cart_minimized = tmol.run_cart_min(pose_stack.clone(), score_function)\n", "\n", "# Score it and print the result.\n", "cart_minimized_score = score_function.render_whole_pose_scoring_module(cart_minimized)(cart_minimized.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after: {float(cart_minimized_score[0].detach()):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"cartesian minimized\": cart_minimized},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"cartesian minimized\": f\"score: {float(cart_minimized_score[0].detach()):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "1096ca19", "metadata": {}, "source": [ "#### Torsion-space minimization" ] }, { "cell_type": "markdown", "id": "5e9caedb", "metadata": {}, "source": [ "Torsion-space minimization moves internal degrees of freedom (bond torsions) instead of raw xyz coordinates, following a `FoldForest` -- the kinematic tree describing how motion at one bond propagates through the rest of the structure. It reuses the same coordinate gradients as above, reparameterized through that tree; the Advanced Kinematics notebook covers FoldForests in depth, so here we just build the default one.\n", "\n", "`tmol.MoveMap` says which internal degrees of freedom are allowed to move; `move_all_named_torsions = True` with jumps left off moves every named torsion (phi/psi/chi) and nothing else." ] }, { "cell_type": "code", "execution_count": 10, "id": "1d311881", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:30.921725Z", "iopub.status.busy": "2026-10-20T01:12:30.921320Z", "iopub.status.idle": "2026-10-20T01:12:34.541047Z", "shell.execute_reply": "2026-10-20T01:12:34.540326Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after (torsion space): 173.53\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score (torsion space): 173.53
\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. torsion minimized
\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": [ "fold_forest = tmol.FoldForest.reasonable_fold_forest(pose_stack)\n", "move_map = tmol.MoveMap.from_pose_stack(pose_stack)\n", "move_map.move_all_named_torsions = True\n", "move_map.move_all_jumps = False\n", "move_map.move_all_root_jumps = False\n", "\n", "kin_minimized = tmol.run_kin_min(\n", " pose_stack.clone(), score_function, fold_forest, move_map, optimizer_kwargs={\"max_iter\": 10}\n", ")\n", "kin_minimized_score = score_function.render_whole_pose_scoring_module(kin_minimized)(kin_minimized.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after (torsion space): {float(kin_minimized_score[0].detach()):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"torsion minimized\": kin_minimized},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"torsion minimized\": f\"score (torsion space): {float(kin_minimized_score[0].detach()):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "db057fef", "metadata": {}, "source": [ "#### Gradients" ] }, { "cell_type": "markdown", "id": "d35b179c", "metadata": {}, "source": [ "After scoring a structure, you can retreive the gradients directly, if you wish to use them outside the minimizer. " ] }, { "cell_type": "code", "execution_count": 11, "id": "ce11223d", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:34.544794Z", "iopub.status.busy": "2026-10-20T01:12:34.544630Z", "iopub.status.idle": "2026-10-20T01:12:34.593506Z", "shell.execute_reply": "2026-10-20T01:12:34.592881Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Cartesian\n", "A 72 ARG O: 103.499 per A\n", "A 72 ARG HG2: 97.927 per A\n", "A 13 ILE N: 72.700 per A\n", "A 13 ILE CB: 61.983 per A\n", "A 13 ILE CA: 51.459 per A\n", "A 71 LEU HA: 39.167 per A\n", "A 71 LEU HD12: 38.615 per A\n", "A 3 ILE C: 37.888 per A\n", "A 15 LEU CA: 37.106 per A\n", "A 15 LEU N: 35.677 per A\n", "Torsion\n", "A 39 ASP omega: -388.844 per radian\n", "A 37 PRO phi: 388.389 per radian\n", "A 39 ASP psi: -356.287 per radian\n", "A 40 GLN phi: -355.432 per radian\n", "A 36 ILE psi: 353.364 per radian\n", "A 63 LYS phi: 338.571 per radian\n", "A 57 SER psi: -317.802 per radian\n", "A 34 GLU psi: 298.237 per radian\n", "A 63 LYS omega: 297.719 per radian\n", "A 62 GLN psi: 290.887 per radian\n" ] } ], "source": [ "# Cartesian gradients\n", "print(\"Cartesian\")\n", "coords = pose_stack.coords.detach().clone().requires_grad_(True)\n", "whole_scorer(coords).sum().backward()\n", "atom_grad = torch.nan_to_num(coords.grad[0]).norm(dim=-1)\n", "for atom in atom_grad.topk(10).indices.tolist():\n", " print(f\"{atom_label(pose_stack, 0, atom)}: {float(atom_grad[atom]):.3f} per A\")\n", "\n", "# Torsion gradients: forward kinematics (dofs -> coords) feeding the scorer\n", "print(\"Torsion\")\n", "kin_module = PoseStackKinematicsModule(\n", " pose_stack, tmol.FoldForest.reasonable_fold_forest(pose_stack)\n", ")\n", "kin_network = KinForestSfxnNetwork(score_function, pose_stack, kin_module)\n", "kin_network().sum().backward()\n", "dof_grad = torch.zeros_like(kin_network.full_dofs)\n", "dof_grad[kin_network.dof_mask] = kin_network.masked_dofs.grad\n", "\n", "labels, nodes = torsion_labels(pose_stack, kin_module, 0)\n", "torsion_grad = dof_grad[nodes.to(dof_grad.device), BondDOFTypes.phi_c]\n", "for i in torsion_grad.abs().topk(10).indices.tolist():\n", " print(f\"{labels[i]}: {float(torsion_grad[i]):.3f} per radian\")" ] }, { "cell_type": "markdown", "id": "2cd2ed7f", "metadata": {}, "source": [ "#### Controlling minimization with movemaps" ] }, { "cell_type": "markdown", "id": "a93ecd4a", "metadata": {}, "source": [ "A movemap can restrict minimization to part of a structure. `tmol.ops.res_mask_to_coord_mask` expands a block-level mask to the atom-level mask `CartesianMoveMap` needs; `MoveMap`'s `set_move_all_{mc,sc}_tors_for_blocks` setters take the block mask directly." ] }, { "cell_type": "markdown", "id": "65370fc6", "metadata": {}, "source": [ "#### Cartesian space" ] }, { "cell_type": "code", "execution_count": 12, "id": "0a7772fa", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:34.595953Z", "iopub.status.busy": "2026-10-20T01:12:34.595701Z", "iopub.status.idle": "2026-10-20T01:12:35.136893Z", "shell.execute_reply": "2026-10-20T01:12:35.136052Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "max movement, frozen half: 0.0000 A\n", "max movement, movable half: 1.1327 A\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", "
unminimized
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
max movement frozen half: 0.0000 A, movable half: 1.1327 A
\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. partial cartesian min
\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": [ "second_half_mask = torch.ones((1, pose_stack.max_n_blocks), dtype=torch.bool, device=device)\n", "second_half_mask[:, :38] = False # freeze the first half of the chain\n", "movable_coords = res_mask_to_coord_mask(pose_stack, second_half_mask)\n", "\n", "cartesian_move_map = tmol.CartesianMoveMap(coord_mask=movable_coords)\n", "partial_cart_min = tmol.run_cart_min(\n", " pose_stack.clone(), score_function,\n", " coord_mask=cartesian_move_map.coord_mask,\n", " optimizer_kwargs={\"max_iter\": 10},\n", ")\n", "moved = (partial_cart_min.coords - pose_stack.coords).norm(dim=-1)\n", "print(f\"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A\")\n", "print(f\"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"partial cartesian min\": partial_cart_min},\n", " notes={\n", " \"original\": \"unminimized\",\n", " \"partial cartesian min\": (\n", " f\"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, \"\n", " f\"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\"\n", " ),\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "4d937187", "metadata": {}, "source": [ "#### Torsion space" ] }, { "cell_type": "code", "execution_count": 13, "id": "1c0a947b", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:35.141130Z", "iopub.status.busy": "2026-10-20T01:12:35.140964Z", "iopub.status.idle": "2026-10-20T01:12:35.750697Z", "shell.execute_reply": "2026-10-20T01:12:35.750107Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "max movement, frozen half: 0.0000 A\n", "max movement, movable half: 1.2537 A\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", "
unminimized
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
max movement frozen half: 0.0000 A, movable half: 1.2537 A
\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. partial torsion min
\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": [ "partial_move_map = tmol.MoveMap.from_pose_stack(pose_stack)\n", "partial_move_map.set_move_all_mc_tors_for_blocks(second_half_mask)\n", "partial_move_map.set_move_all_sc_tors_for_blocks(second_half_mask)\n", "\n", "partial_kin_min = tmol.run_kin_min(\n", " pose_stack.clone(), score_function, fold_forest, partial_move_map, optimizer_kwargs={\"max_iter\": 10}\n", ")\n", "moved = (partial_kin_min.coords.detach() - pose_stack.coords).norm(dim=-1)\n", "print(f\"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A\")\n", "print(f\"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"partial torsion min\": partial_kin_min},\n", " notes={\n", " \"original\": \"unminimized\",\n", " \"partial torsion min\": (\n", " f\"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, \"\n", " f\"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\"\n", " ),\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "b15353af", "metadata": {}, "source": [ "### Packing" ] }, { "cell_type": "markdown", "id": "c0746360", "metadata": {}, "source": [ "#### Defining packer behavior with a PackerTask" ] }, { "cell_type": "markdown", "id": "604a01b7", "metadata": {}, "source": [ "To run packing, we must create a `PackerTask`. This controls what rotamers are to be considered at which positions." ] }, { "cell_type": "code", "execution_count": 14, "id": "3265d647", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:35.754743Z", "iopub.status.busy": "2026-10-20T01:12:35.754585Z", "iopub.status.idle": "2026-10-20T01:12:35.768692Z", "shell.execute_reply": "2026-10-20T01:12:35.768032Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "packer task ready: every position repackable\n" ] } ], "source": [ "default_task = PackerTask(pose_stack, PackerPalette())\n", "default_task.restrict_to_repacking()\n", "default_task.add_conformer_sampler(FixedAAChiSampler())\n", "default_task.add_conformer_sampler(IncludeCurrentSampler())\n", "print(\"packer task ready: every position repackable\")\n" ] }, { "cell_type": "markdown", "id": "411dbcc2", "metadata": {}, "source": [ "#### Run packing TODO - pick an example that actually shows a difference from packing" ] }, { "cell_type": "markdown", "id": "2a849a05", "metadata": {}, "source": [ "`pack_rotamers` runs the actual search and returns a new `PoseStack`." ] }, { "cell_type": "code", "execution_count": 15, "id": "358d07e2", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:35.770413Z", "iopub.status.busy": "2026-10-20T01:12:35.770239Z", "iopub.status.idle": "2026-10-20T01:12:36.430342Z", "shell.execute_reply": "2026-10-20T01:12:36.429447Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after repack: 201.24\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: 201.24
\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. repacked
\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": [ "torch.manual_seed(SEED)\n", "repacked = pack_rotamers(pose_stack, score_function, default_task)\n", "repacked_score = score_function.render_whole_pose_scoring_module(repacked)(repacked.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after repack: {float(repacked_score[0]):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"repacked\": repacked},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"repacked\": f\"score: {float(repacked_score[0]):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "e5e3bdb8", "metadata": {}, "source": [ "`OptHSampler` samples proton chis and amide/ring flips only, leaving every heavy atom at its input position -- add it (never alongside a Dunbrack sampler on the same position) for hydrogen-only packing." ] }, { "cell_type": "code", "execution_count": 16, "id": "db0fc0d2", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:36.434584Z", "iopub.status.busy": "2026-10-20T01:12:36.434427Z", "iopub.status.idle": "2026-10-20T01:12:37.401573Z", "shell.execute_reply": "2026-10-20T01:12:37.400876Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after optH-only pack: 201.24\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: 201.24
\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. optH-only packed
\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": [ "opth_task = PackerTask(pose_stack, PackerPalette())\n", "opth_task.restrict_to_repacking()\n", "opth_task.add_conformer_sampler(IncludeCurrentSampler())\n", "opth_task.add_conformer_sampler(OptHSampler())\n", "opth_task.add_conformer_sampler(FixedAAChiSampler())\n", "\n", "torch.manual_seed(SEED)\n", "opth_packed = pack_rotamers(pose_stack, score_function, opth_task)\n", "opth_score = score_function.render_whole_pose_scoring_module(opth_packed)(opth_packed.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after optH-only pack: {float(opth_score[0]):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"optH-only packed\": opth_packed},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"optH-only packed\": f\"score: {float(opth_score[0]):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "f7b3a69e", "metadata": {}, "source": [ "#### Controlling packing at the residue level" ] }, { "cell_type": "markdown", "id": "c9be90ef", "metadata": {}, "source": [ "`disable_packing_by_block_mask` restricts the search to specific positions; everything else keeps its input conformation untouched." ] }, { "cell_type": "code", "execution_count": 17, "id": "2603293a", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:37.405370Z", "iopub.status.busy": "2026-10-20T01:12:37.405211Z", "iopub.status.idle": "2026-10-20T01:12:38.120498Z", "shell.execute_reply": "2026-10-20T01:12:38.119598Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "max movement outside block 10: 0.0000 A\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 10 unpacked
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
max movement outside block 10: 0.0000 A
\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. single-residue packed
\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": [ "single_residue_task = PackerTask(pose_stack, PackerPalette())\n", "single_residue_task.restrict_to_repacking()\n", "shell = torch.zeros_like(single_residue_task.is_real_block)\n", "shell[0, 10] = True\n", "single_residue_task.disable_packing_by_block_mask(~shell)\n", "single_residue_task.add_conformer_sampler(create_dunbrack_sampler_from_database(param_db, device))\n", "single_residue_task.add_conformer_sampler(FixedAAChiSampler())\n", "single_residue_task.add_conformer_sampler(IncludeCurrentSampler())\n", "\n", "torch.manual_seed(SEED)\n", "single_packed = pack_rotamers(pose_stack, score_function, single_residue_task)\n", "moved = (single_packed.coords - pose_stack.coords).norm(dim=-1)[0]\n", "untouched = torch.ones(pose_stack.max_n_blocks, dtype=torch.bool)\n", "untouched[10] = False\n", "print(f\"max movement outside block 10: {float(torch.nan_to_num(moved[res_mask_to_coord_mask(pose_stack, untouched.unsqueeze(0))[0]]).max()):.4f} A\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"single-residue packed\": single_packed},\n", " notes={\n", " \"original\": \"block 10 unpacked\",\n", " \"single-residue packed\": (\n", " \"max movement outside block 10: \"\n", " f\"{float(torch.nan_to_num(moved[res_mask_to_coord_mask(pose_stack, untouched.unsqueeze(0))[0]]).max()):.4f} A\"\n", " ),\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "48b92d6d", "metadata": {}, "source": [ "### Relax" ] }, { "cell_type": "markdown", "id": "119e87d6", "metadata": {}, "source": [ "`tmol.fast_relax` alternates packing and minimization over a multi-stage schedule, instead of running either once. It needs a `PackerPalette`, a movemap, and a `FoldForest`, all already introduced above." ] }, { "cell_type": "markdown", "id": "3b5a1532", "metadata": {}, "source": [ "#### Relax options (cycles, schedule)" ] }, { "cell_type": "markdown", "id": "a9dd5f8e", "metadata": {}, "source": [ "`num_repeats` controls how many times the full schedule runs; `schedule` is a list of stages, each ramping `fa_rep` (and, if constraints are active, their weight) from mild to full strength -- packing near-clashing structures directly at full repulsive weight tends to get stuck. Omit `schedule` for the built-in default schedule." ] }, { "cell_type": "code", "execution_count": 18, "id": "bbab9a11", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:38.124936Z", "iopub.status.busy": "2026-10-20T01:12:38.124755Z", "iopub.status.idle": "2026-10-20T01:12:38.129047Z", "shell.execute_reply": "2026-10-20T01:12:38.128651Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "a 2-stage schedule; the default schedule has 4 stages\n" ] } ], "source": [ "from tmol.relax import DEFAULT_RELAX_SCHEDULE\n", "\n", "quick_schedule = [\n", " {\"fa_rep_pack_frac\": 0.10, \"fa_rep_min_frac\": 0.20, \"cst_frac\": 1.0},\n", " {\"fa_rep_pack_frac\": 1.00, \"fa_rep_min_frac\": 1.00, \"cst_frac\": 0.0},\n", "]\n", "print(f\"a 2-stage schedule; the default schedule has {len(DEFAULT_RELAX_SCHEDULE)} stages\")\n" ] }, { "cell_type": "markdown", "id": "05a09a40", "metadata": {}, "source": [ "#### Cartesian relax" ] }, { "cell_type": "markdown", "id": "47eacbfa", "metadata": {}, "source": [ "Passing a `CartesianMoveMap` relaxes in cartesian space." ] }, { "cell_type": "code", "execution_count": 19, "id": "d3dd3bf2", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:38.131206Z", "iopub.status.busy": "2026-10-20T01:12:38.130942Z", "iopub.status.idle": "2026-10-20T01:12:47.229833Z", "shell.execute_reply": "2026-10-20T01:12:47.229210Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after cartesian relax: -266.98\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: -266.98
\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. cartesian relaxed
\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": [ "palette = PackerPalette()\n", "full_cartesian_move_map = tmol.CartesianMoveMap()\n", "\n", "torch.manual_seed(SEED)\n", "cart_relaxed = tmol.fast_relax(\n", " pose_stack.clone(), score_function, palette, full_cartesian_move_map, fold_forest,\n", " num_repeats=1, schedule=quick_schedule,\n", ")\n", "cart_relaxed_score = score_function.render_whole_pose_scoring_module(cart_relaxed)(cart_relaxed.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after cartesian relax: {float(cart_relaxed_score[0].detach()):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"cartesian relaxed\": cart_relaxed},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"cartesian relaxed\": f\"score: {float(cart_relaxed_score[0].detach()):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "a87e4c76", "metadata": {}, "source": [ "#### Torsion relax" ] }, { "cell_type": "markdown", "id": "0a5cb5cd", "metadata": {}, "source": [ "Passing a `MoveMap` instead relaxes in torsion space -- but `fast_relax`'s minimization step always runs `run_cart_min` internally unless told otherwise, so a torsion-space relax also needs a `min_fn` callable wrapping `tmol.run_kin_min` with our fold forest and move map." ] }, { "cell_type": "code", "execution_count": 20, "id": "ce2b1c30", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:47.237082Z", "iopub.status.busy": "2026-10-20T01:12:47.236917Z", "iopub.status.idle": "2026-10-20T01:12:55.320258Z", "shell.execute_reply": "2026-10-20T01:12:55.319589Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after torsion relax: -167.26\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: -167.26
\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. torsion relaxed
\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": [ "def kin_min_fn(ps, sfxn, *, fold_forest, move_map, verbose):\n", " return tmol.run_kin_min(ps, sfxn, fold_forest, move_map, verbose=verbose)\n", "\n", "\n", "torch.manual_seed(SEED)\n", "kin_relaxed = tmol.fast_relax(\n", " pose_stack.clone(), score_function, palette, move_map, fold_forest,\n", " num_repeats=1, schedule=quick_schedule, min_fn=kin_min_fn,\n", ")\n", "kin_relaxed_score = score_function.render_whole_pose_scoring_module(kin_relaxed)(kin_relaxed.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after torsion relax: {float(kin_relaxed_score[0].detach()):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"torsion relaxed\": kin_relaxed},\n", " notes={\n", " \"original\": f\"score: {float(total_score[0]):.2f}\",\n", " \"torsion relaxed\": f\"score: {float(kin_relaxed_score[0].detach()):.2f}\",\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "b0fca13e", "metadata": {}, "source": [ "#### Relax with a movemap" ] }, { "cell_type": "markdown", "id": "1946a6df", "metadata": {}, "source": [ "The movemap restricts *minimization* as before, but `fast_relax` also repacks every cycle regardless of the movemap -- freezing a region completely needs `task_operations` to disable packing there too." ] }, { "cell_type": "code", "execution_count": 21, "id": "fb611063", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:12:55.325070Z", "iopub.status.busy": "2026-10-20T01:12:55.324837Z", "iopub.status.idle": "2026-10-20T01:13:01.290532Z", "shell.execute_reply": "2026-10-20T01:13:01.289686Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "max movement, frozen half: 0.0000 A\n", "max movement, movable half: 2.7246 A\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", "
unrelaxed
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
max movement frozen half: 0.0000 A, movable half: 2.7246 A
\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. movemap-restricted relax
\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": [ "restricted_move_map = tmol.CartesianMoveMap(coord_mask=movable_coords)\n", "\n", "\n", "def freeze_first_half(task):\n", " task.disable_packing_by_block_mask(~second_half_mask)\n", "\n", "\n", "torch.manual_seed(SEED)\n", "restricted_relaxed = tmol.fast_relax(\n", " pose_stack.clone(), score_function, palette, restricted_move_map, fold_forest,\n", " num_repeats=1, schedule=quick_schedule, task_operations=[freeze_first_half],\n", ")\n", "moved = (restricted_relaxed.coords - pose_stack.coords).norm(dim=-1)\n", "print(f\"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A\")\n", "print(f\"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\")\n", "\n", "display(tmol.overlay_view(\n", " {\"original\": pose_stack, \"movemap-restricted relax\": restricted_relaxed},\n", " notes={\n", " \"original\": \"unrelaxed\",\n", " \"movemap-restricted relax\": (\n", " f\"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, \"\n", " f\"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A\"\n", " ),\n", " },\n", "))\n" ] }, { "cell_type": "markdown", "id": "20f567d5", "metadata": {}, "source": [ "#### Relax with constraints" ] }, { "cell_type": "markdown", "id": "d828265d", "metadata": {}, "source": [ "`create_mainchain_coordinate_constraints` attaches a harmonic restraint on every mainchain atom to its current position; `fast_relax` ramps the `constraint` weight down over the schedule (`cst_frac`) as long as that `ScoreType`'s weight is nonzero." ] }, { "cell_type": "code", "execution_count": 22, "id": "e67410f1", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T01:13:01.295992Z", "iopub.status.busy": "2026-10-20T01:13:01.295826Z", "iopub.status.idle": "2026-10-20T01:13:10.727875Z", "shell.execute_reply": "2026-10-20T01:13:10.727326Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/runner/work/tmol/tmol/tmol/score/constraint/_constraint_energy_term.py:469: UserWarning: Sparse invariant checks are implicitly disabled. Memory errors (e.g. SEGFAULT) will occur when operating on a sparse tensor which violates the invariants, but checks incur performance overhead. To silence this warning, explicitly opt in or out. See `torch.sparse.check_sparse_tensor_invariants.__doc__` for guidance. (Triggered internally at /__w/pytorch/pytorch/aten/src/ATen/Context.cpp:848.)\n", " sparse_scores = torch.sparse_coo_tensor(\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "before: 201.51 after constrained relax: -268.72\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", "
score: 201.51
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
score: -268.72
\n", "
\n", "

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

\n", "
\n", "\n", "
\n", "
\n", "
constrained start vs. constrained relaxed
\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": [ "score_function.set_weight(ScoreType.constraint, 1.0)\n", "constrained_start = create_mainchain_coordinate_constraints(pose_stack.clone())\n", "\n", "torch.manual_seed(SEED)\n", "constrained_relaxed = tmol.fast_relax(\n", " constrained_start, score_function, palette, full_cartesian_move_map, fold_forest,\n", " num_repeats=1, schedule=quick_schedule, ramp_constraints=True,\n", ")\n", "constrained_score = score_function.render_whole_pose_scoring_module(constrained_relaxed)(constrained_relaxed.coords)\n", "print(f\"before: {float(total_score[0]):.2f} after constrained relax: {float(constrained_score[0].detach()):.2f}\")\n", "\n", "display(tmol.overlay_view(\n", " {\"constrained start\": constrained_start, \"constrained relaxed\": constrained_relaxed},\n", " notes={\n", " \"constrained start\": f\"score: {float(total_score[0]):.2f}\",\n", " \"constrained relaxed\": f\"score: {float(constrained_score[0].detach()):.2f}\",\n", " },\n", "))\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 }