TMol How-To Guide — Selections#
5. Selections#
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.
Setup#
Wheel install, imports, seeding, and device selection live in notebook_setup.py.
[1]:
try:
import google.colab # noqa: F401
except ImportError:
IN_COLAB = False
else:
IN_COLAB = True
if IN_COLAB:
from urllib.request import urlretrieve
urlretrieve(
"https://raw.githubusercontent.com/RosettaCommons/tmol/"
"main/docs/tutorial/notebook_setup.py",
"notebook_setup.py",
)
from notebook_setup import * # noqa: F401,F403
device, param_db = setup()
Environment variable CCD_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
(1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
(2) set it in your current shell with 'export VAR_NAME=path/to/variable'
(3) write it to a .env file in the root of the atomworks.io repository
Environment variable PDB_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
(1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
(2) set it in your current shell with 'export VAR_NAME=path/to/variable'
(3) write it to a .env file in the root of the atomworks.io repository
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 (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.
[2]:
cif_text = fetch_cif("1BL8")
pose_stack = pose_stack_from_file(io.StringIO(cif_text), device, param_db=param_db)
print(f"n_poses={pose_stack.n_poses}, max_n_blocks={pose_stack.max_n_blocks}")
viewer = tmol.view(pose_stack)
viewer.show()
n_poses=1, max_n_blocks=391
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.
Select by chain#
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.)
[3]:
chain_labels = pose_stack.pdb_info.chain_labels[0]
chain_a = get_chain_mask(pose_stack, "A")
print(f"unique chains: {sorted(set(chain_labels.tolist()))}")
print(f"chain A blocks: {int(chain_a.sum())}")
chain_a_atom_mask = res_mask_to_coord_mask(pose_stack, chain_a)
viewer = tmol.view(pose_stack, highlighted=chain_a_atom_mask[pose_stack.real_atoms])
viewer.show()
unique chains: ['A', 'B', 'C', 'D']
chain A blocks: 100
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.
Select by backbone / side chain#
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.
[4]:
sidechain_atoms = build_sidechain_coord_mask(pose_stack)
backbone_atoms = pose_stack.real_atoms & ~sidechain_atoms
print(f"sidechain atoms: {int(sidechain_atoms.sum())}")
print(f"backbone atoms: {int(backbone_atoms.sum())}")
viewer = tmol.view(pose_stack, highlighted=sidechain_atoms[pose_stack.real_atoms])
viewer.show()
sidechain atoms: 4727
backbone atoms: 1164
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.
Select by polymer type#
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.
[5]:
protein_blocks = tmol.get_polymer_type_mask(pose_stack, "amino_acid")
ion_blocks = tmol.get_polymer_type_mask(pose_stack, None)
print(f"protein blocks: {int(protein_blocks.sum())}, ion blocks: {int(ion_blocks.sum())}")
ion_atom_mask = res_mask_to_coord_mask(pose_stack, ion_blocks)
viewer = tmol.view(pose_stack, highlighted=ion_atom_mask[pose_stack.real_atoms])
viewer.show()
protein blocks: 388, ion blocks: 3
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.
Select within k Å of a reference#
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.
[6]:
from tmol.io import atom_array_from_file, pose_stack_from_atom_array
def atoms_within(pose_stack, reference_atom_mask, distance_angstroms, pose=0):
"""Real atoms within `distance_angstroms` of any atom in `reference_atom_mask`."""
all_coords = torch.nan_to_num(pose_stack.coords[pose])
reference_coords = pose_stack.coords[pose][reference_atom_mask[pose]]
nearest_distance = torch.cdist(all_coords, reference_coords).min(dim=1).values
return (nearest_distance <= distance_angstroms) & pose_stack.real_atoms[pose]
# A ligand: benzamidine (BEN) bound to trypsin, from section 1/4.
ligand_pose_stack = pose_stack_from_atom_array(
atom_array_from_file(io.StringIO(fetch_cif("3PTB"))), device, prepare_ligands=True, param_db=param_db
)
ligand_blocks = torch.zeros((1, ligand_pose_stack.max_n_blocks), dtype=torch.bool)
for block in range(ligand_pose_stack.max_n_blocks):
if ligand_pose_stack.block_type(0, block).name3 == "BEN":
ligand_blocks[0, block] = True
near_ligand = atoms_within(ligand_pose_stack, res_mask_to_coord_mask(ligand_pose_stack, ligand_blocks), 5.0)
print(f"atoms within 5 A of the ligand: {int(near_ligand.sum())}")
viewer = tmol.view(ligand_pose_stack, highlighted=near_ligand[ligand_pose_stack.real_atoms[0]])
viewer.show()
==============================
*** Open Babel Warning in WriteMolecule
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.
atoms within 5 A of the ligand: 145
/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
warnings.warn(
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.
Multiple criteria#
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.
[7]:
chain_b = get_chain_mask(pose_stack, "B")
ion_coord_mask = res_mask_to_coord_mask(pose_stack, ion_blocks)
near_ions_8A = atoms_within(pose_stack, ion_coord_mask, 8.0)
near_chain_b_5A = atoms_within(pose_stack, res_mask_to_coord_mask(pose_stack, chain_b), 5.0)
combined = near_ions_8A | near_chain_b_5A
print(f"within 8 A of the bound ions: {int(near_ions_8A.sum())}")
print(f"within 5 A of chain B: {int(near_chain_b_5A.sum())}")
print(f"either condition: {int(combined.sum())}")
viewer = tmol.view(pose_stack, highlighted=combined[pose_stack.real_atoms[0]])
viewer.show()
within 8 A of the bound ions: 325
within 5 A of chain B: 2070
either condition: 2200
/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.
record = np.char.array(np.where(array.hetero, "HETATM", "ATOM"))
3Dmol.js failed to load for some reason. Please check your browser console for error messages.