{ "cells": [ { "cell_type": "markdown", "id": "a134b6d0", "metadata": {}, "source": [ "# TMol How-To Guide — Selections\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/05_selections.ipynb)" ] }, { "cell_type": "markdown", "id": "a46c8e3c", "metadata": {}, "source": [ "## 5. Selections" ] }, { "cell_type": "markdown", "id": "f9c854e2", "metadata": {}, "source": [ "Every selection in TMol is a plain boolean `torch.Tensor` -- block-shaped (`[n_poses, max_n_blocks]`) or atom-shaped (`[n_poses, max_n_atoms]`) -- built by hand from whatever a `PoseStack` already exposes. There's no dedicated \"Selection\" class or query language, just tensors and boolean algebra; this notebook collects the common patterns in one place." ] }, { "cell_type": "markdown", "id": "0dbb0add", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "markdown", "id": "c452efb8", "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": "fe69dc3b", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:15.549729Z", "iopub.status.busy": "2026-10-20T17:52:15.549459Z", "iopub.status.idle": "2026-10-20T17:52:20.551707Z", "shell.execute_reply": "2026-10-20T17:52:20.551010Z" }, "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": "4d92a968", "metadata": {}, "source": [ "`pose_stack.pdb_info` is where most chain/residue-identity selections start: parallel `[pose, block]` arrays `chain_labels`, `residue_labels`, and `residue_insertion_codes`, plus per-atom `atom_occupancy` and `atom_b_factor`. We use [1BL8](https://www.rcsb.org/structure/1BL8) (the protein-protein interface example from section 2) as the running example, for its multiple chains and three bound potassium ions as a second polymer type." ] }, { "cell_type": "code", "execution_count": 2, "id": "e2b80e53", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:20.553998Z", "iopub.status.busy": "2026-10-20T17:52:20.553538Z", "iopub.status.idle": "2026-10-20T17:52:28.761595Z", "shell.execute_reply": "2026-10-20T17:52:28.760860Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "n_poses=1, max_n_blocks=391\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(\"1BL8\")\n", "pose_stack = pose_stack_from_file(io.StringIO(cif_text), 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": "d1c9736c", "metadata": {}, "source": [ "### Select by chain" ] }, { "cell_type": "markdown", "id": "0199ce4e", "metadata": {}, "source": [ "There is no chain-selection helper; compare `pdb_info.chain_labels` directly. (`pose_stack.chain_id64` gives an integer per-pose chain index, if comparing against a label string is inconvenient.)" ] }, { "cell_type": "code", "execution_count": 3, "id": "c1054449", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:28.771869Z", "iopub.status.busy": "2026-10-20T17:52:28.771595Z", "iopub.status.idle": "2026-10-20T17:52:28.863525Z", "shell.execute_reply": "2026-10-20T17:52:28.862937Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "unique chains: ['A', 'B', 'C', 'D']\n", "chain A blocks: 100\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": [ "chain_labels = pose_stack.pdb_info.chain_labels[0]\n", "chain_a = get_chain_mask(pose_stack, \"A\")\n", "print(f\"unique chains: {sorted(set(chain_labels.tolist()))}\")\n", "print(f\"chain A blocks: {int(chain_a.sum())}\")\n", "\n", "chain_a_atom_mask = res_mask_to_coord_mask(pose_stack, chain_a)\n", "viewer = tmol.view(pose_stack, highlighted=chain_a_atom_mask[pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "e0f7f102", "metadata": {}, "source": [ "### Select by backbone / side chain" ] }, { "cell_type": "markdown", "id": "42b5a89d", "metadata": {}, "source": [ "`build_sidechain_coord_mask` is atom-shaped (`[n_poses, max_n_atoms]`) and excludes each residue type's declared main-chain atoms (non-polymers count as all side chain). There's no separate backbone function -- it's the complement, restricted to real atoms." ] }, { "cell_type": "code", "execution_count": 4, "id": "78e3cd92", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:28.874296Z", "iopub.status.busy": "2026-10-20T17:52:28.874119Z", "iopub.status.idle": "2026-10-20T17:52:28.973543Z", "shell.execute_reply": "2026-10-20T17:52:28.972944Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "sidechain atoms: 4727\n", "backbone atoms: 1164\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": [ "sidechain_atoms = build_sidechain_coord_mask(pose_stack)\n", "backbone_atoms = pose_stack.real_atoms & ~sidechain_atoms\n", "print(f\"sidechain atoms: {int(sidechain_atoms.sum())}\")\n", "print(f\"backbone atoms: {int(backbone_atoms.sum())}\")\n", "\n", "viewer = tmol.view(pose_stack, highlighted=sidechain_atoms[pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "87952c67", "metadata": {}, "source": [ "### Select by polymer type" ] }, { "cell_type": "markdown", "id": "4641bfdc", "metadata": {}, "source": [ "Each block type declares `properties.polymer.polymer_type` (`\"amino_acid\"`, `\"nucleic_acid\"`, or `None` for a non-polymer like a ligand or ion), alongside the `is_polymer` flag already used in section 1. `tmol.get_polymer_type_mask` gathers this across every block in one vectorized pass instead of walking them in Python." ] }, { "cell_type": "code", "execution_count": 5, "id": "6b780d74", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:28.984651Z", "iopub.status.busy": "2026-10-20T17:52:28.984373Z", "iopub.status.idle": "2026-10-20T17:52:29.074223Z", "shell.execute_reply": "2026-10-20T17:52:29.073393Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "protein blocks: 388, ion blocks: 3\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": [ "protein_blocks = tmol.get_polymer_type_mask(pose_stack, \"amino_acid\")\n", "ion_blocks = tmol.get_polymer_type_mask(pose_stack, None)\n", "print(f\"protein blocks: {int(protein_blocks.sum())}, ion blocks: {int(ion_blocks.sum())}\")\n", "\n", "ion_atom_mask = res_mask_to_coord_mask(pose_stack, ion_blocks)\n", "viewer = tmol.view(pose_stack, highlighted=ion_atom_mask[pose_stack.real_atoms])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "39de2745", "metadata": {}, "source": [ "### Select within *k* Å of a reference" ] }, { "cell_type": "markdown", "id": "d1eff8bd", "metadata": {}, "source": [ "There is no \"within distance\" utility either -- the real pattern throughout TMol's tutorials is a plain `torch.cdist` between all atom coordinates and the reference atoms' coordinates, thresholded. We write it once and reuse it for a ligand." ] }, { "cell_type": "code", "execution_count": 6, "id": "05628120", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:29.084287Z", "iopub.status.busy": "2026-10-20T17:52:29.084069Z", "iopub.status.idle": "2026-10-20T17:52:32.323169Z", "shell.execute_reply": "2026-10-20T17:52:32.322407Z" } }, "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": [ "atoms within 5 A of the ligand: 145\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", "\n", "def atoms_within(pose_stack, reference_atom_mask, distance_angstroms, pose=0):\n", " \"\"\"Real atoms within `distance_angstroms` of any atom in `reference_atom_mask`.\"\"\"\n", " all_coords = torch.nan_to_num(pose_stack.coords[pose])\n", " reference_coords = pose_stack.coords[pose][reference_atom_mask[pose]]\n", " nearest_distance = torch.cdist(all_coords, reference_coords).min(dim=1).values\n", " return (nearest_distance <= distance_angstroms) & pose_stack.real_atoms[pose]\n", "\n", "\n", "# A ligand: benzamidine (BEN) bound to trypsin, from section 1/4.\n", "ligand_pose_stack = pose_stack_from_atom_array(\n", " atom_array_from_file(io.StringIO(fetch_cif(\"3PTB\"))), device, prepare_ligands=True, param_db=param_db\n", ")\n", "ligand_blocks = torch.zeros((1, ligand_pose_stack.max_n_blocks), dtype=torch.bool)\n", "for block in range(ligand_pose_stack.max_n_blocks):\n", " if ligand_pose_stack.block_type(0, block).name3 == \"BEN\":\n", " ligand_blocks[0, block] = True\n", "near_ligand = atoms_within(ligand_pose_stack, res_mask_to_coord_mask(ligand_pose_stack, ligand_blocks), 5.0)\n", "print(f\"atoms within 5 A of the ligand: {int(near_ligand.sum())}\")\n", "\n", "viewer = tmol.view(ligand_pose_stack, highlighted=near_ligand[ligand_pose_stack.real_atoms[0]])\n", "viewer.show()\n" ] }, { "cell_type": "markdown", "id": "520033c1", "metadata": {}, "source": [ "### Multiple criteria" ] }, { "cell_type": "markdown", "id": "9d3dd0b5", "metadata": {}, "source": [ "Combining selections is plain boolean algebra on the tensors -- `|` for \"or\", `&` for \"and\", `~` for \"not\" -- no special combinator needed. Here we combine two distance criteria: within 8 Å of the bound ions, or within 5 Å of the second protein chain." ] }, { "cell_type": "code", "execution_count": 7, "id": "cfd556f9", "metadata": { "execution": { "iopub.execute_input": "2026-10-20T17:52:32.328680Z", "iopub.status.busy": "2026-10-20T17:52:32.328509Z", "iopub.status.idle": "2026-10-20T17:52:32.427776Z", "shell.execute_reply": "2026-10-20T17:52:32.427144Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "within 8 A of the bound ions: 325\n", "within 5 A of chain B: 2070\n", "either condition: 2200\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": [ "chain_b = get_chain_mask(pose_stack, \"B\")\n", "ion_coord_mask = res_mask_to_coord_mask(pose_stack, ion_blocks)\n", "\n", "near_ions_8A = atoms_within(pose_stack, ion_coord_mask, 8.0)\n", "near_chain_b_5A = atoms_within(pose_stack, res_mask_to_coord_mask(pose_stack, chain_b), 5.0)\n", "combined = near_ions_8A | near_chain_b_5A\n", "\n", "print(f\"within 8 A of the bound ions: {int(near_ions_8A.sum())}\")\n", "print(f\"within 5 A of chain B: {int(near_chain_b_5A.sum())}\")\n", "print(f\"either condition: {int(combined.sum())}\")\n", "\n", "viewer = tmol.view(pose_stack, highlighted=combined[pose_stack.real_atoms[0]])\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 }