TMol How-To Guide — Scoring and Optimization#
2. Scoring and Optimization#
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 section 1.
Setup (TODO - these extra setup steps will go away when the notebooks are combined)#
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
We use ubiquitin (1UBQ) as the running example, loaded the same way as in section 1, plus a ScoreFunction.
[2]:
cif_text = fetch_cif("1UBQ")
pose_stack = pose_stack_from_file(io.StringIO(cif_text), device)
score_function = beta2016_score_function(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=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.
Scoring#
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.
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.
Scoring a PoseStack#
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].
[3]:
# Create our score function.
score_function = beta2016_score_function(device, param_db=param_db)
# Create our scoring module.
whole_scorer = score_function.render_whole_pose_scoring_module(pose_stack)
# The scoring module is what does the actual score evaluation.
# This is separate from the ScoreFunction because it also needs details
# about the PoseStack being scored - mainly the Residue Types being used.
# Feed in the coordinates to the scorer and print the result
total_score = whole_scorer(pose_stack.coords)
print(f"total score: {float(total_score[0]):.2f}")
total score: 201.51
Score a PoseStack, get per-term breakdowns#
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].
score_function.all_score_types() gives the matching term order.
[4]:
weighted_terms = whole_scorer(pose_stack.coords, sum_terms=False, apply_weights=True)
# Get the score types so that we can label them with the names
score_types = score_function.all_score_types()
# Combine the scores and names and then print
term_breakdown = {st.name: float(weighted_terms[i, 0]) for i, st in enumerate(score_types)}
for name, value in sorted(term_breakdown.items(), key=lambda kv: -abs(kv[1]))[:8]:
print(f"{name:>16}: {value:9.2f}")
fa_ljatr: -415.45
fa_lk: 300.99
dunbrack_rotdev: 165.81
lk_ball_iso: -159.97
lk_ball: 158.11
fa_elec: -146.66
fa_ljrep: 126.12
cart_angles: 91.90
Setting score function weights#
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.
[5]:
focused = ScoreFunction(param_db, device)
focused.set_weight(ScoreType.fa_ljrep, 0.55)
focused.set_weight(ScoreType.hbond, 1.0)
focused_scorer = focused.render_whole_pose_scoring_module(pose_stack)
print(f"fa_ljrep + hbond only: {float(focused_scorer(pose_stack.coords)[0]):.2f}")
fa_ljrep + hbond only: 66.10
Block-pair scoring#
The ‘whole_pose’ scoring module above only returns a single energy for each pose.
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).
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.
[6]:
block_pair_scorer = score_function.render_block_pair_scoring_module(pose_stack)
block_pair_scores = block_pair_scorer(pose_stack.coords)
# Print the energy between residues 4 and 5 on pose 0
print(f"Score: {block_pair_scores[0, 4, 5]}")
Score: -0.44379860162734985
Interface energies#
We briefly switch to a two-chain example, 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.
[7]:
interface_cif_text = fetch_cif("1BL8")
interface_pose_stack = pose_stack_from_file(io.StringIO(interface_cif_text), device)
chain_a = get_chain_mask(interface_pose_stack, "A")
chain_b = get_chain_mask(interface_pose_stack, "B")
interface_score = calculate_block_pair_ddg(
interface_pose_stack, chain_a, chain_b,
sfxn=score_function, minimize=False, pack=False,
)
print(f"chain A: {int(chain_a.sum())} blocks, chain B: {int(chain_b.sum())} blocks")
print(f"A-B interface score: {float(interface_score[0]):.2f}")
chain_a_atom_mask = res_mask_to_coord_mask(interface_pose_stack, chain_a)
viewer = tmol.view(interface_pose_stack, highlighted=chain_a_atom_mask[interface_pose_stack.real_atoms])
viewer.show()
chain A: 100 blocks, chain B: 97 blocks
A-B interface score: 83.06
/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.
Per-residue energies#
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.
[8]:
per_residue_scores = per_residue_block_pair_energies(block_pair_scores)
print(f"residue score for residue 10 {float(per_residue_scores[0, 10]):7.2f}")
residue score for residue 10 3.28
Minimization#
tmol provides 2 different methods of minimizing energies - cartesian and torsion-space minimzation.
Cartesian-space minimization#
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.
[9]:
# Minimize the pose_stack
cart_minimized = tmol.run_cart_min(pose_stack.clone(), score_function)
# Score it and print the result.
cart_minimized_score = score_function.render_whole_pose_scoring_module(cart_minimized)(cart_minimized.coords)
print(f"before: {float(total_score[0]):.2f} after: {float(cart_minimized_score[0].detach()):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "cartesian minimized": cart_minimized},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"cartesian minimized": f"score: {float(cart_minimized_score[0].detach()):.2f}",
},
))
before: 201.51 after: -235.66
/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"))
/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"))
/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"))
/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"))
Torsion-space minimization#
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.
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.
[10]:
fold_forest = tmol.FoldForest.reasonable_fold_forest(pose_stack)
move_map = tmol.MoveMap.from_pose_stack(pose_stack)
move_map.move_all_named_torsions = True
move_map.move_all_jumps = False
move_map.move_all_root_jumps = False
kin_minimized = tmol.run_kin_min(
pose_stack.clone(), score_function, fold_forest, move_map, optimizer_kwargs={"max_iter": 10}
)
kin_minimized_score = score_function.render_whole_pose_scoring_module(kin_minimized)(kin_minimized.coords)
print(f"before: {float(total_score[0]):.2f} after (torsion space): {float(kin_minimized_score[0].detach()):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "torsion minimized": kin_minimized},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"torsion minimized": f"score (torsion space): {float(kin_minimized_score[0].detach()):.2f}",
},
))
before: 201.51 after (torsion space): 173.53
/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"))
/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"))
/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"))
/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"))
Gradients#
After scoring a structure, you can retreive the gradients directly, if you wish to use them outside the minimizer.
[11]:
# Cartesian gradients
print("Cartesian")
coords = pose_stack.coords.detach().clone().requires_grad_(True)
whole_scorer(coords).sum().backward()
atom_grad = torch.nan_to_num(coords.grad[0]).norm(dim=-1)
for atom in atom_grad.topk(10).indices.tolist():
print(f"{atom_label(pose_stack, 0, atom)}: {float(atom_grad[atom]):.3f} per A")
# Torsion gradients: forward kinematics (dofs -> coords) feeding the scorer
print("Torsion")
kin_module = PoseStackKinematicsModule(
pose_stack, tmol.FoldForest.reasonable_fold_forest(pose_stack)
)
kin_network = KinForestSfxnNetwork(score_function, pose_stack, kin_module)
kin_network().sum().backward()
dof_grad = torch.zeros_like(kin_network.full_dofs)
dof_grad[kin_network.dof_mask] = kin_network.masked_dofs.grad
labels, nodes = torsion_labels(pose_stack, kin_module, 0)
torsion_grad = dof_grad[nodes.to(dof_grad.device), BondDOFTypes.phi_c]
for i in torsion_grad.abs().topk(10).indices.tolist():
print(f"{labels[i]}: {float(torsion_grad[i]):.3f} per radian")
Cartesian
A 72 ARG O: 103.499 per A
A 72 ARG HG2: 97.927 per A
A 13 ILE N: 72.700 per A
A 13 ILE CB: 61.983 per A
A 13 ILE CA: 51.459 per A
A 71 LEU HA: 39.167 per A
A 71 LEU HD12: 38.615 per A
A 3 ILE C: 37.888 per A
A 15 LEU CA: 37.106 per A
A 15 LEU N: 35.677 per A
Torsion
A 39 ASP omega: -388.844 per radian
A 37 PRO phi: 388.389 per radian
A 39 ASP psi: -356.287 per radian
A 40 GLN phi: -355.432 per radian
A 36 ILE psi: 353.364 per radian
A 63 LYS phi: 338.571 per radian
A 57 SER psi: -317.802 per radian
A 34 GLU psi: 298.237 per radian
A 63 LYS omega: 297.719 per radian
A 62 GLN psi: 290.887 per radian
Controlling minimization with movemaps#
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.
Cartesian space#
[12]:
second_half_mask = torch.ones((1, pose_stack.max_n_blocks), dtype=torch.bool, device=device)
second_half_mask[:, :38] = False # freeze the first half of the chain
movable_coords = res_mask_to_coord_mask(pose_stack, second_half_mask)
cartesian_move_map = tmol.CartesianMoveMap(coord_mask=movable_coords)
partial_cart_min = tmol.run_cart_min(
pose_stack.clone(), score_function,
coord_mask=cartesian_move_map.coord_mask,
optimizer_kwargs={"max_iter": 10},
)
moved = (partial_cart_min.coords - pose_stack.coords).norm(dim=-1)
print(f"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A")
print(f"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A")
display(tmol.overlay_view(
{"original": pose_stack, "partial cartesian min": partial_cart_min},
notes={
"original": "unminimized",
"partial cartesian min": (
f"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, "
f"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A"
),
},
))
max movement, frozen half: 0.0000 A
max movement, movable half: 1.1327 A
/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"))
/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"))
/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"))
/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"))
Torsion space#
[13]:
partial_move_map = tmol.MoveMap.from_pose_stack(pose_stack)
partial_move_map.set_move_all_mc_tors_for_blocks(second_half_mask)
partial_move_map.set_move_all_sc_tors_for_blocks(second_half_mask)
partial_kin_min = tmol.run_kin_min(
pose_stack.clone(), score_function, fold_forest, partial_move_map, optimizer_kwargs={"max_iter": 10}
)
moved = (partial_kin_min.coords.detach() - pose_stack.coords).norm(dim=-1)
print(f"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A")
print(f"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A")
display(tmol.overlay_view(
{"original": pose_stack, "partial torsion min": partial_kin_min},
notes={
"original": "unminimized",
"partial torsion min": (
f"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, "
f"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A"
),
},
))
max movement, frozen half: 0.0000 A
max movement, movable half: 1.2537 A
/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"))
/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"))
/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"))
/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"))
Packing#
Defining packer behavior with a PackerTask#
To run packing, we must create a PackerTask. This controls what rotamers are to be considered at which positions.
[14]:
default_task = PackerTask(pose_stack, PackerPalette())
default_task.restrict_to_repacking()
default_task.add_conformer_sampler(FixedAAChiSampler())
default_task.add_conformer_sampler(IncludeCurrentSampler())
print("packer task ready: every position repackable")
packer task ready: every position repackable
Run packing TODO - pick an example that actually shows a difference from packing#
pack_rotamers runs the actual search and returns a new PoseStack.
[15]:
torch.manual_seed(SEED)
repacked = pack_rotamers(pose_stack, score_function, default_task)
repacked_score = score_function.render_whole_pose_scoring_module(repacked)(repacked.coords)
print(f"before: {float(total_score[0]):.2f} after repack: {float(repacked_score[0]):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "repacked": repacked},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"repacked": f"score: {float(repacked_score[0]):.2f}",
},
))
before: 201.51 after repack: 201.24
/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"))
/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"))
/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"))
/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"))
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.
[16]:
opth_task = PackerTask(pose_stack, PackerPalette())
opth_task.restrict_to_repacking()
opth_task.add_conformer_sampler(IncludeCurrentSampler())
opth_task.add_conformer_sampler(OptHSampler())
opth_task.add_conformer_sampler(FixedAAChiSampler())
torch.manual_seed(SEED)
opth_packed = pack_rotamers(pose_stack, score_function, opth_task)
opth_score = score_function.render_whole_pose_scoring_module(opth_packed)(opth_packed.coords)
print(f"before: {float(total_score[0]):.2f} after optH-only pack: {float(opth_score[0]):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "optH-only packed": opth_packed},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"optH-only packed": f"score: {float(opth_score[0]):.2f}",
},
))
before: 201.51 after optH-only pack: 201.24
/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"))
/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"))
/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"))
/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"))
Controlling packing at the residue level#
disable_packing_by_block_mask restricts the search to specific positions; everything else keeps its input conformation untouched.
[17]:
single_residue_task = PackerTask(pose_stack, PackerPalette())
single_residue_task.restrict_to_repacking()
shell = torch.zeros_like(single_residue_task.is_real_block)
shell[0, 10] = True
single_residue_task.disable_packing_by_block_mask(~shell)
single_residue_task.add_conformer_sampler(create_dunbrack_sampler_from_database(param_db, device))
single_residue_task.add_conformer_sampler(FixedAAChiSampler())
single_residue_task.add_conformer_sampler(IncludeCurrentSampler())
torch.manual_seed(SEED)
single_packed = pack_rotamers(pose_stack, score_function, single_residue_task)
moved = (single_packed.coords - pose_stack.coords).norm(dim=-1)[0]
untouched = torch.ones(pose_stack.max_n_blocks, dtype=torch.bool)
untouched[10] = False
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")
display(tmol.overlay_view(
{"original": pose_stack, "single-residue packed": single_packed},
notes={
"original": "block 10 unpacked",
"single-residue packed": (
"max movement outside block 10: "
f"{float(torch.nan_to_num(moved[res_mask_to_coord_mask(pose_stack, untouched.unsqueeze(0))[0]]).max()):.4f} A"
),
},
))
max movement outside block 10: 0.0000 A
/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"))
/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"))
/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"))
/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"))
Relax#
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.
Relax options (cycles, schedule)#
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.
[18]:
from tmol.relax import DEFAULT_RELAX_SCHEDULE
quick_schedule = [
{"fa_rep_pack_frac": 0.10, "fa_rep_min_frac": 0.20, "cst_frac": 1.0},
{"fa_rep_pack_frac": 1.00, "fa_rep_min_frac": 1.00, "cst_frac": 0.0},
]
print(f"a 2-stage schedule; the default schedule has {len(DEFAULT_RELAX_SCHEDULE)} stages")
a 2-stage schedule; the default schedule has 4 stages
Cartesian relax#
Passing a CartesianMoveMap relaxes in cartesian space.
[19]:
palette = PackerPalette()
full_cartesian_move_map = tmol.CartesianMoveMap()
torch.manual_seed(SEED)
cart_relaxed = tmol.fast_relax(
pose_stack.clone(), score_function, palette, full_cartesian_move_map, fold_forest,
num_repeats=1, schedule=quick_schedule,
)
cart_relaxed_score = score_function.render_whole_pose_scoring_module(cart_relaxed)(cart_relaxed.coords)
print(f"before: {float(total_score[0]):.2f} after cartesian relax: {float(cart_relaxed_score[0].detach()):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "cartesian relaxed": cart_relaxed},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"cartesian relaxed": f"score: {float(cart_relaxed_score[0].detach()):.2f}",
},
))
before: 201.51 after cartesian relax: -266.98
/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"))
/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"))
/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"))
/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"))
Torsion relax#
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.
[20]:
def kin_min_fn(ps, sfxn, *, fold_forest, move_map, verbose):
return tmol.run_kin_min(ps, sfxn, fold_forest, move_map, verbose=verbose)
torch.manual_seed(SEED)
kin_relaxed = tmol.fast_relax(
pose_stack.clone(), score_function, palette, move_map, fold_forest,
num_repeats=1, schedule=quick_schedule, min_fn=kin_min_fn,
)
kin_relaxed_score = score_function.render_whole_pose_scoring_module(kin_relaxed)(kin_relaxed.coords)
print(f"before: {float(total_score[0]):.2f} after torsion relax: {float(kin_relaxed_score[0].detach()):.2f}")
display(tmol.overlay_view(
{"original": pose_stack, "torsion relaxed": kin_relaxed},
notes={
"original": f"score: {float(total_score[0]):.2f}",
"torsion relaxed": f"score: {float(kin_relaxed_score[0].detach()):.2f}",
},
))
before: 201.51 after torsion relax: -167.26
/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"))
/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"))
/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"))
/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"))
Relax with a movemap#
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.
[21]:
restricted_move_map = tmol.CartesianMoveMap(coord_mask=movable_coords)
def freeze_first_half(task):
task.disable_packing_by_block_mask(~second_half_mask)
torch.manual_seed(SEED)
restricted_relaxed = tmol.fast_relax(
pose_stack.clone(), score_function, palette, restricted_move_map, fold_forest,
num_repeats=1, schedule=quick_schedule, task_operations=[freeze_first_half],
)
moved = (restricted_relaxed.coords - pose_stack.coords).norm(dim=-1)
print(f"max movement, frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A")
print(f"max movement, movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A")
display(tmol.overlay_view(
{"original": pose_stack, "movemap-restricted relax": restricted_relaxed},
notes={
"original": "unrelaxed",
"movemap-restricted relax": (
f"max movement frozen half: {float(torch.nan_to_num(moved[~movable_coords]).max()):.4f} A, "
f"movable half: {float(torch.nan_to_num(moved[movable_coords]).max()):.4f} A"
),
},
))
max movement, frozen half: 0.0000 A
max movement, movable half: 2.7246 A
/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"))
/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"))
/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"))
/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"))
Relax with constraints#
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.
[22]:
score_function.set_weight(ScoreType.constraint, 1.0)
constrained_start = create_mainchain_coordinate_constraints(pose_stack.clone())
torch.manual_seed(SEED)
constrained_relaxed = tmol.fast_relax(
constrained_start, score_function, palette, full_cartesian_move_map, fold_forest,
num_repeats=1, schedule=quick_schedule, ramp_constraints=True,
)
constrained_score = score_function.render_whole_pose_scoring_module(constrained_relaxed)(constrained_relaxed.coords)
print(f"before: {float(total_score[0]):.2f} after constrained relax: {float(constrained_score[0].detach()):.2f}")
display(tmol.overlay_view(
{"constrained start": constrained_start, "constrained relaxed": constrained_relaxed},
notes={
"constrained start": f"score: {float(total_score[0]):.2f}",
"constrained relaxed": f"score: {float(constrained_score[0].detach()):.2f}",
},
))
/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(
before: 201.51 after constrained relax: -268.72
/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"))
/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"))
/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"))
/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"))