TMol How-To Guide — Constraints#
6. Constraints#
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.
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
Generating constraints programmatically#
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.
[2]:
cif_text = fetch_cif("1UBQ")
pose_stack = pose_stack_from_file(io.StringIO(cif_text), device, param_db=param_db)
residue_3, residue_4 = pose_stack.block_type(0, 3), pose_stack.block_type(0, 4)
constraint_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)
constraint_atoms[0, 0] = torch.tensor([0, 3, residue_3.atom_to_idx["C"]])
constraint_atoms[0, 1] = torch.tensor([0, 4, residue_4.atom_to_idx["N"]])
constraint_params = torch.tensor([[1.47, 0.1]], dtype=torch.float32, device=device)
distance_constraints = ConstraintSet.create_empty(device=device, n_poses=pose_stack.n_poses).add_constraints(
ConstraintEnergyTerm.harmonic, constraint_atoms, constraint_params
)
print(f"constraints defined: {distance_constraints.constraint_atoms.shape[0]}")
constraint_atom_mask = torch.zeros_like(pose_stack.real_atoms)
for _, block, atom in constraint_atoms[0].tolist():
offset = int(pose_stack.block_coord_offset[0, block])
constraint_atom_mask[0, offset + atom] = True
viewer = tmol.view(pose_stack, highlighted=constraint_atom_mask[pose_stack.real_atoms])
viewer.show()
constraints defined: 1
/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.
Attach constraints to a pose#
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.
[3]:
import attrs
constrained_pose = attrs.evolve(pose_stack, constraint_set=distance_constraints)
score_function = beta2016_score_function(device, param_db=param_db)
score_function.set_weight(ScoreType.constraint, 1.0)
scorer = score_function.render_whole_pose_scoring_module(constrained_pose)
weighted_terms = scorer(constrained_pose.coords, sum_terms=False, apply_weights=True)
constraint_index = score_function.all_score_types().index(ScoreType.constraint)
print(f"constraint energy: {float(weighted_terms[constraint_index, 0]):.3f}")
constraint energy: 1.987
/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.)
sparse_scores = torch.sparse_coo_tensor(
Examples#
Backbone-only coordinate constraints#
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.
[4]:
mainchain_constrained = create_mainchain_coordinate_constraints(pose_stack.clone())
ca_only_constrained = constrain_all_ca(pose_stack.clone())
print(f"mainchain constraints: {mainchain_constrained.constraint_set.constraint_atoms.shape[0]}")
print(f"CA-only constraints: {ca_only_constrained.constraint_set.constraint_atoms.shape[0]}")
ca_constraint_mask = torch.zeros_like(ca_only_constrained.real_atoms)
for pose, block, atom in ca_only_constrained.constraint_set.constraint_atoms[:, 0, :].tolist():
offset = int(ca_only_constrained.block_coord_offset[pose, block])
ca_constraint_mask[pose, offset + atom] = True
viewer = tmol.view(ca_only_constrained, highlighted=ca_constraint_mask[ca_only_constrained.real_atoms])
viewer.show()
mainchain constraints: 228
CA-only constraints: 76
/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.
PPI interface pairwise distance constraints#
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.
[5]:
interface_cif_text = fetch_cif("1BL8")
interface_pose_stack = pose_stack_from_file(io.StringIO(interface_cif_text), device, param_db=param_db)
def protein_blocks_with_ca(pose_stack, chain):
chain_mask = get_chain_mask(pose_stack, chain)[0]
blocks = []
for block in torch.nonzero(chain_mask).flatten().tolist():
block_type = pose_stack.block_type(0, block)
if "CA" in block_type.atom_to_idx:
blocks.append(block)
return blocks
blocks_a = protein_blocks_with_ca(interface_pose_stack, "A")
blocks_b = protein_blocks_with_ca(interface_pose_stack, "B")
def ca_coord(pose_stack, block):
block_type = pose_stack.block_type(0, block)
offset = int(pose_stack.block_coord_offset[0, block])
return pose_stack.coords[0, offset + block_type.atom_to_idx["CA"]]
ca_coords_a = torch.stack([ca_coord(interface_pose_stack, b) for b in blocks_a])
ca_coords_b = torch.stack([ca_coord(interface_pose_stack, b) for b in blocks_b])
pairwise_distance = torch.cdist(ca_coords_a, ca_coords_b)
closest = int(torch.argmin(pairwise_distance))
closest_a, closest_b = blocks_a[closest // len(blocks_b)], blocks_b[closest % len(blocks_b)]
current_distance = float(pairwise_distance.min())
print(f"closest CA-CA pair: blocks {closest_a} (A) and {closest_b} (B), {current_distance:.2f} A apart")
ppi_type_a, ppi_type_b = interface_pose_stack.block_type(0, closest_a), interface_pose_stack.block_type(0, closest_b)
ppi_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)
ppi_atoms[0, 0] = torch.tensor([0, closest_a, ppi_type_a.atom_to_idx["CA"]])
ppi_atoms[0, 1] = torch.tensor([0, closest_b, ppi_type_b.atom_to_idx["CA"]])
ppi_params = torch.tensor([[current_distance, 0.5]], dtype=torch.float32, device=device)
ppi_constraints = ConstraintSet.create_empty(device=device, n_poses=1).add_constraints(
ConstraintEnergyTerm.harmonic, ppi_atoms, ppi_params
)
ppi_constrained_pose = attrs.evolve(interface_pose_stack, constraint_set=ppi_constraints)
print("PPI interface constraint attached")
ppi_atom_mask = torch.zeros_like(interface_pose_stack.real_atoms)
offset_a = int(interface_pose_stack.block_coord_offset[0, closest_a])
offset_b = int(interface_pose_stack.block_coord_offset[0, closest_b])
ppi_atom_mask[0, offset_a + ppi_type_a.atom_to_idx["CA"]] = True
ppi_atom_mask[0, offset_b + ppi_type_b.atom_to_idx["CA"]] = True
viewer = tmol.view(ppi_constrained_pose, highlighted=ppi_atom_mask[interface_pose_stack.real_atoms])
viewer.show()
closest CA-CA pair: blocks 88 (A) and 182 (B), 4.17 A apart
PPI interface constraint attached
/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.
Protein-ligand interface pairwise distance constraints#
Same mechanism, between a protein atom and a ligand atom. 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.
[6]:
from tmol.io import atom_array_from_file, pose_stack_from_atom_array
pl_cif_text = fetch_cif("3PTB")
pl_pose_stack, pl_build_context = pose_stack_from_atom_array(
atom_array_from_file(io.StringIO(pl_cif_text)), device, prepare_ligands=True, param_db=param_db, return_context=True
)
ligand_block = next(
block for block in range(pl_pose_stack.max_n_blocks)
if pl_pose_stack.block_type(0, block).name3 == "BEN"
)
carboxylate_blocks = [
block for block in range(pl_pose_stack.max_n_blocks)
if pl_pose_stack.block_type(0, block).name3 in ("ASP", "GLU")
]
ligand_type = pl_pose_stack.block_type(0, ligand_block)
ligand_offset = int(pl_pose_stack.block_coord_offset[0, ligand_block])
ligand_coords = pl_pose_stack.coords[0, ligand_offset : ligand_offset + len(ligand_type.atoms)]
best = None
for block in carboxylate_blocks:
block_type = pl_pose_stack.block_type(0, block)
offset = int(pl_pose_stack.block_coord_offset[0, block])
block_coords = pl_pose_stack.coords[0, offset : offset + len(block_type.atoms)]
distances = torch.cdist(ligand_coords, block_coords)
nearest = float(distances.min())
if best is None or nearest < best[0]:
flat_index = int(distances.argmin())
ligand_atom, protein_atom = divmod(flat_index, distances.shape[1])
best = (nearest, block, ligand_atom, protein_atom)
distance, closest_block, ligand_atom_index, protein_atom_index = best
closest_type = pl_pose_stack.block_type(0, closest_block)
print(
f"closest ligand-carboxylate contact: {ligand_type.atoms[ligand_atom_index].name} to "
f"{closest_type.name3} block {closest_block} {closest_type.atoms[protein_atom_index].name}, {distance:.2f} A"
)
pl_atoms = torch.zeros((1, 2, 3), dtype=torch.int32, device=device)
pl_atoms[0, 0] = torch.tensor([0, ligand_block, ligand_atom_index])
pl_atoms[0, 1] = torch.tensor([0, closest_block, protein_atom_index])
pl_params = torch.tensor([[distance, 0.5]], dtype=torch.float32, device=device)
pl_constraints = ConstraintSet.create_empty(device=device, n_poses=1).add_constraints(
ConstraintEnergyTerm.harmonic, pl_atoms, pl_params
)
pl_constrained_pose = attrs.evolve(pl_pose_stack, constraint_set=pl_constraints)
print("protein-ligand interface constraint attached")
protein_offset = int(pl_pose_stack.block_coord_offset[0, closest_block])
pl_atom_mask = torch.zeros_like(pl_pose_stack.real_atoms)
pl_atom_mask[0, ligand_offset + ligand_atom_index] = True
pl_atom_mask[0, protein_offset + protein_atom_index] = True
viewer = tmol.view(pl_constrained_pose, highlighted=pl_atom_mask[pl_pose_stack.real_atoms])
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.
closest ligand-carboxylate contact: HN21 to ASP block 170 OD2, 1.97 A
protein-ligand interface constraint attached
/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.