{ "cells": [ { "cell_type": "markdown", "id": "1afec672", "metadata": {}, "source": [ "# TMol How-To Guide — Constraints\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/06_constraints.ipynb)" ] }, { "cell_type": "markdown", "id": "b2162287", "metadata": {}, "source": [ "## 6. Constraints" ] }, { "cell_type": "markdown", "id": "7dd743d4", "metadata": {}, "source": [ "Constraints are extra harmonic (or similar) restraint terms attached to a `PoseStack` and scored like any other energy term, at weight 0 by default. Section 2 already used one ready-made constraint generator; this notebook covers the general mechanism and a few more targeted examples." ] }, { "cell_type": "markdown", "id": "2de0598e", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "markdown", "id": "efb17d3b", "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": "c941c1bd", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:29.616523Z", "iopub.status.busy": "2026-10-19T23:48:29.616277Z", "iopub.status.idle": "2026-10-19T23:48:41.211921Z", "shell.execute_reply": "2026-10-19T23:48:41.210973Z" }, "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": "012dfc0d", "metadata": {}, "source": [ "### Generating constraints programmatically" ] }, { "cell_type": "markdown", "id": "c17e5542", "metadata": {}, "source": [ "`ConstraintSet.create_empty(device, n_poses)` starts empty; `.add_constraints(constraint_function, atoms, params)` returns a new `ConstraintSet` with one more kind of restraint added. `ConstraintEnergyTerm.harmonic` is the simplest constraint function -- a single harmonic distance restraint between two atoms, `params[:, 0]` the target distance and `params[:, 1]` its standard deviation -- and `atoms` is `[constraint, atom_slot, (pose, block, atom)]` for each side." ] }, { "cell_type": "code", "execution_count": 2, "id": "dae7e9b1", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:41.215646Z", "iopub.status.busy": "2026-10-19T23:48:41.215205Z", "iopub.status.idle": "2026-10-19T23:48:47.393705Z", "shell.execute_reply": "2026-10-19T23:48:47.392638Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "constraints defined: 1\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, param_db=param_db)\n", "\n", "residue_3, residue_4 = pose_stack.block_type(0, 3), pose_stack.block_type(0, 4)\n", "constraint_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)\n", "constraint_atoms[0, 0] = torch.tensor([0, 3, residue_3.atom_to_idx[\"C\"]])\n", "constraint_atoms[0, 1] = torch.tensor([0, 4, residue_4.atom_to_idx[\"N\"]])\n", "constraint_params = torch.tensor([[1.47, 0.1]], dtype=torch.float32, device=device)\n", "\n", "distance_constraints = ConstraintSet.create_empty(device=device, n_poses=pose_stack.n_poses).add_constraints(\n", " ConstraintEnergyTerm.harmonic, constraint_atoms, constraint_params\n", ")\n", "print(f\"constraints defined: {distance_constraints.constraint_atoms.shape[0]}\")\n", "\n", "constraint_atom_mask = torch.zeros_like(pose_stack.real_atoms)\n", "for _, block, atom in constraint_atoms[0].tolist():\n", " offset = int(pose_stack.block_coord_offset[0, block])\n", " constraint_atom_mask[0, offset + atom] = True\n", "\n", "viewer = tmol.view(pose_stack, highlighted=constraint_atom_mask[pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "b4cad40c", "metadata": {}, "source": [ "### Attach constraints to a pose" ] }, { "cell_type": "markdown", "id": "102c1766", "metadata": {}, "source": [ "A `ConstraintSet` is a field on `PoseStack` (`pose_stack.constraint_set`), set via `attrs.evolve` since `PoseStack` is immutable. Once attached, it scores automatically through the ordinary `constraint` `ScoreType` -- no separate argument to scoring calls -- as long as that term's weight is nonzero." ] }, { "cell_type": "code", "execution_count": 3, "id": "ce03b83d", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:47.396735Z", "iopub.status.busy": "2026-10-19T23:48:47.396508Z", "iopub.status.idle": "2026-10-19T23:48:50.972258Z", "shell.execute_reply": "2026-10-19T23:48:50.971405Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "constraint energy: 1.987\n" ] }, { "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" ] } ], "source": [ "import attrs\n", "\n", "constrained_pose = attrs.evolve(pose_stack, constraint_set=distance_constraints)\n", "\n", "score_function = beta2016_score_function(device, param_db=param_db)\n", "score_function.set_weight(ScoreType.constraint, 1.0)\n", "scorer = score_function.render_whole_pose_scoring_module(constrained_pose)\n", "weighted_terms = scorer(constrained_pose.coords, sum_terms=False, apply_weights=True)\n", "constraint_index = score_function.all_score_types().index(ScoreType.constraint)\n", "print(f\"constraint energy: {float(weighted_terms[constraint_index, 0]):.3f}\")\n" ] }, { "cell_type": "markdown", "id": "d0a3cd64", "metadata": {}, "source": [ "### Examples" ] }, { "cell_type": "markdown", "id": "b89d1d2a", "metadata": {}, "source": [ "#### Backbone-only coordinate constraints" ] }, { "cell_type": "markdown", "id": "ba9ac843", "metadata": {}, "source": [ "`create_mainchain_coordinate_constraints` (used already in section 2) restrains every declared main-chain atom of every polymer block to its current coordinate; `constrain_all_ca` is the sparser version, one restraint per CA only." ] }, { "cell_type": "code", "execution_count": 4, "id": "74070153", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:50.974422Z", "iopub.status.busy": "2026-10-19T23:48:50.974123Z", "iopub.status.idle": "2026-10-19T23:48:51.228641Z", "shell.execute_reply": "2026-10-19T23:48:51.219982Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "mainchain constraints: 228\n", "CA-only constraints: 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": [ "mainchain_constrained = create_mainchain_coordinate_constraints(pose_stack.clone())\n", "ca_only_constrained = constrain_all_ca(pose_stack.clone())\n", "print(f\"mainchain constraints: {mainchain_constrained.constraint_set.constraint_atoms.shape[0]}\")\n", "print(f\"CA-only constraints: {ca_only_constrained.constraint_set.constraint_atoms.shape[0]}\")\n", "\n", "ca_constraint_mask = torch.zeros_like(ca_only_constrained.real_atoms)\n", "for pose, block, atom in ca_only_constrained.constraint_set.constraint_atoms[:, 0, :].tolist():\n", " offset = int(ca_only_constrained.block_coord_offset[pose, block])\n", " ca_constraint_mask[pose, offset + atom] = True\n", "\n", "viewer = tmol.view(ca_only_constrained, highlighted=ca_constraint_mask[ca_only_constrained.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "f4210314", "metadata": {}, "source": [ "#### PPI interface pairwise distance constraints" ] }, { "cell_type": "markdown", "id": "041cd059", "metadata": {}, "source": [ "[1BL8](https://www.rcsb.org/structure/1BL8) is the two-chain interface example from section 2. We find the closest CA-CA pair across chains A and B and restrain it to its current distance -- a common way to hold an interface together while the rest of a structure is sampled or minimized." ] }, { "cell_type": "code", "execution_count": 5, "id": "559b4ddc", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:51.230596Z", "iopub.status.busy": "2026-10-19T23:48:51.230407Z", "iopub.status.idle": "2026-10-19T23:48:56.258107Z", "shell.execute_reply": "2026-10-19T23:48:56.257338Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "closest CA-CA pair: blocks 88 (A) and 182 (B), 4.17 A apart\n", "PPI interface constraint attached\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, param_db=param_db)\n", "\n", "def protein_blocks_with_ca(pose_stack, chain):\n", " chain_mask = get_chain_mask(pose_stack, chain)[0]\n", " blocks = []\n", " for block in torch.nonzero(chain_mask).flatten().tolist():\n", " block_type = pose_stack.block_type(0, block)\n", " if \"CA\" in block_type.atom_to_idx:\n", " blocks.append(block)\n", " return blocks\n", "\n", "\n", "blocks_a = protein_blocks_with_ca(interface_pose_stack, \"A\")\n", "blocks_b = protein_blocks_with_ca(interface_pose_stack, \"B\")\n", "\n", "\n", "def ca_coord(pose_stack, block):\n", " block_type = pose_stack.block_type(0, block)\n", " offset = int(pose_stack.block_coord_offset[0, block])\n", " return pose_stack.coords[0, offset + block_type.atom_to_idx[\"CA\"]]\n", "\n", "\n", "ca_coords_a = torch.stack([ca_coord(interface_pose_stack, b) for b in blocks_a])\n", "ca_coords_b = torch.stack([ca_coord(interface_pose_stack, b) for b in blocks_b])\n", "pairwise_distance = torch.cdist(ca_coords_a, ca_coords_b)\n", "closest = int(torch.argmin(pairwise_distance))\n", "closest_a, closest_b = blocks_a[closest // len(blocks_b)], blocks_b[closest % len(blocks_b)]\n", "current_distance = float(pairwise_distance.min())\n", "print(f\"closest CA-CA pair: blocks {closest_a} (A) and {closest_b} (B), {current_distance:.2f} A apart\")\n", "\n", "ppi_type_a, ppi_type_b = interface_pose_stack.block_type(0, closest_a), interface_pose_stack.block_type(0, closest_b)\n", "ppi_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)\n", "ppi_atoms[0, 0] = torch.tensor([0, closest_a, ppi_type_a.atom_to_idx[\"CA\"]])\n", "ppi_atoms[0, 1] = torch.tensor([0, closest_b, ppi_type_b.atom_to_idx[\"CA\"]])\n", "ppi_params = torch.tensor([[current_distance, 0.5]], dtype=torch.float32, device=device)\n", "\n", "ppi_constraints = ConstraintSet.create_empty(device=device, n_poses=1).add_constraints(\n", " ConstraintEnergyTerm.harmonic, ppi_atoms, ppi_params\n", ")\n", "ppi_constrained_pose = attrs.evolve(interface_pose_stack, constraint_set=ppi_constraints)\n", "print(\"PPI interface constraint attached\")\n", "\n", "ppi_atom_mask = torch.zeros_like(interface_pose_stack.real_atoms)\n", "offset_a = int(interface_pose_stack.block_coord_offset[0, closest_a])\n", "offset_b = int(interface_pose_stack.block_coord_offset[0, closest_b])\n", "ppi_atom_mask[0, offset_a + ppi_type_a.atom_to_idx[\"CA\"]] = True\n", "ppi_atom_mask[0, offset_b + ppi_type_b.atom_to_idx[\"CA\"]] = True\n", "\n", "viewer = tmol.view(ppi_constrained_pose, highlighted=ppi_atom_mask[interface_pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "6652a025", "metadata": {}, "source": [ "#### Protein-ligand interface pairwise distance constraints" ] }, { "cell_type": "markdown", "id": "ca578048", "metadata": {}, "source": [ "Same mechanism, between a protein atom and a ligand atom. [3PTB](https://www.rcsb.org/structure/3PTB)'s benzamidine binds trypsin's S1 pocket via a salt bridge from its amidinium group to a nearby aspartate or glutamate carboxylate; we find that closest contact and restrain it." ] }, { "cell_type": "code", "execution_count": 6, "id": "db69ab43", "metadata": { "execution": { "iopub.execute_input": "2026-10-19T23:48:56.272669Z", "iopub.status.busy": "2026-10-19T23:48:56.272374Z", "iopub.status.idle": "2026-10-19T23:49:00.288322Z", "shell.execute_reply": "2026-10-19T23:49:00.287503Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "==============================\n", "*** Open Babel Warning in WriteMolecule\n", " No 2D or 3D coordinates exist. Stereochemical information will be stored using an Open Babel extension. To generate 2D or 3D coordinates instead use --gen2D or --gen3D.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "closest ligand-carboxylate contact: HN21 to ASP block 170 OD2, 1.97 A\n", "protein-ligand interface constraint attached\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/runner/work/tmol/tmol/tmol/io/_pose_stack_from_biotite.py:1401: UserWarning: Renumbering chain A of the output AtomArray: biotite's CIF writer does not support numbering that decreases within a chain\n", " warnings.warn(\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": { "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": [ "from tmol.io import atom_array_from_file, pose_stack_from_atom_array\n", "\n", "pl_cif_text = fetch_cif(\"3PTB\")\n", "pl_pose_stack, pl_build_context = pose_stack_from_atom_array(\n", " atom_array_from_file(io.StringIO(pl_cif_text)), device, prepare_ligands=True, param_db=param_db, return_context=True\n", ")\n", "\n", "ligand_block = next(\n", " block for block in range(pl_pose_stack.max_n_blocks)\n", " if pl_pose_stack.block_type(0, block).name3 == \"BEN\"\n", ")\n", "carboxylate_blocks = [\n", " block for block in range(pl_pose_stack.max_n_blocks)\n", " if pl_pose_stack.block_type(0, block).name3 in (\"ASP\", \"GLU\")\n", "]\n", "\n", "ligand_type = pl_pose_stack.block_type(0, ligand_block)\n", "ligand_offset = int(pl_pose_stack.block_coord_offset[0, ligand_block])\n", "ligand_coords = pl_pose_stack.coords[0, ligand_offset : ligand_offset + len(ligand_type.atoms)]\n", "\n", "best = None\n", "for block in carboxylate_blocks:\n", " block_type = pl_pose_stack.block_type(0, block)\n", " offset = int(pl_pose_stack.block_coord_offset[0, block])\n", " block_coords = pl_pose_stack.coords[0, offset : offset + len(block_type.atoms)]\n", " distances = torch.cdist(ligand_coords, block_coords)\n", " nearest = float(distances.min())\n", " if best is None or nearest < best[0]:\n", " flat_index = int(distances.argmin())\n", " ligand_atom, protein_atom = divmod(flat_index, distances.shape[1])\n", " best = (nearest, block, ligand_atom, protein_atom)\n", "\n", "distance, closest_block, ligand_atom_index, protein_atom_index = best\n", "closest_type = pl_pose_stack.block_type(0, closest_block)\n", "print(\n", " f\"closest ligand-carboxylate contact: {ligand_type.atoms[ligand_atom_index].name} to \"\n", " f\"{closest_type.name3} block {closest_block} {closest_type.atoms[protein_atom_index].name}, {distance:.2f} A\"\n", ")\n", "\n", "pl_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)\n", "pl_atoms[0, 0] = torch.tensor([0, ligand_block, ligand_atom_index])\n", "pl_atoms[0, 1] = torch.tensor([0, closest_block, protein_atom_index])\n", "pl_params = torch.tensor([[distance, 0.5]], dtype=torch.float32, device=device)\n", "\n", "pl_constraints = ConstraintSet.create_empty(device=device, n_poses=1).add_constraints(\n", " ConstraintEnergyTerm.harmonic, pl_atoms, pl_params\n", ")\n", "pl_constrained_pose = attrs.evolve(pl_pose_stack, constraint_set=pl_constraints)\n", "print(\"protein-ligand interface constraint attached\")\n", "\n", "protein_offset = int(pl_pose_stack.block_coord_offset[0, closest_block])\n", "\n", "pl_atom_mask = torch.zeros_like(pl_pose_stack.real_atoms)\n", "pl_atom_mask[0, ligand_offset + ligand_atom_index] = True\n", "pl_atom_mask[0, protein_offset + protein_atom_index] = True\n", "\n", "viewer = tmol.view(pl_constrained_pose, highlighted=pl_atom_mask[pl_pose_stack.real_atoms])\n", "viewer.show()\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 }