Tutorial 05 — Minimization, constraints, and kinematics#

Minimize a small structure with Cartesian or kinematic movement, choose movable degrees of freedom, and add coordinate or distance constraints. Assumes scoring and gradients.

The short CPU runs check execution, not convergence. Larger systems benefit from CUDA.

Setup#

In Colab, select T4 GPU, then Run all. For local execution, follow the installation guide. Setup installs TMol and downloads the fixtures on first use.

The examples use a checked-in eight-residue 1UBQ slice on the selected CPU or CUDA device. A fixed constraint weight and ten minimizer iterations keep the Cartesian and kinematic runs suitable for documentation smoke tests; later sections explain why finite output is not evidence of convergence.

[1]:

try:
    import google.colab  # noqa: F401
except ImportError:
    IN_COLAB = False
else:
    IN_COLAB = True

if IN_COLAB:
    from urllib.request import urlopen

    exec(
        urlopen(
            "https://raw.githubusercontent.com/uw-ipd/tmol/"
            "master/docs/tutorial/colab_setup.py"
        ).read(),
        globals(),
    )
    setup_colab(["tmol/tests/data/cif/1UBQ.cif"])

[2]:

from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from pathlib import Path
import warnings

import attrs
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from IPython.display import display
from biotite.structure.io import load_structure

import tmol
from tmol.io import pose_stack_from_biotite
from tmol.kinematics import CartesianMoveMap, EdgeType, FoldForest, MoveMap
from tmol.ops import res_mask_to_coord_mask
from tmol.optimization import run_cart_min, run_kin_min
from tmol.pose import ConstraintSet, PoseStackBuilder
from tmol.score import ScoreType, beta2016_score_function
from tmol.score.constraint import (
    ConstraintEnergyTerm,
    create_mainchain_coordinate_constraints,
)

SEED = 20260807
np.random.seed(SEED)
torch.manual_seed(SEED)
warnings.filterwarnings(
    "ignore", message=r"Sparse invariant checks are implicitly disabled.*"
)
device = (
    torch.device("cuda", torch.cuda.current_device())
    if torch.cuda.is_available()
    else torch.device("cpu")
)


def show_table(frame):
    """Use a sortable docs table when available, otherwise display pandas."""
    try:
        from itables import show
    except ImportError:
        return display(frame)
    return show(frame)


def score_components(pose_stack, score_function):
    scorer = score_function.render_whole_pose_scoring_module(pose_stack)
    weighted_terms = scorer(
        pose_stack.coords, sum_terms=False, apply_weights=True
    )[:, 0]
    score_types = score_function.all_score_types()
    constraint_index = next(
        index
        for index, score_type in enumerate(score_types)
        if score_type == ScoreType.constraint
    )
    constraint_score = weighted_terms[constraint_index]
    non_constraint_score = weighted_terms.sum() - constraint_score
    return {
        "non_constraint_score_units": float(non_constraint_score.detach().cpu()),
        "constraint_score_units": float(constraint_score.detach().cpu()),
        "reported_total_score_units": float(weighted_terms.sum().detach().cpu()),
    }


def common_atom_mask(reference, mobile, atom_mask=None):
    finite = torch.isfinite(reference.coords).all(dim=-1)
    finite &= torch.isfinite(mobile.coords).all(dim=-1)
    return finite if atom_mask is None else finite & atom_mask


def rms_displacement(reference, mobile, atom_mask=None):
    """Raw RMS displacement in the input frame; no superposition is applied."""
    mask = common_atom_mask(reference, mobile, atom_mask)
    delta = mobile.coords[mask] - reference.coords[mask]
    return float(torch.sqrt(torch.mean(torch.sum(delta * delta, dim=-1))).cpu())


def kabsch_aligned_rmsd(reference, mobile, atom_mask=None):
    """RMSD after least-squares rigid-body superposition of mobile onto reference."""
    mask = common_atom_mask(reference, mobile, atom_mask)
    reference_xyz = reference.coords[mask]
    mobile_xyz = mobile.coords[mask]
    reference_centered = reference_xyz - reference_xyz.mean(dim=0)
    mobile_centered = mobile_xyz - mobile_xyz.mean(dim=0)

    covariance = mobile_centered.T @ reference_centered
    u, _, vh = torch.linalg.svd(covariance)
    handedness = torch.sign(torch.det(u @ vh))
    correction = torch.eye(3, device=mobile_xyz.device, dtype=mobile_xyz.dtype)
    correction[-1, -1] = handedness
    rotation = u @ correction @ vh
    aligned_mobile = mobile_centered @ rotation
    delta = aligned_mobile - reference_centered
    return float(torch.sqrt(torch.mean(torch.sum(delta * delta, dim=-1))).cpu())


def declared_mainchain_atom_mask(pose_stack):
    """Select atoms named in each residue type's TMol polymer declaration."""
    mask = torch.zeros_like(pose_stack.real_atoms)
    atom_names = []
    seen_atom_names = set()
    for pose_index in range(pose_stack.n_poses):
        for block_index in range(pose_stack.max_n_blocks):
            block_type_index = int(
                pose_stack.block_type_ind64[pose_index, block_index].item()
            )
            if block_type_index < 0:
                continue
            block_type = pose_stack.packed_block_types.active_block_types[
                block_type_index
            ]
            polymer = block_type.properties.polymer
            if polymer is None:
                continue
            atom_offset = int(
                pose_stack.block_coord_offset64[pose_index, block_index].item()
            )
            for atom_name in polymer.mainchain_atoms:
                atom_index = atom_offset + block_type.atom_to_idx[atom_name]
                if bool(pose_stack.real_atoms[pose_index, atom_index]):
                    mask[pose_index, atom_index] = True
                    if atom_name not in seen_atom_names:
                        atom_names.append(atom_name)
                        seen_atom_names.add(atom_name)
    return mask, tuple(atom_names)

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

Build a small pose and compare FoldForests#

FoldForest.reasonable_fold_forest() derives polymer and jump edges from backbone connectivity. Each edge is (type, start block, end block, jump index), and the edge array is batched over poses.

The same contiguous chain is also written explicitly with FoldForest.from_edges(). Finally, a per-residue-root or “dandelion” forest attaches every real block to the virtual root. That representation matches neural-network models that predict one frame per residue while retaining internal side-chain geometry; it changes the movable DOFs and is not a generic replacement for a polymer forest.

[3]:

repo_root = Path.cwd()
if not (repo_root / "tmol/tests/data/cif/1UBQ.cif").exists():
    repo_root = Path(tmol.__file__).resolve().parents[1]
cif_path = repo_root / "tmol" / "tests" / "data" / "cif" / "1UBQ.cif"
atom_array = load_structure(str(cif_path), model=1, include_bonds=True)
protein_slice = atom_array[(atom_array.chain_id == "A") & (atom_array.res_id <= 8)]
pose_diagnostics = StringIO()
try:
    with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):
        start = pose_stack_from_biotite(protein_slice, device, no_optH=True)
except Exception:
    print(pose_diagnostics.getvalue())
    raise

score_function = beta2016_score_function(device)
score_function.set_weight(ScoreType.constraint, 1.0)
fold_forest = FoldForest.reasonable_fold_forest(start)

explicit_edges = np.full((1, 2, 4), -1, dtype=np.int64)
explicit_edges[0, 0] = [int(EdgeType.root_jump), -1, 0, -1]
explicit_edges[0, 1] = [
    int(EdgeType.polymer),
    0,
    start.max_n_blocks - 1,
    -1,
]
explicit_fold_forest = FoldForest.from_edges(explicit_edges)

# One root edge per real block: useful for per-residue NN frames.
dandelion_edges = np.full(
    (start.n_poses, start.max_n_blocks, 4), -1, dtype=np.int64
)
dandelion_edges[:, :, 0] = int(EdgeType.root_jump)
dandelion_edges[:, :, 1] = -1
dandelion_edges[:, :, 2] = np.arange(start.max_n_blocks)[None, :]
per_residue_root_forest = FoldForest.from_edges(dandelion_edges)
assert int(per_residue_root_forest.n_edges[0]) == start.max_n_blocks

edge_names = {int(edge): edge.name for edge in EdgeType}
edge_rows = []
for forest_name, forest in (
    ("automatic polymer", fold_forest),
    ("explicit polymer", explicit_fold_forest),
    ("per-residue root / dandelion", per_residue_root_forest),
):
    for pose_i in range(start.n_poses):
        for edge_i in range(int(forest.n_edges[pose_i])):
            edge_type, begin, end, jump_index = forest.edges[pose_i, edge_i]
            edge_rows.append(
                {
                    "forest": forest_name,
                    "pose": pose_i,
                    "edge": edge_i,
                    "type": edge_names[int(edge_type)],
                    "start": int(begin),
                    "end": int(end),
                    "jump_index": int(jump_index),
                }
            )

print("input:", cif_path.name)
print("automatic FoldForest edge tensor shape:", fold_forest.edges.shape)
show_table(pd.DataFrame(edge_rows))

input: 1UBQ.cif
automatic FoldForest edge tensor shape: (1, 2, 4)

^ⓘforest	pose	edge	type	start	end	jump_index
automatic polymer	0	0	root_jump	-1	0	-1
automatic polymer	0	1	polymer	0	7	-1
explicit polymer	0	0	root_jump	-1	0	-1
explicit polymer	0	1	polymer	0	7	-1
per-residue root / dandelion	0	0	root_jump	-1	0	-1
per-residue root / dandelion	0	1	root_jump	-1	1	-1
per-residue root / dandelion	0	2	root_jump	-1	2	-1
per-residue root / dandelion	0	3	root_jump	-1	3	-1
per-residue root / dandelion	0	4	root_jump	-1	4	-1
per-residue root / dandelion	0	5	root_jump	-1	5	-1
(2 more rows not shown)

The automatic and explicit polymer forests each contain a root_jump whose start is -1, followed by a polymer edge covering the contiguous 1UBQ slice. The dandelion forest instead has one root_jump per real block and no polymer edge. A jump_index of -1 is expected for root-jump and polymer edges; only ordinary jumps receive contiguous non-negative jump indices. Missing, duplicated, or overlapping coverage indicates a malformed forest and should be investigated before minimization.

Every FoldForest has an explicit virtual root at the origin. reasonable_fold_forest() follows backbone up/down connections, turns same-chain gaps into jumps, and roots separate chains independently. Its current automatic builder intentionally ignores non-polymer connections such as disulfides.

Minimize with and without restraints#

run_cart_min accepts (pose_stack, sfxn, coord_mask=None, optimizer_cls=..., optimizer_kwargs=None, verbose=False). run_kin_min accepts (pose_stack, sfxn, ff, mm, optimizer_cls=..., optimizer_kwargs=None, verbose=False, kin_dtype=...).

Distance restraint#

ConstraintSet identifies atoms by (pose, block, atom-within-block). Restrain the first-to-last Cα distance to its input value with σ = 0.5 Å, then replicate the constraint across two poses with add_constraints_to_all_poses().

Main-chain coordinate restraints#

create_mainchain_coordinate_constraints() follows each type’s mainchain_atoms: N, CA, C for this fixture, excluding O. It uses σ = 0.5 Å and retains existing constraints. The score weight here is 1.0.

Movable coordinates#

Cartesian runs move all finite atom coordinates. An unrestrained isolated molecule has six rigid-body null modes. Raw displacement includes frame drift; Kabsch RMSD removes global translation and rotation.

Kinematic runs fix ordinary and root jumps and enable named backbone/side-chain torsions. Default ideal kinematics does not independently optimize lengths and angles. Cartesian runs can change them, penalized by beta2016’s cart_lengths, cart_angles, cart_torsions, and cart_impropers (weight 0.5 each), plus cart_hxltorsions (weight 1.0). The two runs therefore optimize different degrees of freedom.

[4]:

cartesian_move_map = CartesianMoveMap(
    coord_mask=torch.isfinite(start.coords).all(dim=-1)
)
kinematic_move_map = MoveMap.from_pose_stack(start)
kinematic_move_map.move_all_named_torsions = True  # named MC and SC torsions
kinematic_move_map.move_all_jumps = False
kinematic_move_map.move_all_root_jumps = False

mainchain_mask, restrained_atom_names = declared_mainchain_atom_mask(start)
constraint_diagnostics = StringIO()
try:
    with redirect_stdout(constraint_diagnostics), redirect_stderr(constraint_diagnostics):
        constrained_start = create_mainchain_coordinate_constraints(start.clone())
except Exception:
    print(constraint_diagnostics.getvalue())
    raise

# One low-level Cα distance restraint, expressed in block-local atom indices.
first_block, last_block = 0, start.max_n_blocks - 1
first_ca = start.block_type(0, first_block).atom_to_idx["CA"]
last_ca = start.block_type(0, last_block).atom_to_idx["CA"]
first_coord_index = int(start.block_coord_offset64[0, first_block]) + first_ca
last_coord_index = int(start.block_coord_offset64[0, last_block]) + last_ca
target_distance = torch.linalg.vector_norm(
    start.coords[0, first_coord_index] - start.coords[0, last_coord_index]
)
distance_atoms = torch.tensor(
    [[[0, first_block, first_ca], [0, last_block, last_ca]]],
    dtype=torch.int32,
    device=device,
)
distance_params = torch.zeros((1, 4), dtype=torch.float32, device=device)
distance_params[0, 0] = target_distance
distance_params[0, 1] = 0.5
distance_constraint_set = ConstraintSet.create_empty(device, 1).add_constraints(
    ConstraintEnergyTerm.harmonic,
    distance_atoms,
    distance_params,
)
distance_constrained_start = attrs.evolve(
    start, constraint_set=distance_constraint_set
)

constraint_batch = PoseStackBuilder.from_poses([start] * 2, device)
batch_constraint_set = ConstraintSet.create_empty(
    device, constraint_batch.n_poses
).add_constraints_to_all_poses(
    ConstraintEnergyTerm.harmonic,
    distance_atoms,
    distance_params,
)
distance_constrained_batch = attrs.evolve(
    constraint_batch, constraint_set=batch_constraint_set
)
assert len(distance_constraint_set.constraint_function_inds) == 1
assert len(batch_constraint_set.constraint_function_inds) == 2

CONSTRAINT_SIGMA_A = 0.5
constraint_weight = float(score_function.get_weight(ScoreType.constraint))
print(
    f"distance restraint: block {first_block} CA ↔ block {last_block} CA; "
    f"target={float(target_distance):.3f} Å; sigma=0.5 Å"
)
print(
    "distance constraints in one pose / replicated batch:",
    len(distance_constraint_set.constraint_function_inds),
    "/",
    len(batch_constraint_set.constraint_function_inds),
)
print(
    "declared-main-chain coordinate restraints:",
    len(constrained_start.constraint_set.constraint_function_inds),
)
print("coordinate-restraint atom names:", ", ".join(restrained_atom_names))
print("coordinate-restraint count:", int(mainchain_mask.sum().item()))
print(f"coordinate-restraint sigma: {CONSTRAINT_SIGMA_A:.1f} Å")
print(f"coordinate-restraint score weight: {constraint_weight:.1f}")
print("O is not TMol-declared main chain here and is not constrained.")
print("Weight 1.0 is intentional; constraint weights are protocol-specific.")

# Strictly a short smoke-test budget for tutorial/CI runtime. Finite output and
# a lower objective after ten iterations do not demonstrate convergence.
TUTORIAL_MAX_ITER = 10
optimizer_kwargs = {"max_iter": TUTORIAL_MAX_ITER}
print(f"smoke-test optimizer max_iter: {TUTORIAL_MAX_ITER} (convergence not assessed)")

# Every branch starts from the same coordinates, score function, and iteration
# budget. Within each coordinate model, only the presence of restraints changes.
cart_unconstrained = run_cart_min(
    start.clone(),
    score_function,
    coord_mask=cartesian_move_map.coord_mask,
    optimizer_kwargs=optimizer_kwargs,
)
cart_constrained = run_cart_min(
    constrained_start.clone(),
    score_function,
    coord_mask=cartesian_move_map.coord_mask,
    optimizer_kwargs=optimizer_kwargs,
)
kin_unconstrained = run_kin_min(
    start.clone(),
    score_function,
    fold_forest,
    kinematic_move_map,
    optimizer_kwargs=optimizer_kwargs,
)
kin_constrained = run_kin_min(
    constrained_start.clone(),
    score_function,
    fold_forest,
    kinematic_move_map,
    optimizer_kwargs=optimizer_kwargs,
)

distance restraint: block 0 CA ↔ block 7 CA; target=23.155 Å; sigma=0.5 Å
distance constraints in one pose / replicated batch: 1 / 2
declared-main-chain coordinate restraints: 24
coordinate-restraint atom names: N, CA, C
coordinate-restraint count: 24
coordinate-restraint sigma: 0.5 Å
coordinate-restraint score weight: 1.0
O is not TMol-declared main chain here and is not constrained.
Weight 1.0 is intentional; constraint weights are protocol-specific.
smoke-test optimizer max_iter: 10 (convergence not assessed)

[5]:

variants = {
    "Cartesian, unconstrained": cart_unconstrained,
    "Cartesian, constrained": cart_constrained,
    "Kinematic, unconstrained": kin_unconstrained,
    "Kinematic, constrained": kin_constrained,
}
mainchain_mask, _ = declared_mainchain_atom_mask(start)
start_components = score_components(start, score_function)
rows = [
    {
        "variant": "start",
        **start_components,
        "non_constraint_score_change_units": 0.0,
        "constraint_score_change_units": 0.0,
        "objective_change_units": 0.0,
        "all_atom_raw_RMS_displacement_A": 0.0,
        "all_atom_Kabsch_RMSD_A": 0.0,
        "mainchain_raw_RMS_displacement_A": 0.0,
        "mainchain_Kabsch_RMSD_A": 0.0,
        "finite_smoke_test_output": True,
        "convergence_status": "not assessed",
    }
]
for name, pose_stack in variants.items():
    components = score_components(pose_stack, score_function)
    all_atom_raw = rms_displacement(start, pose_stack)
    all_atom_aligned = kabsch_aligned_rmsd(start, pose_stack)
    mainchain_raw = rms_displacement(start, pose_stack, mainchain_mask)
    mainchain_aligned = kabsch_aligned_rmsd(start, pose_stack, mainchain_mask)
    numeric_outputs = (
        *components.values(),
        all_atom_raw,
        all_atom_aligned,
        mainchain_raw,
        mainchain_aligned,
    )
    rows.append(
        {
            "variant": name,
            **components,
            "non_constraint_score_change_units": (
                components["non_constraint_score_units"]
                - start_components["non_constraint_score_units"]
            ),
            "constraint_score_change_units": (
                components["constraint_score_units"]
                - start_components["constraint_score_units"]
            ),
            "objective_change_units": (
                components["reported_total_score_units"]
                - start_components["reported_total_score_units"]
            ),
            "all_atom_raw_RMS_displacement_A": all_atom_raw,
            "all_atom_Kabsch_RMSD_A": all_atom_aligned,
            "mainchain_raw_RMS_displacement_A": mainchain_raw,
            "mainchain_Kabsch_RMSD_A": mainchain_aligned,
            "finite_smoke_test_output": bool(np.isfinite(numeric_outputs).all()),
            "convergence_status": "not assessed (10-iteration smoke test)",
        }
    )
comparison_frame = pd.DataFrame(rows)
show_table(comparison_frame)

fig, axes = plt.subplots(2, 3, figsize=(18, 8))
axes = axes.ravel()
axes[0].barh(
    comparison_frame["variant"], comparison_frame["non_constraint_score_units"]
)
axes[0].set(xlabel="weighted score units", title="Non-constraint score")
axes[1].barh(
    comparison_frame["variant"], comparison_frame["constraint_score_units"]
)
axes[1].set(xlabel="weighted score units", title="Constraint penalty")
axes[2].barh(
    comparison_frame["variant"],
    comparison_frame["all_atom_raw_RMS_displacement_A"],
)
axes[2].set(xlabel="raw RMS displacement (Å)", title="All atoms: input frame")
axes[3].barh(
    comparison_frame["variant"], comparison_frame["all_atom_Kabsch_RMSD_A"]
)
axes[3].set(xlabel="Kabsch RMSD (Å)", title="All atoms: aligned")
axes[4].barh(
    comparison_frame["variant"],
    comparison_frame["mainchain_raw_RMS_displacement_A"],
)
axes[4].set(xlabel="raw RMS displacement (Å)", title="Main chain: input frame")
axes[5].barh(
    comparison_frame["variant"], comparison_frame["mainchain_Kabsch_RMSD_A"]
)
axes[5].set(xlabel="Kabsch RMSD (Å)", title="Main chain: aligned")
plt.tight_layout()
plt.show()

try:
    restraint_viewer = tmol.view(
        start,
        highlighted=mainchain_mask[start.real_atoms],
        highlight_color="#7b2cbf",
    )
    restraint_viewer.show()
except ImportError:
    print("Install py3Dmol to inspect the restrained main-chain atoms.")

viewer_structures = {"before": start, **variants}
viewer_notes = {"before": "Starting 1UBQ slice"}
for row in rows[1:]:
    viewer_notes[row["variant"]] = (
        f"objective change {row['objective_change_units']:+.3f} score units; "
        f"non-constraint {row['non_constraint_score_units']:.3f}; "
        f"constraint {row['constraint_score_units']:.3f}; "
        f"main-chain raw/aligned "
        f"{row['mainchain_raw_RMS_displacement_A']:.3f}/"
        f"{row['mainchain_Kabsch_RMSD_A']:.3f} Å; convergence not assessed"
    )
display(tmol.switchable_view(viewer_structures, notes=viewer_notes))

/tmp/ipykernel_4012/382234577.py:78: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /__w/pytorch/pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:820.)
  return float(torch.sqrt(torch.mean(torch.sum(delta * delta, dim=-1))).cpu())

^ⓘvariant	non_constraint_score_units	constraint_score_units	reported_total_score_units	non_constraint_score_change_units	constraint_score_change_units	objective_change_units	all_atom_raw_RMS_displacement_A	all_atom_Kabsch_RMSD_A	mainchain_raw_RMS_displacement_A	mainchain_Kabsch_RMSD_A	finite_smoke_test_output	convergence_status
start	28.661953	0.000000	28.661953	0.000000	0.000000	0.000000	0.000000	0.000000	0.000000	0.000000	True	not assessed
Cartesian, unconstrained	9.868545	0.000000	9.868545	-18.793408	0.000000	-18.793408	0.094823	0.094822	0.067072	0.062319	True	not assessed (10-iteration smoke test)
Cartesian, constrained	9.873143	0.340653	10.213797	-18.788810	0.340653	-18.448156	0.092888	0.092869	0.059569	0.056067	True	not assessed (10-iteration smoke test)
Kinematic, unconstrained	13.958851	0.000000	13.958851	-14.703102	0.000000	-14.703102	4.730233	1.445364	4.063869	0.768433	True	not assessed (10-iteration smoke test)
Kinematic, constrained	20.235886	1.326277	21.562162	-8.426067	1.326277	-7.099791	0.354134	0.311590	0.117539	0.092123	True	not assessed (10-iteration smoke test)

../_images/tutorial_05_minimization_constraints_kinematics_9_2.png

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

Structure:

Starting 1UBQ slice

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

objective change -18.793 score units; non-constraint 9.869; constraint 0.000; main-chain raw/aligned 0.067/0.062 Å; convergence not assessed

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

objective change -18.448 score units; non-constraint 9.873; constraint 0.341; main-chain raw/aligned 0.060/0.056 Å; convergence not assessed

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

objective change -14.703 score units; non-constraint 13.959; constraint 0.000; main-chain raw/aligned 4.064/0.768 Å; convergence not assessed

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

objective change -7.100 score units; non-constraint 20.236; constraint 1.326; main-chain raw/aligned 0.118/0.092 Å; convergence not assessed

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

Check finite scores and structural metrics. Ten iterations do not establish convergence, even if the objective falls; assess stopping diagnostics and repeatability with a suitable budget.

Main-chain RMSD uses the restrained N/CA/C atoms shown in purple. All-atom RMSD includes side-chain minimization; this notebook does no packing. A large raw-versus-aligned displacement indicates rigid-body drift in unrestrained Cartesian runs. The kinematic root is fixed.

Compare restrained and unrestrained runs within each coordinate model. Cartesian and kinematic endpoints need not agree.

Restraint-guided rescue of a displaced local segment#

The preceding comparison starts from a deposited structure. This controlled rescue experiment creates an explicit defect: all atoms in zero-based blocks 3–4 are translated by the same vector while blocks 2–6 define the movable local shell. Three short Cartesian runs then start from exactly that perturbed structure:

no restraints;
reference main-chain coordinate restraints copied from the unperturbed pose; or
one reference Cα–Cα distance restraint linking fixed-side block 1 to displaced block 4.

The experiment asks how restraint information trades against the non-constraint score and recovery toward the known reference. It is a method diagnostic, not loop modeling or structure prediction.

[6]:

local_blocks = torch.zeros_like(start.block_type_ind, dtype=torch.bool)
local_blocks[:, 2:7] = True
displaced_blocks = torch.zeros_like(local_blocks)
displaced_blocks[:, 3:5] = True
local_coord_mask = res_mask_to_coord_mask(start, local_blocks) & start.real_atoms
displaced_coord_mask = (
    res_mask_to_coord_mask(start, displaced_blocks) & start.real_atoms
)
local_mainchain_mask = mainchain_mask & local_coord_mask

perturbed = start.clone()
displacement = torch.tensor(
    [0.8, -0.6, 0.4], device=device, dtype=perturbed.coords.dtype
)
perturbed.coords[displaced_coord_mask] += displacement

# The coordinate-constraint targets were created from the unperturbed start.
coordinate_rescue_start = attrs.evolve(
    perturbed.clone(), constraint_set=constrained_start.constraint_set
)

anchor_block, rescued_block = 1, 4
anchor_ca = start.block_type(0, anchor_block).atom_to_idx["CA"]
rescued_ca = start.block_type(0, rescued_block).atom_to_idx["CA"]
anchor_coord = int(start.block_coord_offset64[0, anchor_block]) + anchor_ca
rescued_coord = int(start.block_coord_offset64[0, rescued_block]) + rescued_ca
rescue_target_distance = torch.linalg.vector_norm(
    start.coords[0, anchor_coord] - start.coords[0, rescued_coord]
)
rescue_atoms = torch.tensor(
    [[[0, anchor_block, anchor_ca], [0, rescued_block, rescued_ca]]],
    dtype=torch.int32,
    device=device,
)
rescue_params = torch.zeros((1, 4), dtype=torch.float32, device=device)
rescue_params[0, 0] = rescue_target_distance
rescue_params[0, 1] = 0.5
rescue_distance_set = ConstraintSet.create_empty(device, 1).add_constraints(
    ConstraintEnergyTerm.harmonic,
    rescue_atoms,
    rescue_params,
)
distance_rescue_start = attrs.evolve(
    perturbed.clone(), constraint_set=rescue_distance_set
)

rescue_budget = {"max_iter": 15}
rescue_variants = {
    "perturbed input": perturbed,
    "unconstrained rescue": run_cart_min(
        perturbed.clone(),
        score_function,
        coord_mask=local_coord_mask,
        optimizer_kwargs=rescue_budget,
    ),
    "reference coordinate rescue": run_cart_min(
        coordinate_rescue_start,
        score_function,
        coord_mask=local_coord_mask,
        optimizer_kwargs=rescue_budget,
    ),
    "reference distance rescue": run_cart_min(
        distance_rescue_start,
        score_function,
        coord_mask=local_coord_mask,
        optimizer_kwargs=rescue_budget,
    ),
}

rescue_rows = []
for label, pose_stack in rescue_variants.items():
    components = score_components(pose_stack, score_function)
    observed_distance = torch.linalg.vector_norm(
        pose_stack.coords[0, anchor_coord] - pose_stack.coords[0, rescued_coord]
    )
    rescue_rows.append(
        {
            "state": label,
            **components,
            "local_mainchain_RMSD_from_reference_A": rms_displacement(
                start, pose_stack, local_mainchain_mask
            ),
            "reference_distance_error_A": abs(
                float(observed_distance.detach().cpu())
                - float(rescue_target_distance.detach().cpu())
            ),
            "movable_atoms": int(local_coord_mask.sum().item()),
            "optimizer_budget": (
                "input only"
                if label == "perturbed input"
                else "15-iteration smoke test; convergence not assessed"
            ),
        }
    )
rescue_frame = pd.DataFrame(rescue_rows)
show_table(rescue_frame)

fig, axis = plt.subplots(figsize=(7, 4.5))
axis.scatter(
    rescue_frame["local_mainchain_RMSD_from_reference_A"],
    rescue_frame["non_constraint_score_units"],
    s=70,
)
for row in rescue_frame.itertuples():
    axis.annotate(
        row.state,
        (
            row.local_mainchain_RMSD_from_reference_A,
            row.non_constraint_score_units,
        ),
        xytext=(5, 5),
        textcoords="offset points",
        fontsize=8,
    )
axis.set(
    xlabel="local main-chain RMSD from reference (Å)",
    ylabel="weighted non-constraint score",
    title="Score improvement and structural recovery are distinct",
)
axis.grid(alpha=0.3)
plt.tight_layout()
plt.show()

display(
    tmol.switchable_view(
        {"reference": start, **rescue_variants},
        notes={
            "reference": "Unperturbed eight-residue reference",
            **{
                row.state: (
                    f"local main-chain RMSD "
                    f"{row.local_mainchain_RMSD_from_reference_A:.3f} Å; "
                    f"distance error {row.reference_distance_error_A:.3f} Å; "
                    f"non-constraint score {row.non_constraint_score_units:.2f}"
                )
                for row in rescue_frame.itertuples()
            },
        },
    )
)

^ⓘstate	non_constraint_score_units	constraint_score_units	reported_total_score_units	local_mainchain_RMSD_from_reference_A	reference_distance_error_A	movable_atoms	optimizer_budget
perturbed input	110.912704	0.000000	110.912704	0.681175	0.241167	91	input only
unconstrained rescue	32.149914	0.000000	32.149914	0.560058	0.307968	91	15-iteration smoke test; convergence not assessed
reference coordinate rescue	35.354950	6.437084	41.792034	0.327543	0.192638	91	15-iteration smoke test; convergence not assessed
reference distance rescue	32.167385	0.305045	32.472431	0.555575	0.276155	91	15-iteration smoke test; convergence not assessed

../_images/tutorial_05_minimization_constraints_kinematics_12_1.png

Structure:

Unperturbed eight-residue reference

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

local main-chain RMSD 0.681 Å; distance error 0.241 Å; non-constraint score 110.91

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

local main-chain RMSD 0.560 Å; distance error 0.308 Å; non-constraint score 32.15

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

local main-chain RMSD 0.328 Å; distance error 0.193 Å; non-constraint score 35.35

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

local main-chain RMSD 0.556 Å; distance error 0.276 Å; non-constraint score 32.17

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

Score improvement need not match structural recovery. Coordinate restraints generally pull the main chain toward its reference; one distance restraint cannot determine a unique conformation. Assess recovery and convergence separately from the 15-iteration execution check.

Next#

Complete packing, then combine both operations in FastRelax.

Exercises#

Change the Cartesian coordinate mask so only residues 3–5 move, then compare the displacement table.
Disable side-chain torsions in the kinematic MoveMap and enable only main-chain torsions for selected blocks.
Replace the one Cα–Cα harmonic restraint with a tested ConstraintEnergyTerm.bounded distance constraint and verify its two-pose replication.
Build a two-chain slice and compare its automatic forest with an explicit FoldForest.from_edges() definition.
Apply create_mainchain_coordinate_constraints() to a mixed protein–nucleic-acid pose and inspect the atoms selected from each residue type’s declared main chain.

Download this notebook