Tutorial 04 — Packing and mutation scans#

Open In Colab

Repack side chains, inspect rotamers, and run a small mutation scan and regional design example. Assumes scoring.

Repacking preserves residue identities; design allows selected identities to change. The scan composes lower-level APIs. CPU and GPU packers can select different outcomes.

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 setup uses the sampler construction exercised by TMol’s packer tests and existing score/pack/minimize tutorial. The ten-residue slice keeps CPU execution and two-outcome comparisons practical. NumPy and PyTorch seeds are initialized, but the current CPU annealer uses the C rand() stream and is not controlled by torch.manual_seed(). Hydrogen coordinates are optimized during pose construction before scores are interpreted.

[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
from time import perf_counter
import warnings

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.database import ParameterDatabase
from tmol.io import pose_stack_from_biotite
from tmol.numeric import coord_dihedrals
from tmol.ops import (
    build_sidechain_coord_mask,
    compute_block_adjacency,
    compute_block_centroids_and_furthest_dist,
    res_mask_to_coord_mask,
)
from tmol.pack import PackerPalette, PackerTask, SetPackerTask, pack_rotamers
from tmol.pack.rotamer import FixedAAChiSampler, IncludeCurrentSampler, build_rotamers
from tmol.pack.rotamer.dunbrack import create_dunbrack_sampler_from_database
from tmol.pose import PoseStackBuilder
from tmol.score import beta2016_score_function

SEED = 20260807
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(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")
)
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"

param_db = ParameterDatabase.get_default()
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 <= 10)]
pose_diagnostics = StringIO()
try:
    with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):
        # Optimize hydrogen coordinates before comparing interpreted scores.
        pose_stack = pose_stack_from_biotite(
            protein_slice, device, param_db=param_db, no_optH=False
        )
except Exception:
    print(pose_diagnostics.getvalue())
    raise
score_function = beta2016_score_function(device, param_db=param_db)


def show_table(frame):
    try:
        from itables import show
    except ImportError:
        display(frame)
        return None
    return show(frame)


def total_score(pose):
    scorer = score_function.render_whole_pose_scoring_module(pose)
    return float(scorer(pose.coords).detach().cpu()[0])

print(f"device={device}; input={cif_path.name}; slice blocks={pose_stack.max_n_blocks}")

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

device=cpu; input=1UBQ.cif; slice blocks=10

Fixed-sequence regional repacking#

PackerPalette defines the residue-type space considered from each starting block. PackerTask.restrict_to_repacking() keeps each position’s name3 fixed. The task then uses Dunbrack rotamers, fixed-amino-acid chi samples, and the current conformation.

Here residues 3–8 (zero-based block indices 2–7) are allowed to repack. disable_packing_by_block_mask() receives True where packing should be disabled.

[3]:

palette = PackerPalette()
repack_task = PackerTask(pose_stack, palette)
repack_task.restrict_to_repacking()

packing_region = torch.zeros(
    (pose_stack.n_poses, pose_stack.max_n_blocks),
    dtype=torch.bool,
    device=device,
)
packing_region[:, 2:8] = True
repack_task.disable_packing_by_block_mask(~packing_region)
repack_task.add_conformer_sampler(
    create_dunbrack_sampler_from_database(param_db, device)
)
repack_task.add_conformer_sampler(FixedAAChiSampler())
repack_task.add_conformer_sampler(IncludeCurrentSampler())

start_score = total_score(pose_stack)
torch.manual_seed(SEED)
repacked_pose = pack_rotamers(pose_stack, score_function, repack_task)
repacked_score = total_score(repacked_pose)

packing_atom_mask = res_mask_to_coord_mask(pose_stack, packing_region)
packing_sidechain_mask = packing_atom_mask & build_sidechain_coord_mask(pose_stack)
sidechain_displacement = torch.linalg.vector_norm(
    repacked_pose.coords - pose_stack.coords, dim=-1
)[packing_sidechain_mask]
sidechain_rms_displacement = float(
    torch.sqrt(torch.mean(sidechain_displacement.square())).detach().cpu()
)
max_sidechain_displacement = float(sidechain_displacement.max().detach().cpu())
changed_sidechain_atoms = int((sidechain_displacement > 1e-3).sum().detach().cpu())

repack_frame = pd.DataFrame(
    [
        {
            "stage": "input",
            "beta2016": start_score,
            "score_change": 0.0,
            "packed_sidechain_RMS_displacement_A": 0.0,
            "max_sidechain_displacement_A": 0.0,
            "changed_sidechain_atoms": 0,
        },
        {
            "stage": "regional repack",
            "beta2016": repacked_score,
            "score_change": repacked_score - start_score,
            "packed_sidechain_RMS_displacement_A": sidechain_rms_displacement,
            "max_sidechain_displacement_A": max_sidechain_displacement,
            "changed_sidechain_atoms": changed_sidechain_atoms,
        },
    ]
)
show_table(repack_frame)
ⓘstage beta2016 score_change packed_sidechain_RMS_displacement_A max_sidechain_displacement_A changed_sidechain_atoms
input28.7787110.0000000.0000000.0000000
regional repack20.820780-7.9579320.4918511.48222874

Only blocks 3–8 may change side-chain conformers. Sequence and backbone coordinates stay fixed. The table reports scores, RMS/maximum side-chain displacement, and atoms moving more than 0.001 Å.

CPU and CUDA use different searches:

  • CPU: one trajectory, 20 outer cycles with 5 × N_rotamers random trials per cycle, then a quench. It uses C rand(), which torch.manual_seed() does not seed.

  • CUDA: 500 high-temperature trajectories, 1,250 low-temperature trajectories, then 312 quenches. It uses PyTorch’s CUDA Philox generator.

Treat these as different stochastic protocols; equal seeds do not imply matching outcomes. Interactive viewing requires py3Dmol.

[4]:

try:
    display(
        tmol.switchable_view(
            {"input": pose_stack, "regional repack": repacked_pose},
            notes={
                "input": f"weighted score: {start_score:.3f}",
                "regional repack": (
                    f"score {repacked_score:.3f}; Δ {repacked_score - start_score:+.3f}; "
                    f"side-chain RMS motion {sidechain_rms_displacement:.3f} Å"
                ),
            },
            width=720,
            height=420,
        )
    )
except ImportError as exc:
    print("Interactive repacking comparison unavailable:", exc)
weighted score: 28.779

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

Expand side-chain sampling#

PackerTask.or_expand_chi() adds samples at approximately ±1 standard deviation around supported wells. Indices are zero-based: 0 is χ1 and 1 is χ2.

Rotamer counts are deterministic for a fixed task; packing outcomes are stochastic. Extra candidates need not improve the selected score. CUDA shows two seeded outcomes, CPU one. IncludeCurrentSampler retains the input conformation. These short runs do not establish convergence.

[5]:

expanded_task = PackerTask(pose_stack, palette)
expanded_task.restrict_to_repacking()
expanded_task.disable_packing_by_block_mask(~packing_region)
expanded_task.add_conformer_sampler(
    create_dunbrack_sampler_from_database(param_db, device)
)
expanded_task.add_conformer_sampler(FixedAAChiSampler())
expanded_task.add_conformer_sampler(IncludeCurrentSampler())
# Chi indices are zero-based: 0 is chi1 and 1 is chi2.
expanded_task.or_expand_chi(0)
expanded_task.or_expand_chi(1)

torch.manual_seed(SEED)
expanded_repacked_pose = pack_rotamers(
    pose_stack, score_function, expanded_task
)
expanded_score = total_score(expanded_repacked_pose)

# Keep this deliberately small. CUDA gets a second seeded outcome because
# torch.manual_seed() controls its Philox stream; the CPU C RNG has no exposed
# seed API, so repeating under a different torch seed would not be a controlled
# replicate.
sampling_outcome_rows = [
    {
        "sampling": "standard",
        "outcome": 1,
        "weighted_score_units": repacked_score,
    },
    {
        "sampling": "expanded χ1 + χ2",
        "outcome": 1,
        "weighted_score_units": expanded_score,
    },
]
if device.type == "cuda":
    torch.manual_seed(SEED + 1)
    standard_repeat_pose = pack_rotamers(pose_stack, score_function, repack_task)
    torch.manual_seed(SEED + 1001)
    expanded_repeat_pose = pack_rotamers(pose_stack, score_function, expanded_task)
    sampling_outcome_rows.extend(
        [
            {
                "sampling": "standard",
                "outcome": 2,
                "weighted_score_units": total_score(standard_repeat_pose),
            },
            {
                "sampling": "expanded χ1 + χ2",
                "outcome": 2,
                "weighted_score_units": total_score(expanded_repeat_pose),
            },
        ]
    )
else:
    print("CPU packing shown once per setting: torch seeds do not control C rand().")
sampling_outcome_frame = pd.DataFrame(sampling_outcome_rows)
show_table(sampling_outcome_frame)
sampling_outcome_summary = (
    sampling_outcome_frame.groupby("sampling", as_index=False)
    .agg(
        score_mean=("weighted_score_units", "mean"),
        score_min=("weighted_score_units", "min"),
        score_max=("weighted_score_units", "max"),
    )
)
sampling_outcome_summary["observed_range"] = (
    sampling_outcome_summary["score_max"]
    - sampling_outcome_summary["score_min"]
)
show_table(sampling_outcome_summary)
try:
    display(
        tmol.switchable_view(
            {
                "standard rotamers": repacked_pose,
                "expanded χ1 + χ2": expanded_repacked_pose,
            },
            notes={
                "standard rotamers": f"outcome 1 weighted score {repacked_score:.3f}",
                "expanded χ1 + χ2": f"outcome 1 weighted score {expanded_score:.3f}",
            },
        )
    )
except ImportError as exc:
    print("Interactive sampling comparison unavailable:", exc)

CPU packing shown once per setting: torch seeds do not control C rand().
ⓘsampling outcome weighted_score_units
standard120.820780
expanded χ1 + χ2120.678461
ⓘsampling score_mean score_min score_max observed_range
expanded χ1 + χ220.67846120.67846120.6784610.0
standard20.82078020.82078020.8207800.0
outcome 1 weighted score 20.821

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

Inspect deterministic candidate expansion#

Expanded chi sampling changes the discrete search space before annealing. The table below builds both rotamer sets explicitly and reports deterministic candidate counts at each packed position. A wrapped χ1-versus-χ2 plot then shows candidate geometry for one representative block without treating internal rotamer storage order as a physical coordinate.

These torsions describe the candidate set, not a trajectory or score result: simulated annealing evaluates assignments stochastically and returns only the selected final conformation.

[6]:

standard_rotamer_pose, standard_rotamers = build_rotamers(
    pose_stack,
    SetPackerTask.from_packer_task(repack_task),
    pose_stack.packed_block_types.chem_db,
)
expanded_rotamer_pose, expanded_rotamers = build_rotamers(
    pose_stack,
    SetPackerTask.from_packer_task(expanded_task),
    pose_stack.packed_block_types.chem_db,
)

rotamer_count_rows = []
for block_index in range(2, 8):
    block_type = pose_stack.block_type(0, block_index)
    rotamer_count_rows.append(
        {
            "block_index": block_index,
            "residue": block_type.name3,
            "standard_candidates": int(
                standard_rotamers.n_rots_for_block[0, block_index]
            ),
            "expanded_candidates": int(
                expanded_rotamers.n_rots_for_block[0, block_index]
            ),
        }
    )
rotamer_count_frame = pd.DataFrame(rotamer_count_rows)
rotamer_count_frame["additional_candidates"] = (
    rotamer_count_frame["expanded_candidates"]
    - rotamer_count_frame["standard_candidates"]
)
show_table(rotamer_count_frame)


def named_chi_frame(label, rotamer_pose, rotamer_set, block_index):
    rows = []
    n_rotamers = int(rotamer_set.n_rots_for_block[0, block_index])
    rotamer_offset = int(rotamer_set.rot_offset_for_block[0, block_index])
    for local_rotamer in range(n_rotamers):
        rotamer_index = rotamer_offset + local_rotamer
        block_type_index = int(rotamer_set.block_type_ind_for_rot[rotamer_index])
        block_type = rotamer_pose.packed_block_types.active_block_types[
            block_type_index
        ]
        coordinate_offset = int(rotamer_set.coord_offset_for_rot[rotamer_index])
        row = {"sampling": label, "rotamer": local_rotamer}
        for torsion_index in block_type.sc_torsions:
            atom_indices = block_type.ordered_torsions[torsion_index, :, 0]
            if np.any(atom_indices < 0):
                continue
            coordinates = rotamer_set.coords[
                coordinate_offset
                + torch.as_tensor(atom_indices, dtype=torch.long, device=device)
            ].double()
            radians = coord_dihedrals(
                coordinates[0:1],
                coordinates[1:2],
                coordinates[2:3],
                coordinates[3:4],
            )[0]
            torsion_name = block_type.torsions[torsion_index].name
            row[f"{torsion_name}_degrees"] = float(torch.rad2deg(radians).cpu())
        rows.append(row)
    return pd.DataFrame(rows)


# Block 5 is LYS in this 1UBQ slice and therefore demonstrates both chi1 and chi2.
rotamer_inspection_block = 5
rotamer_geometry_frame = pd.concat(
    [
        named_chi_frame(
            "standard",
            standard_rotamer_pose,
            standard_rotamers,
            rotamer_inspection_block,
        ),
        named_chi_frame(
            "expanded",
            expanded_rotamer_pose,
            expanded_rotamers,
            rotamer_inspection_block,
        ),
    ],
    ignore_index=True,
)
show_table(rotamer_geometry_frame)

chi1_column = next(
    (
        column
        for column in rotamer_geometry_frame.columns
        if column.lower().startswith("chi1") and column.endswith("_degrees")
    ),
    None,
)
chi2_column = next(
    (
        column
        for column in rotamer_geometry_frame.columns
        if column.lower().startswith("chi2") and column.endswith("_degrees")
    ),
    None,
)
if chi1_column is not None and chi2_column is not None:
    wrapped_geometry = rotamer_geometry_frame.copy()
    wrapped_geometry["chi1_wrapped_degrees"] = (
        (wrapped_geometry[chi1_column] + 180.0) % 360.0 - 180.0
    )
    wrapped_geometry["chi2_wrapped_degrees"] = (
        (wrapped_geometry[chi2_column] + 180.0) % 360.0 - 180.0
    )
    fig, ax = plt.subplots(figsize=(5.5, 5.0))
    for label, values in wrapped_geometry.groupby("sampling"):
        ax.scatter(
            values["chi1_wrapped_degrees"],
            values["chi2_wrapped_degrees"],
            label=label,
            alpha=0.7,
        )
    ax.set(
        xlim=(-180, 180),
        ylim=(-180, 180),
        xlabel="wrapped χ1 (degrees)",
        ylabel="wrapped χ2 (degrees)",
        title=f"Candidate side-chain geometry at block {rotamer_inspection_block}",
    )
    ax.set_aspect("equal", adjustable="box")
    ax.grid(alpha=0.3)
    ax.legend()
    plt.tight_layout()
    plt.show()
ⓘblock_index residue standard_candidates expanded_candidates additional_candidates
2ILE75548
3PHE1310996
4VAL4106
5LYS41361320
6THR3710972
7LEU75548
ⓘsampling rotamer chi1_degrees chi2_degrees chi3_degrees chi4_degrees
standard0-176.578568177.852829179.383850-180.000000
standard1177.91906768.981468176.515289178.500015
standard2-178.95297271.885971178.46687369.606934
standard3-58.962368-64.175339-175.639450-177.932907
standard4-170.350662-82.596512178.522079-177.400085
standard5-177.339264176.993179-176.099976-66.127388
standard6-65.590263-179.024490-179.595016179.459991
standard7-178.528931177.198334175.33515964.900024
standard8-58.109982-64.209007-173.958755-66.999977
standard9-173.606400178.399933-68.899933-68.500175
(392 more rows not shown)
../_images/tutorial_04_packing_and_mutation_scan_11_2.png

Local mutation scan#

The helper restricts per_block_is_block_type_allowed by turning entries from True to False; PackerTask has no public single-identity restriction method. Non-target positions retain their original type, and the target retains only types matching the requested name3.

Repack the target and its compute_block_adjacency() shell: centroid distance below the sum of both furthest-atom radii plus the default 5 Å margin.

Batch identities together, with two outcomes each on CUDA and one on CPU. Compare each identity’s mean with the independently repacked WT mean; do not pair individual outcomes.

[7]:

def restrict_batched_substitutions(task, block_index, name3_by_pose):
    """Keep the original sequence except for one requested identity per pose."""
    if len(name3_by_pose) != task.per_block_is_block_type_allowed.shape[0]:
        raise ValueError("Provide exactly one requested identity per pose")
    considered = task.per_block_considered_block_types
    keep = task.per_block_considered_block_types_is_orig.detach().clone()
    for pose_index, name3 in enumerate(name3_by_pose):
        target_considered = (
            considered[pose_index, block_index].detach().cpu().tolist()
        )
        target_keep = torch.tensor(
            [
                index >= 0 and task.pbt.active_block_types[index].name3 == name3
                for index in target_considered
            ],
            dtype=torch.bool,
            device=task.device,
        )
        if not torch.any(target_keep):
            raise ValueError(f"{name3!r} is not considered at block {block_index}")
        keep[pose_index, block_index] = target_keep
    task.per_block_is_block_type_allowed = torch.logical_and(
        task.per_block_is_block_type_allowed, keep
    )
    return task


def add_standard_samplers(task):
    task.add_conformer_sampler(
        create_dunbrack_sampler_from_database(param_db, device)
    )
    task.add_conformer_sampler(FixedAAChiSampler())
    task.add_conformer_sampler(IncludeCurrentSampler())
    return task

[8]:

target_block = 4
original_type_index = int(pose_stack.block_type_ind64[0, target_block].item())
original_name3 = pose_stack.packed_block_types.active_block_types[
    original_type_index
].name3
variants = list(dict.fromkeys([original_name3, "ALA", "LEU", "PHE"]))
target_label = (
    f"{pose_stack.pdb_info.chain_labels[0, target_block]}:"
    f"{pose_stack.pdb_info.residue_labels[0, target_block]}"
    f"{pose_stack.pdb_info.residue_insertion_codes[0, target_block]}"
)

# Build one explicit local shell with the existing score_utils geometry helpers.
block_centroids, block_radii = compute_block_centroids_and_furthest_dist(pose_stack)
block_adjacency = compute_block_adjacency(block_centroids, block_radii)
mutation_shell = block_adjacency[:, target_block].clone()
mutation_shell[:, target_block] = True
shell_block_indices = torch.nonzero(mutation_shell[0], as_tuple=False).flatten().tolist()
shell_labels = [
    f"{pose_stack.pdb_info.chain_labels[0, block]}:"
    f"{pose_stack.pdb_info.residue_labels[0, block]}"
    f"{pose_stack.pdb_info.residue_insertion_codes[0, block]}"
    for block in shell_block_indices
]
print(
    f"mutation target block {target_block} ({target_label}); local repacking shell "
    f"has {len(shell_block_indices)} blocks: {', '.join(shell_labels)}"
)


def target_to_rest_score(pose, block_index):
    scorer = score_function.render_block_pair_scoring_module(pose)
    matrix = scorer(pose.coords, sum_terms=True, apply_weights=True)[0]
    partners = torch.ones(matrix.shape[0], dtype=torch.bool, device=matrix.device)
    partners[block_index] = False
    return float(
        (matrix[block_index, partners].sum() + matrix[partners, block_index].sum())
        .detach()
        .cpu()
    )


# CUDA uses two seeded outcomes per identity. On CPU, show one batch member
# per identity because torch.manual_seed() does not control the C RNG.
n_scan_outcomes = 2 if device.type == "cuda" else 1
scan_requests = [
    (outcome, variant)
    for outcome in range(1, n_scan_outcomes + 1)
    for variant in variants
]
requested_variants = [variant for _, variant in scan_requests]
mutation_batch = PoseStackBuilder.from_poses(
    [pose_stack] * len(scan_requests), device
)
mutation_scan_task = PackerTask(mutation_batch, PackerPalette())
restrict_batched_substitutions(
    mutation_scan_task, target_block, requested_variants
)
mutation_shell_batch = mutation_shell.expand(len(scan_requests), -1).clone()
mutation_scan_task.disable_packing_by_block_mask(~mutation_shell_batch)
add_standard_samplers(mutation_scan_task)

torch.manual_seed(SEED)  # Controls CUDA annealing, not the CPU C rand() stream.
if device.type == "cuda":
    torch.cuda.synchronize()
scan_start = perf_counter()
packed_mutation_batch = pack_rotamers(
    mutation_batch, score_function, mutation_scan_task
)
if device.type == "cuda":
    torch.cuda.synchronize()
scan_seconds = perf_counter() - scan_start

scan_rows = []
variant_poses = {}
for pose_index, (outcome, variant) in enumerate(scan_requests):
    packed_variant = packed_mutation_batch.split(pose_index)
    state_label = f"{variant} (repacked WT)" if variant == original_name3 else variant
    if outcome == 1:
        variant_poses[state_label] = packed_variant

    # The monotonic identity helper must leave every non-target identity unchanged.
    for block_index in range(packed_variant.max_n_blocks):
        if block_index != target_block:
            assert (
                packed_variant.block_type(0, block_index).name3
                == pose_stack.block_type(0, block_index).name3
            )

    scan_rows.append(
        {
            "block_index": target_block,
            "pdb_label": target_label,
            "variant": variant,
            "state": state_label,
            "outcome": outcome,
            "total_weighted_score_units": total_score(packed_variant),
            "target_to_rest_weighted_score_units": target_to_rest_score(
                packed_variant, target_block
            ),
            "shared_batch_seconds": scan_seconds,
        }
    )

scan_frame = pd.DataFrame(scan_rows)
show_table(scan_frame.sort_values(["variant", "outcome"]))

# Summarize each identity independently, then compare distribution means.
scan_summary = (
    scan_frame.groupby(["variant", "state"], as_index=False)
    .agg(
        total_score_mean=("total_weighted_score_units", "mean"),
        total_score_min=("total_weighted_score_units", "min"),
        total_score_max=("total_weighted_score_units", "max"),
        target_to_rest_mean=("target_to_rest_weighted_score_units", "mean"),
        target_to_rest_min=("target_to_rest_weighted_score_units", "min"),
        target_to_rest_max=("target_to_rest_weighted_score_units", "max"),
    )
)
scan_summary["total_score_range"] = (
    scan_summary["total_score_max"] - scan_summary["total_score_min"]
)
scan_summary["target_to_rest_range"] = (
    scan_summary["target_to_rest_max"] - scan_summary["target_to_rest_min"]
)
wt_summary = scan_summary.loc[scan_summary["variant"] == original_name3].iloc[0]
scan_summary["total_difference_of_means_vs_repacked_wt"] = (
    scan_summary["total_score_mean"] - wt_summary["total_score_mean"]
)
scan_summary["target_to_rest_difference_of_means_vs_repacked_wt"] = (
    scan_summary["target_to_rest_mean"] - wt_summary["target_to_rest_mean"]
)
scan_summary = scan_summary.sort_values(
    "total_difference_of_means_vs_repacked_wt"
)
show_table(scan_summary)

mutation target block 4 (A:5); local repacking shell has 7 blocks: A:2, A:3, A:4, A:5, A:6, A:7, A:8
ⓘ block_index pdb_label variant state outcome total_weighted_score_units target_to_rest_weighted_score_units shared_batch_seconds
14A:5ALAALA122.9564721.3562410.236953
24A:5LEULEU119.289846-1.8347080.236953
34A:5PHEPHE120.958065-1.2219410.236953
04A:5VALVAL (repacked WT)118.264778-2.8326450.236953
ⓘ variant state total_score_mean total_score_min total_score_max target_to_rest_mean target_to_rest_min target_to_rest_max total_score_range target_to_rest_range total_difference_of_means_vs_repacked_wt target_to_rest_difference_of_means_vs_repacked_wt
3VALVAL (repacked WT)18.26477818.26477818.264778-2.832645-2.832645-2.8326450.00.00.0000000.000000
1LEULEU19.28984619.28984619.289846-1.834708-1.834708-1.8347080.00.01.0250680.997937
2PHEPHE20.95806520.95806520.958065-1.221941-1.221941-1.2219410.00.02.6932871.610704
0ALAALA22.95647222.95647222.9564721.3562411.3562411.3562410.00.04.6916944.188887
[9]:

plot_summary = scan_summary.sort_values(
    "target_to_rest_difference_of_means_vs_repacked_wt"
).reset_index(drop=True)
fig, ax = plt.subplots(figsize=(8, 4.2))
for y_position, row in plot_summary.iterrows():
    outcome_offsets = (
        scan_frame.loc[
            scan_frame["variant"] == row["variant"],
            "target_to_rest_weighted_score_units",
        ].to_numpy()
        - wt_summary["target_to_rest_mean"]
    )
    ax.hlines(
        y_position,
        outcome_offsets.min(),
        outcome_offsets.max(),
        color="0.45",
        linewidth=2,
    )
    ax.scatter(outcome_offsets, np.full_like(outcome_offsets, y_position), color="0.35")
    ax.scatter(
        row["target_to_rest_difference_of_means_vs_repacked_wt"],
        y_position,
        color="tab:blue",
        marker="D",
        zorder=3,
    )
ax.axvline(0.0, color="black", linewidth=1)
ax.set_yticks(np.arange(len(plot_summary)), plot_summary["state"])
ax.set(
    xlabel=(
        "target-to-rest weighted score relative to independently summarized "
        "repacked-WT mean"
    ),
    title=f"Two-outcome local-shell scan at {target_label}: outcomes, range, and mean",
)
ax.grid(axis="x", alpha=0.25)
plt.tight_layout()
plt.show()

shell_atom_mask = res_mask_to_coord_mask(pose_stack, mutation_shell)
try:
    shell_viewer = tmol.view(
        pose_stack,
        highlighted=shell_atom_mask[pose_stack.real_atoms],
        highlight_color="#ff7f0e",
    )
    print(
        f"Orange: local repacking shell ({len(shell_block_indices)} blocks); "
        f"mutation target: block {target_block} ({target_label})."
    )
    shell_viewer.show()
except ImportError as exc:
    print("Interactive local-shell viewer unavailable:", exc)
../_images/tutorial_04_packing_and_mutation_scan_15_0.png 
Orange: local repacking shell (7 blocks); mutation target: block 4 (A:5).

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

The orange selection marks the repacking shell; the author label identifies the target. shared_batch_seconds measures the full packing call, not one variant.

The plot shows outcomes, ranges, and means relative to repacked WT. Target-to-rest scores combine both matrix orientations before comparing means. Zero observed spread does not establish determinism or convergence. One CPU or two CUDA outcomes are illustrative only.

[10]:

lowest_mean_state = scan_summary.iloc[0]
print(
    "lowest mean total score in this two-outcome local scan:",
    lowest_mean_state["state"],
)
summary_by_state = scan_summary.set_index("state")
try:
    display(
        tmol.switchable_view(
            variant_poses,
            notes={
                state: (
                    f"outcome 1 structure; independently summarized mean score "
                    f"{summary_by_state.loc[state, 'total_score_mean']:.3f}; "
                    f"observed range "
                    f"{summary_by_state.loc[state, 'total_score_range']:.3f}"
                )
                for state in variant_poses
            },
            width=720,
            height=420,
        )
    )
except ImportError as exc:
    print("Interactive mutation comparison unavailable:", exc)

lowest mean total score in this two-outcome local scan: VAL (repacked WT)
outcome 1 structure; independently summarized mean score 18.265; observed range 0.000

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

Regional sequence design#

Design positions may change identity; adjacent shell blocks may repack; blocks outside the shell stay frozen. Checks use each block’s coordinate offset to verify the frozen coordinates within float32 reconstruction precision.

The three batch members each receive a complete anneal. CUDA uses independent Philox-driven searches; CPU uses its C-RNG search and may return identical results. Batch members are not trajectories within one annealer, nor controlled replicates across devices.

[11]:

def restrict_regional_design(task, allowed_name3_by_block):
    """Keep the original identity elsewhere and allow a small design alphabet."""
    considered = task.per_block_considered_block_types
    keep = task.per_block_considered_block_types_is_orig.detach().clone()
    for pose_index in range(keep.shape[0]):
        for block_index, allowed_names in allowed_name3_by_block.items():
            considered_types = (
                considered[pose_index, block_index].detach().cpu().tolist()
            )
            allowed = torch.tensor(
                [
                    index >= 0
                    and task.pbt.active_block_types[index].name3 in allowed_names
                    for index in considered_types
                ],
                dtype=torch.bool,
                device=task.device,
            )
            if not torch.any(allowed):
                raise ValueError(f"No requested identities at block {block_index}")
            keep[pose_index, block_index] = allowed
    task.per_block_is_block_type_allowed = torch.logical_and(
        task.per_block_is_block_type_allowed, keep
    )
    return task


design_blocks = [4, 6]
design_alphabet = {
    block_index: {
        pose_stack.block_type(0, block_index).name3,
        "ALA",
        "LEU",
        "PHE",
    }
    for block_index in design_blocks
}

# The union of adjacency neighborhoods defines which blocks may pack.
design_block_mask = torch.zeros(
    (pose_stack.n_poses, pose_stack.max_n_blocks),
    dtype=torch.bool,
    device=device,
)
design_block_mask[:, design_blocks] = True
design_shell = torch.any(block_adjacency[:, design_blocks, :], dim=1)
design_shell |= design_block_mask
design_repack_only = design_shell & ~design_block_mask
design_frozen = ~design_shell

def design_pdb_label(block_index):
    return (
        f"{pose_stack.pdb_info.chain_labels[0, block_index]}:"
        f"{pose_stack.pdb_info.residue_labels[0, block_index]}"
        f"{pose_stack.pdb_info.residue_insertion_codes[0, block_index]}"
    )


design_role_rows = []
for block_index in range(pose_stack.max_n_blocks):
    if bool(design_block_mask[0, block_index]):
        role = "design"
    elif bool(design_repack_only[0, block_index]):
        role = "repack-only"
    else:
        role = "frozen"
    design_role_rows.append(
        {
            "block_index": block_index,
            "pdb_label": design_pdb_label(block_index),
            "starting_identity": pose_stack.block_type(0, block_index).name3,
            "role": role,
        }
    )
design_role_frame = pd.DataFrame(design_role_rows)
show_table(design_role_frame)

n_design_batch_members = 3
design_batch = PoseStackBuilder.from_poses(
    [pose_stack] * n_design_batch_members, device
)
design_task = PackerTask(design_batch, PackerPalette())
restrict_regional_design(design_task, design_alphabet)
design_shell_batch = design_shell.expand(n_design_batch_members, -1).clone()
design_task.disable_packing_by_block_mask(~design_shell_batch)
add_standard_samplers(design_task)

torch.manual_seed(SEED)
if device.type == "cuda":
    torch.cuda.synchronize()
design_start = perf_counter()
designed_batch = pack_rotamers(design_batch, score_function, design_task)
if device.type == "cuda":
    torch.cuda.synchronize()
design_seconds = perf_counter() - design_start

design_rows = []
design_poses = {}
frozen_block_indices = torch.nonzero(
    design_frozen[0], as_tuple=False
).flatten().tolist()
for pose_index in range(n_design_batch_members):
    designed_pose = designed_batch.split(pose_index)
    label = f"batch member {pose_index + 1}"
    identities = [
        designed_pose.block_type(0, block_index).name3
        for block_index in design_blocks
    ]

    # Compare each frozen block in local coordinates; mutations can shift later offsets.
    max_frozen_coordinate_change = 0.0
    for block_index in frozen_block_indices:
        input_block_type = pose_stack.block_type(0, block_index)
        output_block_type = designed_pose.block_type(0, block_index)
        assert output_block_type.name == input_block_type.name
        n_block_atoms = len(input_block_type.atoms)
        input_offset = int(pose_stack.block_coord_offset64[0, block_index])
        output_offset = int(designed_pose.block_coord_offset64[0, block_index])
        input_block_coords = pose_stack.coords[
            0, input_offset : input_offset + n_block_atoms
        ]
        output_block_coords = designed_pose.coords[
            0, output_offset : output_offset + n_block_atoms
        ]
        # Pose reconstruction may introduce a few micro-angstroms of float32
        # roundoff even when packing is disabled for the block.
        torch.testing.assert_close(
            output_block_coords, input_block_coords, rtol=1e-6, atol=1e-5
        )
        block_change = float(
            torch.linalg.vector_norm(
                output_block_coords - input_block_coords, dim=-1
            )
            .max()
            .detach()
            .cpu()
        )
        max_frozen_coordinate_change = max(
            max_frozen_coordinate_change, block_change
        )

    design_poses[label] = designed_pose
    design_rows.append(
        {
            "batch_member": pose_index + 1,
            "block_4_identity": identities[0],
            "block_6_identity": identities[1],
            "sequence_pair": "-".join(identities),
            "total_score": total_score(designed_pose),
            "max_frozen_coordinate_change_A": max_frozen_coordinate_change,
            "shared_batch_seconds": design_seconds,
        }
    )

design_frame = pd.DataFrame(design_rows).sort_values("total_score")
show_table(design_frame)

design_shell_atom_mask = res_mask_to_coord_mask(pose_stack, design_shell)
try:
    design_region_viewer = tmol.view(
        pose_stack,
        highlighted=design_shell_atom_mask[pose_stack.real_atoms],
        highlight_color="#ff7f0e",
    )
    print(
        "Orange: complete design packing shell; see the role table for "
        "design, repack-only, and frozen assignments."
    )
    design_region_viewer.show()
    display(
        tmol.switchable_view(
            design_poses,
            notes={
                f"batch member {row['batch_member']}": (
                    f"{row['sequence_pair']}; max frozen coordinate change "
                    f"{row['max_frozen_coordinate_change_A']:.3f} Å"
                )
                for row in design_rows
            },
            width=720,
            height=420,
        )
    )
except ImportError as exc:
    print("Interactive design comparison unavailable:", exc)
ⓘblock_index pdb_label starting_identity role
0A:1METfrozen
1A:2GLNrepack-only
2A:3ILErepack-only
3A:4PHErepack-only
4A:5VALdesign
5A:6LYSrepack-only
6A:7THRdesign
7A:8LEUrepack-only
8A:9THRrepack-only
9A:10GLYrepack-only
ⓘbatch_member block_4_identity block_6_identity sequence_pair total_score max_frozen_coordinate_change_A shared_batch_seconds
1VALTHRVAL-THR15.6725120.00.313489
2VALTHRVAL-THR15.6725120.00.313489
3VALTHRVAL-THR15.6725120.00.313489
Orange: complete design packing shell; see the role table for design, repack-only, and frozen assignments.

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

VAL-THR; max frozen coordinate change 0.000 Å

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

Next#

Complete minimization, then FastRelax.

Exercises#

  1. Change the regional repacking mask and verify that disabled blocks retain their input coordinates.

  2. Change the adjacency constant, redisplay the mutation and design shells, and record which author-labeled blocks change role.

  3. Increase the scan to at least five independent outcomes per identity; summarize each distribution before computing differences of means, and do not call any score difference a \(\Delta\Delta G\). Remember that torch.manual_seed() still does not seed the CPU C RNG.

  4. Add Cartesian minimization after each packed variant and state how that changes the computational experiment.

  5. Sketch an external scheduler that shards positions across GPUs while preserving one result row per position, identity, outcome, and score-function version.

References#

Download this notebook