Tutorial 10 — Ligand pose sensitivity#

Open In Colab

Displace a crystallographic ligand, score the decoys in one batch, and locally minimize three selected poses. Compare interaction scores with ligand heavy-atom displacement.

Complete ligand preparation first; this notebook reuses its ADA/LG1 fixture and parameters.

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 checked-in CIF and matching .tmol file are the authoritative coordinate and chemical inputs. The notebook runs on CPU or CUDA, initializes NumPy and PyTorch seeds, and uses the fixtures downloaded during Colab setup instead of querying a live structure database.

[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/protein_ligand_test/ada.tmol.nomin.cif",
            "tmol/tests/data/protein_ligand_test/ada.xtal-lig.mmff94.tmol",
        ]
    )

[2]:

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

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

import tmol
from tmol.database import ParameterDatabase
from tmol.io import build_context_from_biotite, pose_stack_from_biotite
from tmol.ligand import inject_params_file
from tmol.ops import build_sidechain_coord_mask, res_mask_to_coord_mask
from tmol.optimization import run_cart_min
from tmol.pose import PoseStackBuilder
from tmol.score import beta2016_score_function

SEED = 20260810
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")
)
LIGAND_NAME = "LG1"


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


def block_mask_for_name3(pose, name3):
    """Select blocks with one residue name across a PoseStack."""
    mask = torch.zeros_like(pose.block_type_ind, dtype=torch.bool)
    for pose_index in range(pose.n_poses):
        for block_index in range(pose.max_n_blocks):
            if int(pose.block_type_ind64[pose_index, block_index]) < 0:
                continue
            mask[pose_index, block_index] = (
                pose.block_type(pose_index, block_index).name3 == name3
            )
    return mask


def heavy_atom_mask_for_blocks(pose, block_mask):
    """Expand selected blocks to their non-hydrogen coordinate atoms."""
    mask = torch.zeros_like(pose.real_atoms)
    for pose_index, block_index in torch.nonzero(block_mask, as_tuple=False).tolist():
        block_type_index = int(pose.block_type_ind64[pose_index, block_index].item())
        block_type = pose.block_type(pose_index, block_index)
        offset = int(pose.block_coord_offset64[pose_index, block_index])
        n_atoms = len(block_type.atoms)
        is_hydrogen = pose.packed_block_types.atom_is_hydrogen[
            block_type_index, :n_atoms
        ].bool()
        mask[pose_index, offset : offset + n_atoms] = ~is_hydrogen
    return mask & pose.real_atoms


def ligand_protein_interactions(pose, score_function):
    """Return both-orientation weighted ligand–protein scores per pose."""
    ligand = block_mask_for_name3(pose, LIGAND_NAME)
    protein = (pose.block_type_ind64 >= 0) & ~ligand
    pair_mask = (ligand[:, :, None] & protein[:, None, :]) | (
        protein[:, :, None] & ligand[:, None, :]
    )
    scorer = score_function.render_block_pair_scoring_module(pose)
    with torch.no_grad():
        weighted = scorer(
            pose.coords, sum_terms=False, apply_weights=True
        ).sum(dim=0)
    return (weighted * pair_mask).sum(dim=(1, 2))


def ligand_rmsd_from_native(pose, native_heavy_coords):
    """Measure ligand heavy-atom RMSD in the fixed protein coordinate frame."""
    ligand = block_mask_for_name3(pose, LIGAND_NAME)
    heavy = heavy_atom_mask_for_blocks(pose, ligand)
    values = []
    for pose_index in range(pose.n_poses):
        delta = pose.coords[pose_index, heavy[pose_index]] - native_heavy_coords
        values.append(torch.sqrt(torch.mean(torch.sum(delta * delta, dim=-1))))
    return torch.stack(values)

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 one ligand-aware context#

The pose and score function use the same extended parameter database. Hydrogen optimization is disabled because the experiment isolates controlled rigid-body ligand placement and short local Cartesian refinement; every compared state starts from the same prepared coordinates and chemistry.

[3]:

repo_root = Path.cwd()
data_dir = repo_root / "tmol" / "tests" / "data" / "protein_ligand_test"
if not (data_dir / "ada.tmol.nomin.cif").exists():
    repo_root = Path(tmol.__file__).resolve().parents[1]
    data_dir = repo_root / "tmol" / "tests" / "data" / "protein_ligand_test"

complex_path = data_dir / "ada.tmol.nomin.cif"
params_path = data_dir / "ada.xtal-lig.mmff94.tmol"
parameter_database = inject_params_file(ParameterDatabase.get_default(), params_path)
atom_array = biotite.structure.io.load_structure(
    str(complex_path), model=1, include_bonds=True
)

diagnostics = StringIO()
try:
    with redirect_stdout(diagnostics), redirect_stderr(diagnostics):
        build_context = build_context_from_biotite(
            atom_array,
            device,
            param_db=parameter_database,
            prepare_ligands=False,
        )
        native_pose = pose_stack_from_biotite(
            atom_array,
            device,
            context=build_context,
            no_optH=True,
        )
except Exception:
    print(diagnostics.getvalue())
    raise

score_function = beta2016_score_function(
    device, param_db=build_context.parameter_database
)
native_ligand = block_mask_for_name3(native_pose, LIGAND_NAME)
if int(native_ligand.sum().item()) != 1:
    raise RuntimeError("Expected exactly one LG1 ligand block")
native_ligand_coords = res_mask_to_coord_mask(native_pose, native_ligand)
native_heavy_mask = heavy_atom_mask_for_blocks(native_pose, native_ligand)
native_heavy_coords = native_pose.coords[native_heavy_mask]

print(f"device: {device}")
print(f"coordinate input: {complex_path.name}")
print(f"authoritative chemistry: {params_path.name}")
print("ligand heavy atoms:", int(native_heavy_mask.sum().item()))

device: cpu
coordinate input: ada.tmol.nomin.cif
authoritative chemistry: ada.xtal-lig.mmff94.tmol
ligand heavy atoms: 19

Generate a matched decoy series#

Every decoy preserves the protein and the ligand’s internal geometry. Only the ligand receives a declared rotation about its centroid and translation in the protein frame. These seven states are sensitivity probes, not samples from a docking search distribution.

[4]:

decoy_specs = [
    ("deposited", 0.0, (0.00, 0.00, 0.00)),
    ("small rotation", 15.0, (0.25, 0.00, 0.00)),
    ("small shift", -20.0, (0.00, 0.50, 0.00)),
    ("mixed 1", 30.0, (0.75, 0.25, 0.00)),
    ("mixed 2", -45.0, (1.00, 0.50, 0.25)),
    ("large shift", 60.0, (1.50, 0.00, 0.00)),
    ("far decoy", 90.0, (2.00, 0.50, 0.00)),
]
ligand_xyz = native_pose.coords[native_ligand_coords]
ligand_center = ligand_xyz.mean(dim=0)
decoy_poses = []
for _, angle_degrees, translation_xyz in decoy_specs:
    decoy = native_pose.clone()
    angle = torch.as_tensor(
        np.deg2rad(angle_degrees),
        device=device,
        dtype=decoy.coords.dtype,
    )
    rotation = torch.eye(3, device=device, dtype=decoy.coords.dtype)
    rotation[0, 0] = torch.cos(angle)
    rotation[0, 1] = -torch.sin(angle)
    rotation[1, 0] = torch.sin(angle)
    rotation[1, 1] = torch.cos(angle)
    translation = torch.as_tensor(
        translation_xyz, device=device, dtype=decoy.coords.dtype
    )
    decoy.coords[native_ligand_coords] = (
        (ligand_xyz - ligand_center) @ rotation.T + ligand_center + translation
    )
    decoy_poses.append(decoy)

decoy_batch = PoseStackBuilder.from_poses(decoy_poses, device)
interaction_scores = ligand_protein_interactions(decoy_batch, score_function).detach()
total_scorer = score_function.render_whole_pose_scoring_module(decoy_batch)
with torch.no_grad():
    total_scores = total_scorer(decoy_batch.coords).detach()
ligand_rmsd = ligand_rmsd_from_native(decoy_batch, native_heavy_coords).detach()

decoy_frame = pd.DataFrame(
    [
        {
            "pose_index": pose_index,
            "state": label,
            "rotation_degrees": angle,
            "translation_A": float(np.linalg.norm(translation)),
            "ligand_heavy_atom_RMSD_A": float(ligand_rmsd[pose_index].cpu()),
            "ligand_protein_interaction_score": float(
                interaction_scores[pose_index].cpu()
            ),
            "whole_pose_score": float(total_scores[pose_index].cpu()),
        }
        for pose_index, (label, angle, translation) in enumerate(decoy_specs)
    ]
)
show_table(decoy_frame)

fig, axis = plt.subplots(figsize=(7, 4.5))
axis.scatter(
    decoy_frame["ligand_heavy_atom_RMSD_A"],
    decoy_frame["ligand_protein_interaction_score"],
    color="#3b82f6",
)
for row in decoy_frame.itertuples():
    axis.annotate(
        row.state,
        (row.ligand_heavy_atom_RMSD_A, row.ligand_protein_interaction_score),
        xytext=(4, 4),
        textcoords="offset points",
        fontsize=8,
    )
axis.set(
    xlabel="ligand heavy-atom RMSD from deposited pose (Å)",
    ylabel="weighted ligand–protein interaction score",
    title="Controlled pose sensitivity in one batched score call",
)
axis.grid(alpha=0.3)
plt.tight_layout()
plt.show()

display(
    tmol.switchable_view(
        {
            label: decoy_batch.split(index)
            for index, (label, _, _) in enumerate(decoy_specs)
        },
        notes={
            row.state: (
                f"ligand RMSD {row.ligand_heavy_atom_RMSD_A:.2f} Å; "
                f"interaction {row.ligand_protein_interaction_score:.2f}"
            )
            for row in decoy_frame.itertuples()
        },
    )
)
ⓘpose_index state rotation_degrees translation_A ligand_heavy_atom_RMSD_A ligand_protein_interaction_score whole_pose_score
0deposited0.00.0000000.00000012.2112671099.029663
1small rotation15.00.2500000.832808607.7517701694.570190
2small shift-20.00.5000001.193737415.8389891502.657471
3mixed 130.00.7905691.7478931741.3051762828.123535
4mixed 2-45.01.1456442.6729541377.4916992464.310303
5large shift60.01.5000003.4025792660.2473143747.066162
6far decoy90.02.0615534.7883923613.1816414699.999512
../_images/tutorial_10_ligand_pose_sensitivity_7_1.png
ligand RMSD 0.00 Å; interaction 12.21

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

The deposited pose has zero displacement by construction. Increasing displacement generally creates steric or solvation penalties in this controlled series, but monotonicity is not guaranteed for arbitrary transformations or complexes. The interaction quantity combines both block-matrix orientations within one bound complex.

Locally minimize three diagnostic states#

The deposited pose, best-scoring non-deposited pose, and worst-scoring non-deposited pose are selected without manual cherry-picking. They are assembled into one batch. Every ligand atom and protein side-chain atom initially within 5 Å of that pose’s ligand may move; protein main-chain atoms remain fixed. A short shared iteration budget is a local response diagnostic, not a converged docking protocol.

[5]:

best_non_native = 1 + int(torch.argmin(interaction_scores[1:]).item())
worst_non_native = 1 + int(torch.argmax(interaction_scores[1:]).item())
selected_indices = [0, best_non_native, worst_non_native]
if len(set(selected_indices)) != 3:
    raise RuntimeError("Expected three distinct diagnostic states")

selected = PoseStackBuilder.from_poses(
    [decoy_batch.split(index) for index in selected_indices], device
)
selected_labels = [decoy_specs[index][0] for index in selected_indices]
selected_ligand = block_mask_for_name3(selected, LIGAND_NAME)
selected_ligand_coords = res_mask_to_coord_mask(selected, selected_ligand)
selected_sidechains = build_sidechain_coord_mask(selected)
near_ligand = torch.zeros_like(selected.real_atoms)
for pose_index in range(selected.n_poses):
    ligand_coords = selected.coords[pose_index, selected_ligand_coords[pose_index]]
    distances = (
        torch.cdist(selected.coords[pose_index].nan_to_num(), ligand_coords)
        .min(dim=1)
        .values
    )
    near_ligand[pose_index] = (
        (distances <= 5.0)
        & selected_sidechains[pose_index]
        & selected.real_atoms[pose_index]
    )
movable = selected_ligand_coords | near_ligand

refinement_scorer = score_function.render_whole_pose_scoring_module(selected)
with torch.no_grad():
    whole_pose_before = refinement_scorer(selected.coords).detach()
interaction_before = ligand_protein_interactions(selected, score_function).detach()
rmsd_before = ligand_rmsd_from_native(selected, native_heavy_coords).detach()
refined = run_cart_min(
    selected,
    score_function,
    coord_mask=movable,
    optimizer_kwargs={"max_iter": 15},
)
if not bool(torch.isfinite(refined.coords[refined.real_atoms]).all()):
    raise RuntimeError("Local refinement produced non-finite coordinates")
with torch.no_grad():
    whole_pose_after = refinement_scorer(refined.coords).detach()
interaction_after = ligand_protein_interactions(refined, score_function).detach()
rmsd_after = ligand_rmsd_from_native(refined, native_heavy_coords).detach()

refinement_frame = pd.DataFrame(
    [
        {
            "state": label,
            "movable_atoms": int(movable[pose_index].sum().item()),
            "interaction_before": float(interaction_before[pose_index].cpu()),
            "interaction_after": float(interaction_after[pose_index].cpu()),
            "interaction_change": float(
                (interaction_after[pose_index] - interaction_before[pose_index]).cpu()
            ),
            "whole_pose_before": float(whole_pose_before[pose_index].cpu()),
            "whole_pose_after": float(whole_pose_after[pose_index].cpu()),
            "whole_pose_change": float(
                (whole_pose_after[pose_index] - whole_pose_before[pose_index]).cpu()
            ),
            "ligand_RMSD_before_A": float(rmsd_before[pose_index].cpu()),
            "ligand_RMSD_after_A": float(rmsd_after[pose_index].cpu()),
            "optimizer_budget": "15-iteration smoke test; convergence not assessed",
        }
        for pose_index, label in enumerate(selected_labels)
    ]
)
show_table(refinement_frame)

viewer_states = {}
viewer_notes = {}
for pose_index, label in enumerate(selected_labels):
    before_label = f"{label} — before"
    after_label = f"{label} — locally minimized"
    viewer_states[before_label] = selected.split(pose_index)
    viewer_states[after_label] = refined.split(pose_index)
    viewer_notes[before_label] = (
        f"interaction {float(interaction_before[pose_index].cpu()):.2f}; "
        f"ligand RMSD {float(rmsd_before[pose_index].cpu()):.2f} Å"
    )
    viewer_notes[after_label] = (
        f"interaction {float(interaction_after[pose_index].cpu()):.2f}; "
        f"ligand RMSD {float(rmsd_after[pose_index].cpu()):.2f} Å; "
        f"whole-pose Δ {float((whole_pose_after[pose_index] - whole_pose_before[pose_index]).cpu()):+.2f}"
    )
display(tmol.switchable_view(viewer_states, notes=viewer_notes))
ⓘstate movable_atoms interaction_before interaction_after interaction_change whole_pose_before whole_pose_after whole_pose_change ligand_RMSD_before_A ligand_RMSD_after_A optimizer_budget
deposited19912.211267-20.595383-32.8066481099.0297851047.216553-51.8132320.0000000.19986715-iteration smoke test; convergence not assessed
small shift198415.838989-7.532748-423.3717351502.6574711068.502930-434.1545411.1937370.89947215-iteration smoke test; convergence not assessed
far decoy2133613.18164158.827770-3554.3537604699.9995121492.751221-3207.2482914.7883924.33990515-iteration smoke test; convergence not assessed
interaction 12.21; ligand RMSD 0.00 Å

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

Local minimization may lower the objective without recovering the deposited pose.

Exercises#

  1. Add rotations around the other two principal axes without changing the native member.

  2. Increase the local minimizer budget and check whether score and geometry stabilize.

  3. Repeat the selected-state refinement with several pocket cutoffs.

  4. Decompose the ligand–protein interaction into weighted score terms.

  5. Construct independent rigid transformations before examining scores, then report rank correlation rather than selecting transformations after inspection.

References#

Download this notebook