TMol How-To Guide — Advanced Topics#

Open In Colab

8. Advanced Topics#

Two unrelated deep-end topics: customizing which identities the packer offers at a position (beyond the single-target restrict_absent_name3s from section 1), and building a PoseStack directly from fixed-width atom37 tensors – the representation most structure-prediction models output – while keeping the whole thing differentiable.

Setup#

Wheel install, imports, seeding, and device selection live in notebook_setup.py.

[1]:

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

if IN_COLAB:
    from urllib.request import urlretrieve

    urlretrieve(
        "https://raw.githubusercontent.com/RosettaCommons/tmol/"
        "main/docs/tutorial/notebook_setup.py",
        "notebook_setup.py",
    )

from notebook_setup import *  # noqa: F401,F403

device, param_db = setup()

Environment variable CCD_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
  (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
  (2) set it in your current shell with 'export VAR_NAME=path/to/variable'
  (3) write it to a .env file in the root of the atomworks.io repository
Environment variable PDB_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
  (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
  (2) set it in your current shell with 'export VAR_NAME=path/to/variable'
  (3) write it to a .env file in the root of the atomworks.io repository

Extending the packer for design#

PackerPalette#

PackerPalette decides, per position, which identities are ever candidates; restrict_absent_name3s (section 1) then narrows that set further from outside. Subclass PackerPalette and override block_types_from_original to change the candidate set itself – call super().block_types_from_original(pbt, orig) for the defaults, then filter or reorder them. Unlike restrict_absent_name3s, this applies uniformly everywhere the palette is used, not one position at a time.

A worked design example#

A palette that only ever offers a small hydrophobic alphabet (plus whatever the position’s own original identity is, so a non-matching position can still “keep its own”):

[2]:

class HydrophobicPalette(PackerPalette):
    """Allow a small hydrophobic alphabet plus each position's own identity."""

    def __init__(self, allowed_name3s=("ALA", "VAL", "LEU", "ILE", "PHE")):
        super().__init__()
        self.allowed_name3s = frozenset(allowed_name3s)

    def block_types_from_original(self, pbt, orig):
        n_allowed, allowed, is_original = super().block_types_from_original(pbt, orig)
        in_alphabet = torch.tensor(
            [block_type.name3 in self.allowed_name3s for block_type in pbt.active_block_types],
            dtype=torch.bool, device=pbt.device,
        )
        real_choice = allowed >= 0
        keep = torch.zeros_like(allowed, dtype=torch.bool)
        keep[real_choice] = in_alphabet[allowed[real_choice]]
        keep = (keep | is_original) & real_choice

        order = torch.argsort((~keep).to(torch.int64), dim=2, stable=True)
        allowed = torch.gather(allowed, 2, order)
        is_original = torch.gather(is_original, 2, order)
        keep = torch.gather(keep, 2, order)
        allowed = torch.where(keep, allowed, torch.full_like(allowed, -1))
        is_original = is_original & keep
        n_allowed = keep.sum(dim=2)
        return n_allowed, allowed, is_original

[3]:

cif_text = fetch_cif("1UBQ")
pose_stack = pose_stack_from_file(io.StringIO(cif_text), device, param_db=param_db)

viewer = tmol.view(pose_stack)
viewer.show()

score_function = beta2016_score_function(device, param_db=param_db)

target_block = 5
palette = HydrophobicPalette()
design_task = PackerTask(pose_stack, palette)
design_region = torch.zeros_like(design_task.is_real_block)
design_region[:, target_block] = True
design_task.disable_packing_by_block_mask(~design_region)
design_task.add_conformer_sampler(create_dunbrack_sampler_from_database(param_db, device))
design_task.add_conformer_sampler(FixedAAChiSampler())
design_task.add_conformer_sampler(IncludeCurrentSampler())

torch.manual_seed(SEED)
designed = pack_rotamers(pose_stack, score_function, design_task)
before, after = pose_stack.block_type(0, target_block).name3, designed.block_type(0, target_block).name3
print(f"block {target_block}: {before} -> {after}")
print(f"designed identity allowed by the palette: {after in palette.allowed_name3s or after == before}")

display(tmol.overlay_view(
    {"original": pose_stack, "designed": designed},
    notes={
        "original": f"block {target_block}: {before}",
        "designed": f"block {target_block}: {after}",
    },
))

/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.

block 5: LYS -> LEU
designed identity allowed by the palette: True

/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"))

Populate a PoseStack from atom37#

“atom37” is the fixed-width, 37-atom-slot-per-residue layout most structure-prediction models (AlphaFold-style) use. tmol.io.pose_stack_from_atom37 builds a PoseStack directly from such a tensor plus a PoseBuildContext, entirely on tensor ops with no AtomArray involved. Because of that, gradients flow straight back from the built pose’s coordinates to the atom37 input – a model’s output coordinates can be scored by TMol and the loss backpropagated through pose construction with nothing detached along the way.

We still need an AtomArray once here, to resolve chemistry and pull a real structure’s canonical-form tensors (res_types, chain_id, bond records) for re-populating into atom37 slots – a real caller already has these tensors from its own model and skips this step entirely.

[4]:

from tmol.io import (
    atom_array_from_file, build_context_from_atom_array, canonical_form_from_biotite,
    pose_stack_from_atom37, atom37_slot_map_for_ordering,
)

atom37_cif_text = fetch_cif("1UBQ")
atom37_structure = atom_array_from_file(io.StringIO(atom37_cif_text))
atom37_context = build_context_from_atom_array(atom37_structure, device, param_db=param_db)
canonical_ordering = atom37_context.canonical_ordering

canonical_form = canonical_form_from_biotite(
    atom37_structure, device, co=canonical_ordering, missing_density_distance_threshold=0.0,
)


def full_atom_layout(canonical_ordering):
    """Every canonical atom gets its own atom37 slot, in order -- the simplest possible layout."""
    return {
        name3: list(canonical_ordering.restypes_ordered_atom_names[name3])
        for name3 in canonical_ordering.restype_io_equiv_classes
    }


def scatter_to_slots(coords, res_types, slot_map):
    """Spread canonical-form coordinates out into atom37 slot order."""
    n_poses, n_res = res_types.shape
    n_slots = slot_map.shape[1]
    is_real = res_types >= 0
    canonical_atom = slot_map[torch.where(is_real, res_types, torch.zeros_like(res_types)).to(torch.int64)]
    routed = (canonical_atom >= 0) & is_real.unsqueeze(-1)
    slots = torch.full((n_poses, n_res, n_slots, 3), float("nan"), dtype=torch.float32, device=coords.device)
    pose_index = torch.arange(n_poses, device=coords.device).reshape(-1, 1, 1).expand(n_poses, n_res, n_slots)
    res_index = torch.arange(n_res, device=coords.device).reshape(1, -1, 1).expand(n_poses, n_res, n_slots)
    slots[routed] = coords[pose_index[routed], res_index[routed], canonical_atom[routed]]
    return slots


slot_map = atom37_slot_map_for_ordering(canonical_ordering, full_atom_layout(canonical_ordering), device)
atom37_coords = scatter_to_slots(canonical_form.coords, canonical_form.res_types, slot_map)
atom37_coords.requires_grad_(True)
print(f"atom37 tensor shape: {tuple(atom37_coords.shape)} (pose, residue, slot, xyz)")

atom37 tensor shape: (1, 76, 37, 3) (pose, residue, slot, xyz)

pose_stack_from_atom37 takes the atom37 coordinates, a residue-type tensor, a chain-id tensor, and a PoseBuildContext – the same kind of context object built in notebook 4, here without any ligand.

[5]:

atom37_pose_stack = pose_stack_from_atom37(
    atom37_coords, canonical_form.res_types, canonical_form.chain_id, atom37_context,
    slot_map=slot_map, disulfides=canonical_form.disulfides, cyclic_bonds=canonical_form.cyclic_bonds,
    covalent_bonds=canonical_form.covalent_bonds, res_not_connected=canonical_form.res_not_connected,
)
print(f"n_poses={atom37_pose_stack.n_poses}, max_n_blocks={atom37_pose_stack.max_n_blocks}")
print(f"pose coordinates require grad: {atom37_pose_stack.coords.requires_grad}")

viewer = tmol.view(atom37_pose_stack)
viewer.show()

n_poses=1, max_n_blocks=76
pose coordinates require grad: True

/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 that pose and calling .backward() reaches all the way back to the original atom37 tensor – the gradient a model’s own training loop would need.

[6]:

atom37_score_function = beta2016_score_function(device, param_db=atom37_context.parameter_database)
atom37_scorer = atom37_score_function.render_whole_pose_scoring_module(atom37_pose_stack)
atom37_total_score = atom37_scorer(atom37_pose_stack.coords)
print(f"score: {float(atom37_total_score[0].detach()):.2f}")

atom37_total_score.sum().backward()
print(f"gradient reached the atom37 input: {torch.count_nonzero(atom37_coords.grad).item() > 0}")

score: 261.03
gradient reached the atom37 input: True

Download this notebook