Enzyme design tutorial: from a theozyme to hit hunting and diversification#

🚀 Try it live — no local setup needed

In this tutorial you will use the RFD4-Proteina’s Python API along and Atomworks’ selection syntax to design enzymes and then evaluate them using RFD4-proteina’s built-in metrics.

Starting from just a handful of catalytic residues arranged around a ligand, you will design a scaffold that holds them exactly where they need to be.

Goal. Take a theozyme, run the full enzyme design workflow, and read off the results.

Inspect the input: see what coordinates and identities are fixed, and the region left to design.
Design a scaffold around the catalytic constellation and score it against the AME (Atomic Motif Enzyme) criteria.
Select the best design and inspect its saved conditions before refinement.
Query the design with the selection language (optional).
Redesign the sequence on that design with inverse folding.
Refine it with partial diffusion.

Each section opens with a settings block holding the parameters that matter there. Outputs go to directories named after those parameters, so sweeping a value keeps your earlier results. See the CLI reference for every knob.

Before you get started, install the native model environment with make setup, then add the optional development kernel with bash setup.sh --venv and run notebook cells in that environment. Sampling goes to Slurm where one is available; on a host with no scheduler, the Setup cell below detects that and runs rfd4-generate directly instead. Either way, scoring designs against the AME criteria needs a refolding oracle: this tutorial defaults to OpenFold3, free and Apache-licensed with no gate. See https://github.com/RosettaCommons/RFD4-Proteina/blob/main/docs/installation.rst to configure it.

New to this? What’s a “theozyme”?

“Theozyme” = theoretical enzyme. Instead of starting from a whole folded protein, you strip everything away except the handful of atoms that actually do the chemistry (an acid group here, a base there, something to stabilize a charged intermediate) arranged in the exact 3D positions a reaction needs, floating around the substrate with no protein scaffold holding them together yet. It’s the minimal chemical machinery, not an enzyme.

Think lock and key: the protein you’re designing is the lock, the substrate is the key, and the theozyme’s fixed atoms are the pins inside the lock. The working parts that actually engage the key, not the lock’s shape itself. This tutorial’s whole job is to build a lock (a real, physically valid protein) around pins that are already known to be correctly placed.

Setup#

Paths, checkpoint, Slurm resources (used only if Slurm is available), and a preflight check for the refolding oracle. Every path has a default and takes an environment variable override, so the notebook runs elsewhere without you editing it.

The pass criteria and the helper functions below are also part of this Setup step. Everything up to 1. Inspect the input is setup, make sure to run these cells before moving forward with the tutorial.

Have AlphaFold3’s licensed weights instead?

The Setup cell’s auto-detection uses whichever refolding oracle it finds at AF3_MODEL_DIR: if that directory holds an OpenFold3 provenance file (of3_ported_weights.provenance.json), it’s treated as OpenFold3; otherwise it’s treated as AlphaFold3. So if you have access to AlphaFold3’s gated weights and would rather use those, just point AF3_MODEL_DIR at them instead of the OpenFold3 weights, no other change is needed. See docs/installation.rst to configure either path.

[ ]:

import warnings
warnings.filterwarnings("ignore", message=".*NVML.*")
import logging
logging.getLogger("atomworks.constants").setLevel(logging.ERROR)
import os
from pathlib import Path

# Locate the repository by walking up to the `.project-root` marker; $PROJECT_ROOT is the
# fallback when the notebook is run from outside the tree.
REPO_ROOT = Path.cwd()
while not (REPO_ROOT / ".project-root").exists() and REPO_ROOT.parent != REPO_ROOT:
    REPO_ROOT = REPO_ROOT.parent
if not (REPO_ROOT / ".project-root").exists():
    _root = os.environ.get("PROJECT_ROOT")
    if not _root:
        raise RuntimeError(
            "Could not find the repository root. Run this notebook from inside the "
            "RFD4-Proteina checkout (it contains a `.project-root` marker), or set the "
            "PROJECT_ROOT environment variable to the checkout path."
        )
    REPO_ROOT = Path(_root)

# Keep shell settings; local use needs no private runtime profile.
os.environ.setdefault("CLUSTER", "local")
DATA_PATH = os.environ.setdefault("DATA_PATH", str(REPO_ROOT / "data_path"))

# Tutorial assets: pre-built theozymes and the helper modules used below.
TUTORIAL_DIR = Path(os.environ.get("ENZYME_TUTORIAL_DIR", REPO_ROOT / "docs/guides/tutorials/enzyme_design"))
OUT_ROOT = Path(os.environ.get("ENZYME_TUTORIAL_RUNS", REPO_ROOT / "runs/enzyme_design_tutorial_v3"))
OUT_ROOT.mkdir(parents=True, exist_ok=True)

# Auto-discover OpenFold3 assets downloaded by the setup script.
# (`docs/installation.rst`), if nothing has already set these -- a
# named cluster's `.ipd/.env.<cluster>` sets these directly and is left alone by setdefault.
_of3_root = Path(DATA_PATH) / "openfold3"
# A cache dir, not a data path: safe (and necessary) to create rather than require pre-existing.
os.environ.setdefault("AF3_JAX_CACHE_DIR", str(_of3_root / "jax_cache"))
Path(os.environ["AF3_JAX_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
if not os.environ.get("AF3_IMAGE") and (_of3_root / "runtime").exists():
    _sifs = sorted((_of3_root / "runtime").glob("*.sif"))
    if _sifs:
        os.environ["AF3_IMAGE"] = str(_sifs[0])
if not os.environ.get("AF3_MODEL_DIR"):
    _of3_dir = _of3_root / "weights"
    if (_of3_dir / "of3_ported_weights.provenance.json").exists():
        os.environ["AF3_MODEL_DIR"] = str(_of3_dir)

CKPT_PATH = os.environ.get("LATENT_CKPT_PATH", f"{DATA_PATH}/flow_checkpoints")
# Always run the repository's current SOTA checkpoint: read it from configs/inference.yaml so
# the tutorial tracks whatever the code ships. Override with $LATENT_CKPT_NAME. The resolved
# name is printed in this block's output below.
def _repo_ckpt_name():
    import re
    inf = REPO_ROOT / "configs/inference.yaml"
    if inf.exists():
        m = re.search(r"^\s*ckpt_name:\s*(\S+)", inf.read_text(), re.M)
        if m:
            return m.group(1)
    raise RuntimeError("Could not read ckpt_name from configs/inference.yaml; set LATENT_CKPT_NAME.")
CKPT_NAME = os.environ.get("LATENT_CKPT_NAME") or _repo_ckpt_name()

PARTITION = os.environ.get("SLURM_PARTITION", "gpu")
GRES      = os.environ.get("SLURM_GRES", "gpu:large:1")
MEM       = "32g"
TIME      = "3:00:00"
NJOBS     = 5      # replicates are split across this many array tasks, one GPU each; Slurm only
OVERWRITE_SAMPLING = False   # False reuses any existing outputs and skips sampling; True re-runs it

# Metric tags, matching the benchmark configuration so the numbers are comparable.
# `design` supplies the clash metric the third pass criterion needs; without it `ame.pass`
# is silently zero. `ame_refolding` needs a working refolding oracle -- see the preflight
# check below, which surfaces a missing one now rather than after several minutes of sampling.
METRICS_TAGS = "[design,unindexed,3d_coordinate_conditioning,ame_refolding]"

import re, sys, glob, shutil, time, subprocess
import numpy as np, pandas as pd

WT = str(REPO_ROOT)                     # kept for the Hydra command below
PY = f"{REPO_ROOT}/.venv/bin/python3"   # interpreter for scripts invoked below
sys.path.insert(0, str(TUTORIAL_DIR))
sys.path.insert(0, f"{REPO_ROOT}/src")
# Subprocesses do not inherit sys.path; pass it through PYTHONPATH.
os.environ["PYTHONPATH"] = f"{REPO_ROOT}/src:" + os.environ.get("PYTHONPATH", "")
os.environ.setdefault("PROJECT_ROOT", str(REPO_ROOT))
# CCD_MIRROR_PATH / PDB_MIRROR_PATH are left unset: neither this tutorial nor plain generation
# needs them, and pointing either at a path that does not exist turns a harmless "not configured"
# warning into a hard crash the moment anything does.

# --- Slurm vs local fallback -------------------------------------------------------------------
# Sampling submits to Slurm (`scripts/launch_generate.sh`) when one is available. On a standalone
# GPU host (e.g. a single cloud instance) there is no scheduler to submit to, so this runs
# `rfd4-generate` directly instead -- slower per batch of replicates (no parallel array tasks
# across GPUs), but otherwise identical.
HAS_SLURM = shutil.which("sbatch") is not None

# --- Refolding oracle preflight ------------------------------------------------------------------
# `ame_refolding` in METRICS_TAGS makes every later step depend on a working refolding oracle
# (AF3, or its free/ungated alternative OpenFold3). Check that now instead of letting the first
# `launch()` call run sampling to completion and then fail deep in a stack trace.
_AF3_VARS = ("AF3_IMAGE", "AF3_MODEL_DIR", "AF3_JAX_CACHE_DIR")
_af3_missing = [v for v in _AF3_VARS if not os.environ.get(v)]
_af3_bad_paths = [v for v in _AF3_VARS if os.environ.get(v) and not Path(os.environ[v]).exists()]
_IS_OF3 = (
    not _af3_missing and not _af3_bad_paths
    and (Path(os.environ["AF3_MODEL_DIR"]) / "of3_ported_weights.provenance.json").exists()
)
# `of3.yaml` only renames the artifact namespace (oracle_name: af3 -> of3); the actual OF3 vs AF3
# model code path is auto-detected from the provenance file next to the weights. The override is
# only added when OF3 is detected, so an AF3 setup (no provenance file present) is untouched.
_OF3_OVERRIDE = "metrics/refolding_oracles=of3 " if _IS_OF3 else ""

if _af3_missing or _af3_bad_paths:
    print(
        "⚠️  Refolding oracle not configured "
        f"(missing: {_af3_missing or 'none'}; not found on disk: {_af3_bad_paths or 'none'}).\n"
        "   Every launch() call below will run sampling to completion and then fail when it "
        "tries to refold for ame.pass scoring.\n"
        "   No AF3 (Google DeepMind gated weights) access? OpenFold3 is a free, Apache-licensed "
        "alternative using the same runtime -- see "
        "'https://github.com/RosettaCommons/RFD4-Proteina/blob/main/docs/installation.rst' for the setup commands."
    )
else:
    print(f"refolding oracle: configured ({'OpenFold3' if _IS_OF3 else 'AF3'})")

def run(cmd, echo=True, env=None):
    e = {**os.environ, **(env or {})}
    p = subprocess.run(cmd, shell=True, cwd=TUTORIAL_DIR, capture_output=True, text=True, env=e)
    if echo and p.stdout.strip(): print(p.stdout.rstrip())
    if p.returncode != 0 and p.stderr.strip(): print("STDERR:", p.stderr[-1200:])
    return p.stdout

def tag(v):
    """Render a parameter value as a filename-safe suffix (1.4 -> 1p4)."""
    return str(v).replace(".", "p").replace("-", "m")


# TSR k is the primary temperature control for each modality; higher k means lower temperature.
BB_CA_TSR_K_KEY = "generation.model.bb_ca.simulation_step_params.tsr_k"
LOCAL_LATENTS_TSR_K_KEY = "generation.model.local_latents.simulation_step_params.tsr_k"


def override(ov, name, key, value, label):
    """Append a Hydra override when the value is set, and extend the run name accordingly.

    A value of None leaves the configuration default in place and the name unchanged, so the
    output directory records only what was actually varied.
    """
    if value is None:
        return name
    ov.append(f"{key}={value}")
    return name + f"_{label}{tag(value)}"

STAGE_TIMES = []          # (run_name, seconds) per sampling stage

def launch(out_dir, run_name, inputs, num_replicates, overrides=(), batch_size=8,
           njobs=None, timeout=5400):
    """Submit a sampling run (Slurm) or run it directly (no Slurm), then return its metrics CSV.

    On a Slurm cluster, replicates are spread over `njobs` array tasks, one GPU each, and merged
    by the concatenation job the launch script queues behind them, so a run finishes in roughly
    the time of `num_replicates / njobs` sequential samples. Without Slurm (`HAS_SLURM` is
    False), all `num_replicates` run sequentially in one `rfd4-generate` call on this GPU, batched
    by `batch_size` -- correct, just without the cross-GPU parallelism.

    Every call submits/runs; nothing is reused unless `OVERWRITE_SAMPLING` is False and a prior
    run's output already exists. Output directories carry the parameter values, so different
    settings never collide, while repeating one setting resamples in place.

    Args:
        overrides: extra Hydra overrides carrying this section's settings into the job.
        njobs: number of Slurm array tasks; defaults to NJOBS, capped at the replicate count.
            Ignored when running without Slurm.
    """
    csv = f"{out_dir}/inference_metrics/per_example_metrics.csv"
    if not OVERWRITE_SAMPLING and os.path.exists(csv):
        print(f"[skip] {run_name}: reusing {len(pd.read_csv(csv))} existing designs "
              f"(set OVERWRITE_SAMPLING=True to re-sample)")
        return csv
    t_start = time.time()                       # wall clock: submit/start -> metrics on disk
    # Clear any previous run for these settings. Without this the poll below would match the
    # stale metrics file and return before the new run has produced anything.
    shutil.rmtree(out_dir, ignore_errors=True)
    os.makedirs(f"{out_dir}/inference_samples", exist_ok=True)

    if HAS_SLURM:
        n_tasks = min(njobs or NJOBS, num_replicates)
        cmd = (f'bash "{WT}/scripts/launch_generate.sh" {os.environ["CLUSTER"]} '
               f'--njobs {n_tasks} --partition {PARTITION} --gres {GRES} --mem {MEM} --time {TIME} '
               f'--code-dir "{WT}" '
               f'inference_experiment=inference_on_user_inputs '
               f'inputs="{inputs}" num_replicates={num_replicates} batch_size={batch_size} '
               f'metrics_tags={METRICS_TAGS} {_OF3_OVERRIDE}'
               + " ".join(overrides) + " "
               f'ckpt_path={CKPT_PATH} ckpt_name={CKPT_NAME} '
               f'out_dir="{out_dir}" run_name_={run_name}')
        out = run(cmd, echo=False)                      # the launch script prints a long banner
        jid = re.search(r"Submitted batch job (\d+)", out)
        print(f"[submit] {run_name}: {num_replicates} designs over {n_tasks} tasks"
              + (f" (job {jid.group(1)})" if jid else "")
              + (f"  {' '.join(overrides)}" if overrides else ""))
        deadline = time.time() + timeout
        while time.time() < deadline:
            if os.path.exists(csv):
                dt = time.time() - t_start
                STAGE_TIMES.append((run_name, dt))
                print(f"[done] {run_name}: {len(pd.read_csv(csv))} designs in {dt/60:.1f} min")
                return csv
            time.sleep(60)
        raise TimeoutError(f"{run_name} produced no metrics before the timeout")

    # No Slurm: run synchronously on this GPU using the notebook kernel's environment.
    generate_command = Path(sys.executable).with_name("rfd4-generate")
    if os.environ.get("CLUSTER", "local") != "local" and (REPO_ROOT / ".ipd").is_dir():
        generate_command = REPO_ROOT / "src/rfproteina/generate.py"
    cmd = (f'"{generate_command}" '
           f'inference_experiment=inference_on_user_inputs '
           f'inputs="{inputs}" num_replicates={num_replicates} batch_size={batch_size} '
           f'metrics_tags={METRICS_TAGS} {_OF3_OVERRIDE}'
           + " ".join(overrides) + " "
           f'ckpt_path={CKPT_PATH} ckpt_name={CKPT_NAME} '
           f'out_dir="{out_dir}" run_name_={run_name}')
    print(f"[submit] {run_name}: {num_replicates} designs, batch_size={batch_size}, "
          f"running locally on this GPU (no Slurm detected)"
          + (f"  {' '.join(overrides)}" if overrides else ""))
    run(cmd, echo=True)   # blocks until rfd4-generate exits; CLUSTER is already in os.environ
    dt = time.time() - t_start
    if not os.path.exists(csv):
        raise RuntimeError(f"{run_name} produced no metrics -- see output above for the error")
    STAGE_TIMES.append((run_name, dt))
    print(f"[done] {run_name}: {len(pd.read_csv(csv))} designs in {dt/60:.1f} min")
    return csv


print("repository   :", REPO_ROOT)
print("tutorial dir :", TUTORIAL_DIR)
print("run outputs  :", OUT_ROOT)
print("checkpoint   :", CKPT_NAME)
print("scheduler    :", "Slurm" if HAS_SLURM else "none (running rfd4-generate on this GPU)")

Pass criteria#

AME is the in silico enzyme motif scaffolding benchmark for de novo enzyme design and atom level active site placement. A design passes, with ame.pass == 1, when it clears all three criteria at once.

Criterion	Threshold
`refolded_backbone_rmsd`	≤ 2.0 Å
`refolded_motif_allatom_rmsd`	≤ 1.5 Å
`clash_metrics` protein and ligand	≤ 0

ame_matrix() prints one row per design showing the exact values, green where a criterion is met and red where it is not. Quality metrics follow uncoloured. Deviation from the parent scaffold, where it applies, is grouped separately on a blue background.

New to this? What do these numbers actually mean?

RMSD (root-mean-square deviation) is just “how far apart, on average, are two versions of the same atoms, measured in Ångströms (Å).” Lower is better: 0 would mean a perfect match, and a couple of Å means the atoms have drifted apart by roughly an atom’s width or two.

The two thresholds check different things. refolded_backbone_rmsd asks “does the overall protein shape hold together?” An independent AI structure-prediction model (e.g. OpenFold3 or AlphaFold3) is given only the sequence and asked to predict the structure from scratch; if its answer matches what was generated, that’s a strong sign the design is real and not just something that looked good to the model that made it. refolded_motif_allatom_rmsd asks the harder, more important question: did the exact catalytic atoms survive that independent re-check in the right place? A design can hold its overall shape together while still failing this tighter check.

The third gate, clash_metrics, isn’t an RMSD at all, it just checks that no protein atom and ligand atom are sitting on top of each other. A physical overlap like that means the structure isn’t buildable in reality, no matter how good the two RMSDs look.

[ ]:

AME = {
    ("refolded_backbone_rmsd.rmsd", "<="): 2.0,
    ("refolded_motif_allatom_rmsd.rmsd", "<="): 1.5,
    ("clash_metrics.frac_atoms_clashing_inter_residue_polymer_ligand_1p5A", "<="): 0.0,
}
LABEL = {
    "refolded_backbone_rmsd.rmsd": "refold_bb <=2.0",
    "refolded_motif_allatom_rmsd.rmsd": "refold_motif <=1.5",
    "clash_metrics.frac_atoms_clashing_inter_residue_polymer_ligand_1p5A": "clash_lig <=0",
}

def _ok(v, op, t):
    if pd.isna(v): return None
    return {"<=": v <= t, ">=": v >= t, "<": v < t, ">": v > t, "==": v == t}[op]

def ame_matrix(csv, sort_key="design_unindexed_motif_allatom_rmsd.rmsd",
               context=("design_motif_allatom_rmsd.rmsd", "oracle.overall_plddt"),
               reference=None):
    """Tabulate one row per design, grouped by shading.

    Pass criteria are coloured green or red, quality metrics are left plain, and deviation from
    the parent scaffold is shaded blue. Deviation measures distance from the starting point
    rather than design quality, so it forms its own group and takes no part in the verdict.

    Args:
        reference: parent structure; when given, adds the two deviation columns.
    """
    DEV = ("ca_rmsd_to_parent", "seq_ident_to_parent")
    if reference is not None:
        d = deviation_table(csv, reference).reset_index(drop=True)
    else:
        d = pd.read_csv(csv).reset_index(drop=True)
    # Capture rep and job together: with --njobs>1 each array task restarts replicate
    # numbering, so `rep\d+` alone collides after the tasks are concatenated, and the styled
    # table below sets this as its index (Styler requires a unique index).
    _base = d["example_id"].astype(str).str.extract(r"(rep\d+(?:_job\d+)?|_\d+_job)", expand=False)
    _base = _base.fillna(pd.Series([f"s{i}" for i in range(len(d))], index=d.index))
    _cc = _base.groupby(_base).cumcount()
    d["_id"] = [f"{b}#{k}" if (_base == b).sum() > 1 else b for b, k in zip(_base, _cc)]
    verd = {}
    for (col, op), t in AME.items():
        verd[col] = d[col].map(lambda v: _ok(v, op, t)) if col in d else pd.Series([None]*len(d))
    d["_npass"] = [sum(1 for (c, _op) in AME if verd[c].iloc[i] == True) for i in range(len(d))]
    sk = sort_key if sort_key in d else "design_motif_allatom_rmsd.rmsd"
    order = d.sort_values(["_npass", sk], ascending=[False, True]).index

    for c, _op in AME:                       # a metric absent from this run shows as n/a
        if c not in d:
            d[c] = float("nan")
    ame_cols = [c for (c, _op) in AME]
    qual_cols = [c for c in context if c in d]
    dev_cols = [c for c in DEV if c in d]
    cols = ame_cols + qual_cols + dev_cols
    disp = d.loc[order, ["_id"] + cols + ["_npass"]].copy()
    names = ([LABEL.get(c, c.split(".")[0]) for c in ame_cols + qual_cols]
             + list(dev_cols))
    disp.columns = ["sample"] + names + ["PASS /3"]
    disp = disp.set_index("sample")
    ver_by_disp = {LABEL[c]: [verd[c].loc[i] for i in order] for (c, op), t in AME.items()}
    dev_names = set(names[len(ame_cols) + len(qual_cols):])

    def color(df):
        out = pd.DataFrame("", index=df.index, columns=df.columns)
        for col in df.columns:
            if col == "PASS /3":
                out[col] = "font-weight:bold;background-color:#dfe7ef;color:#123"; continue
            if col in dev_names:                       # deviation from parent: shaded blue
                out[col] = "background-color:#dce9f5;color:#123a5c;border-left:2px solid #6a9fd0"
                continue
            if col not in ver_by_disp: continue        # quality metrics: unshaded
            for r, v in zip(df.index, ver_by_disp[col]):
                out.loc[r, col] = ("background-color:#e6e6e6;color:#555" if v is None
                    else "background-color:#a8dba8;color:#14431a" if v
                    else "background-color:#e79a9a;color:#6a0f0f")
        return out

    def fmt(v):
        if pd.isna(v): return "n/a"
        try: return f"{float(v):.4g}"
        except Exception: return str(v)

    n_pass = int((d["_npass"] == 3).sum())
    print(f"ame.pass {n_pass}/{len(d)} = {n_pass/len(d)*100:.0f}%")
    return disp.style.apply(color, axis=None).format(fmt), d


def describe_conditions(path, title=""):
    """Print what a CIF fixes: conditioned atom counts and the catalytic residues.

    Used on both the theozyme and the design prepared for chaining, so the two are read
    with one implementation and can be compared directly.
    """
    from atomworks.io.utils.annotator import ensure_annotations
    from rfproteina.inference_utils.datasets import DesignProblemDataset

    aa = DesignProblemDataset(dir_path_or_file_paths=str(path), transform=None,
                              num_replicates=1)[0]["atom_array"]
    ensure_annotations(aa, "is_ligand")
    prot = np.asarray(aa.mask("is_polymer"))
    crd, seq, idx = (np.asarray(aa.mask(q)) for q in
                     ("sa_mask('C_CRD')", "sa_mask('C_SEQ')", "sa_mask('C_IDX')"))
    motif = prot & crd
    if title:
        print(title)
    print(f"  atoms {aa.array_length()}  |  protein {int(prot.sum())}  ligand {int((~prot).sum())}")
    print(f"  C_CRD coordinates fixed: {int(crd.sum())}  "
          f"(protein {int((crd & prot).sum())} + ligand {int((crd & ~prot).sum())})")
    print(f"  C_SEQ identity fixed:    {int(seq.sum())}  "
          f"(protein {int((seq & prot).sum())} + ligand {int((seq & ~prot).sum())})")
    print(f"  designable residues:     {len(set(aa.res_id[prot & idx].tolist()))}")
    for rid in sorted(set(aa.res_id[motif].astype(int))):
        m = motif & (aa.res_id == rid)
        print(f"    res {rid:>4} {aa.res_name[m][0]:<4} {sorted(aa.atom_name[m].tolist())}")
    return aa

Helpers#

These functions are defined here and folded by default. Expand a block to read it.

pick_best_design copies the top ranked passing design with its saved conditioning targets unchanged.

[ ]:

def pick_best_design(out_dir, verbose=False):
    """Copy the top-ranked passing design with all saved annotations unchanged.

    Rank by ``ame.pass``, then by ``design_unindexed_motif_allatom_rmsd``.
    Return the copied structure path, preserving its compression suffix.
    """
    import shutil
    from pathlib import Path

    PASS, KEY = "ame.pass", "design_unindexed_motif_allatom_rmsd.rmsd"
    best_dir = Path(out_dir) / "best"
    ranked = (pd.read_csv(f"{out_dir}/inference_metrics/per_example_metrics.csv")
              .sort_values([PASS, KEY], ascending=[False, True]).reset_index(drop=True))
    passed = ranked[ranked[PASS] == 1]
    if passed.empty:
        raise RuntimeError(f"no design passed ame.pass out of {len(ranked)}; review the run")
    best = passed.iloc[0]
    source = Path(best["designed_structure_path"])
    best_dir.mkdir(parents=True, exist_ok=True)
    dst = best_dir / source.name
    shutil.copy2(source, dst)
    ranked.to_csv(best_dir / "selection.csv", index=False)

    if verbose:
        print(f"selected {best['example_id']}: {len(passed)} of {len(ranked)} passed, "
              f"motif placed to {best[KEY]:.3f} A")
        print(f"  {dst} (saved conditioning targets preserved unchanged)")
        print(f"  {best_dir / 'selection.csv'} (full ranking)")
    return str(dst)

deviation_table adds backbone RMSD and sequence identity for each design, relative to the parent scaffold. Neither CCD_MIRROR_PATH nor PDB_MIRROR_PATH is needed for anything in this tutorial; the Setup cell suppresses the benign “not set” warning for both so it doesn’t clutter output here.

[ ]:

from __future__ import annotations
import biotite.structure as struc
from atomworks.io import parse
from atomworks.io.utils.annotator import ensure_annotations
from rfproteina.utils.parse_presets import ANNOTATED_CIF


def _ca_and_seq(path):
    """Return the protein C-alpha coordinates and per-residue names, skipping ligands and gaps."""
    aa = parse(str(path), config=ANNOTATED_CIF)["assemblies"]["1"][0]
    ensure_annotations(aa, "is_ligand")
    sel = aa.mask("(atom_name == 'CA') & ~hetero & ~has_nan_coord()")
    return aa.coord[sel], aa.res_name[sel]


def deviation(generated, reference):
    """``ca_rmsd_to_parent`` and ``seq_ident_to_parent`` for one structure.

    NaN when the residue counts differ, rather than truncating to a common length and comparing
    positions that do not correspond.
    """
    gc, gs = _ca_and_seq(generated)
    rc, rs = _ca_and_seq(reference)
    if len(gc) != len(rc) or len(gc) == 0:
        return {"ca_rmsd_to_parent": float("nan"), "seq_ident_to_parent": float("nan"), "n_res": len(gc)}
    fitted, _ = struc.superimpose(rc, gc)
    rmsd = float(np.sqrt(((fitted - rc) ** 2).sum(-1).mean()))
    ident = float((gs == rs).mean())
    return {"ca_rmsd_to_parent": rmsd, "seq_ident_to_parent": ident, "n_res": len(gc)}


def deviation_table(metrics_csv, reference):
    """Add ``ca_rmsd_to_parent`` and ``seq_ident_to_parent`` columns to a metrics table."""
    d = pd.read_csv(metrics_csv)
    rows = [deviation(p, reference) for p in d["designed_structure_path"]]
    for k in ("ca_rmsd_to_parent", "seq_ident_to_parent"):
        d[k] = [r[k] for r in rows]
    return d

switchable_view gives you a viewer with one button per design and a toggle for all side chains.

[ ]:

from __future__ import annotations
import io, json, uuid, warnings
import biotite.structure.io.pdb as biopdb
from IPython.display import HTML

VIEWER_JS = "https://3Dmol.org/build/3Dmol-min.js"


def _to_pdb(sub):
    f = biopdb.PDBFile()
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*chararray class is deprecated.*")
        f.set_structure(sub)
    buf = io.StringIO()
    f.write(buf)
    return buf.getvalue()


def motif_model(cif_path):
    """Write the catalytic residues, whole, as one PDB model.

    Returns ``{"pdb": str, "tips": [[resi, atom_name], ...], "n_rest": int}``. Keeping the
    residues in one model lets bonds be perceived across the tip / non-tip boundary. Tip atoms
    are identified to the viewer by residue number and atom name, carried in the PDB itself.
    """
    aa = parse(str(cif_path), config=ANNOTATED_CIF)["assemblies"]["1"][0]
    ensure_annotations(aa, "is_ligand")
    tip = np.asarray(aa.mask("sa_mask('SA_CRDSRC') & ~hetero"))

    chain = np.asarray(aa.chain_iid).astype(str)
    keys = set(zip(chain[tip].tolist(), aa.res_id[tip].astype(int).tolist(), strict=False))
    same_res = np.array([(c, int(r)) in keys for c, r in zip(chain, aa.res_id, strict=False)], dtype=bool)
    motif = same_res & ~np.asarray(aa.mask("hetero"))

    tips = [[int(r), str(n)] for r, n in zip(aa.res_id[tip], aa.atom_name[tip], strict=False)]
    return {"pdb": _to_pdb(aa[motif]), "tips": tips, "n_rest": int((motif & ~tip).sum())}


def switchable_view(samples, width=820, height=560):
    """A viewer with one button per sample and a toggle for all side chains.

    ``samples`` are entries of ``{"label", "cif", "motif", "tips", "note"}``, where ``motif`` and
    ``tips`` come from :func:`motif_model`. Returns an ``IPython.display.HTML`` that survives
    being saved into the notebook.
    """
    uid = "v" + uuid.uuid4().hex[:10]

    payload = ",\n".join(
        "{{label:{l},cif:{c},motif:{m},tips:{t},note:{n}}}".format(
            l=json.dumps(s["label"]),
            c=json.dumps(s["cif"]),
            m=json.dumps(s["motif"]),
            t=json.dumps(s["tips"]),
            n=json.dumps(s.get("note", "")),
        )
        for s in samples
    )

    buttons = "".join(
        f'<button id="{uid}_b{i}" onclick="{uid}_show({i})" '
        f'style="margin:2px 4px 2px 0;padding:4px 10px;border:1px solid #bbb;'
        f'border-radius:4px;background:#f4f4f4;cursor:pointer;font:12px monospace">'
        f'{s["label"]}</button>'
        for i, s in enumerate(samples)
    )

    return HTML(f"""
<div style="font:13px/1.5 system-ui,sans-serif">
  <div style="margin-bottom:4px">{buttons}</div>
  <div style="margin-bottom:6px">
    <button id="{uid}_sc" onclick="{uid}_toggle()"
      style="padding:4px 10px;border:1px solid #bbb;border-radius:4px;background:#f4f4f4;
             cursor:pointer;font:12px monospace">show all side chains</button>
  </div>
  <div id="{uid}_note" style="margin:4px 0 8px;color:#333;font:12px monospace"></div>
  <div id="{uid}" style="width:{width}px;height:{height}px;position:relative"></div>
  <div style="margin-top:6px;color:#555;font-size:12px">
    light grey cartoon = scaffold &nbsp;|&nbsp; orange = substrate &nbsp;|&nbsp;
    green = conditioned tip atoms &nbsp;|&nbsp; grey sticks = rest of the catalytic residues
  </div>
</div>
<script src="{VIEWER_JS}"></script>
<script>
(function() {{
  var data = [
{payload}
  ];
  var cur = 0, showAll = false, viewer = null;

  function draw() {{
    var s = data[cur];
    viewer.clear();
    viewer.addModel(s.cif, "cif");                       // model 0: the whole design
    viewer.setStyle({{}}, {{cartoon: {{color: "lightgrey"}}}});
    if (showAll) {{                        // every side chain, thin sticks, elements coloured
      viewer.addStyle({{model: 0, hetflag: false}},
                      {{stick: {{colorscheme: "grayCarbon", radius: 0.08}}}});
    }}
    viewer.setStyle({{hetflag: true}}, {{stick: {{colorscheme: "orangeCarbon", radius: 0.30}}}});

    viewer.addModel(s.motif, "pdb");                     // model 1: catalytic residues, whole
    viewer.setStyle({{model: 1}}, {{stick: {{colorscheme: "grayCarbon", radius: 0.12}}}});
    // Tip atoms are addressed by residue number and atom name rather than by serial, which the
    // viewer assigns independently of the order they were written in.
    for (var t = 0; t < s.tips.length; t++) {{
      var sel = {{model: 1, resi: s.tips[t][0], atom: s.tips[t][1]}};
      viewer.addStyle(sel, {{stick: {{colorscheme: "greenCarbon", radius: 0.24}}}});
      viewer.addStyle(sel, {{sphere: {{scale: 0.32, colorscheme: "greenCarbon"}}}});
    }}

    viewer.zoomTo({{hetflag: true}}); viewer.zoom(0.45); viewer.render();
    document.getElementById("{uid}_note").textContent = s.note;
    for (var i = 0; i < data.length; i++) {{
      var b = document.getElementById("{uid}_b" + i);
      if (b) {{ b.style.background = (i === cur) ? "#cfe3cf" : "#f4f4f4";
                b.style.fontWeight = (i === cur) ? "bold" : "normal"; }}
    }}
    var sc = document.getElementById("{uid}_sc");
    sc.style.background = showAll ? "#cfe3cf" : "#f4f4f4";
    sc.textContent = showAll ? "hide other side chains" : "show all side chains";
  }}

  window["{uid}_show"] = function(i) {{ cur = i; draw(); }};
  window["{uid}_toggle"] = function() {{ showAll = !showAll; draw(); }};

  function build() {{
    var el = document.getElementById("{uid}");
    if (!el || typeof $3Dmol === "undefined") {{ setTimeout(build, 120); return; }}
    viewer = $3Dmol.createViewer(el, {{backgroundColor: "white"}});
    draw();
  }}
  build();
}})();
</script>
""")


def samples_from_metrics(df, n=6, sort_by=None, note_cols=()):
    """Take the first ``n`` designs from a metrics table and prepare them for the viewer."""
    d = df.sort_values(sort_by) if sort_by and sort_by in df else df
    out = []
    for i, (_, row) in enumerate(d.head(n).iterrows()):
        path = row["designed_structure_path"]
        mm = motif_model(path)
        with open(path) as f:
            cif = f.read()
        bits = [
            f"{c.split('.')[0]}={row[c]:.3f}"
            for c in note_cols
            if c in row and isinstance(row[c], int | float) and row[c] == row[c]
        ]
        label = str(row.get("example_id", i)).split("__")[-1].replace("_job0", "")
        out.append({"label": label, "cif": cif, "motif": mm["pdb"], "tips": mm["tips"], "note": "  ".join(bits)})
    return out


def selection_gallery(cif_path, selections, width=820, height=560):
    """Interactive selection viewer: one button per atomworks selection query.

    Each button highlights exactly the atoms ``aa.mask("<query>")`` returns — the mask is evaluated
    here in Python and baked into the viewer as ``(chain, res_id, atom_name)`` triples, so what you
    see is what the selection selects, including coordinate/annotation predicates a viewer's own
    selection language could not express. ``selections`` is a list of ``(label, query)`` pairs.
    """
    aa = parse(str(cif_path), config=ANNOTATED_CIF)["assemblies"]["1"][0]
    ensure_annotations(aa, "is_ligand")
    chain = np.asarray(aa.chain_id).astype(str)
    resi = aa.res_id.astype(int)
    aname = np.asarray(aa.atom_name).astype(str)
    items = []
    for label, q in selections:
        m = np.asarray(aa.mask(q))
        atoms = [[str(chain[i]), int(resi[i]), str(aname[i])] for i in np.where(m)[0]]
        items.append({"label": label, "query": q, "n": int(m.sum()), "atoms": atoms})

    uid = "g" + uuid.uuid4().hex[:10]
    payload = ",\n".join(
        "{{label:{l},query:{q},n:{n},atoms:{a}}}".format(
            l=json.dumps(it["label"]), q=json.dumps(it["query"]),
            n=json.dumps(it["n"]), a=json.dumps(it["atoms"]))
        for it in items)
    buttons = "".join(
        f'<button id="{uid}_b{i}" onclick="{uid}_show({i})" '
        f'style="margin:2px 4px 2px 0;padding:4px 10px;border:1px solid #bbb;border-radius:4px;'
        f'background:#f4f4f4;cursor:pointer;font:12px monospace">{it["label"]}</button>'
        for i, it in enumerate(items))

    return HTML(f"""
<div style="font:13px/1.5 system-ui,sans-serif">
  <div style="margin-bottom:4px">{buttons}</div>
  <div id="{uid}_note" style="margin:4px 0 8px;color:#333;font:12px monospace">&nbsp;</div>
  <div id="{uid}" style="width:{width}px;height:{height}px;position:relative"></div>
  <div style="margin-top:6px;color:#555;font-size:12px">
    grey cartoon = design &nbsp;|&nbsp; green = atoms returned by <code>aa.mask(query)</code>
  </div>
</div>
<script src="{VIEWER_JS}"></script>
<script>
(function() {{
  var data = [
{payload}
  ];
  var pdb = {json.dumps(_to_pdb(aa))};
  var cur = 0, viewer = null;
  function draw() {{
    var s = data[cur];
    viewer.clear();
    viewer.addModel(pdb, "pdb");
    viewer.setStyle({{}}, {{cartoon: {{color: "lightgrey"}}}});
    viewer.setStyle({{hetflag: true}}, {{stick: {{colorscheme: "orangeCarbon", radius: 0.22}}}});
    var ors = [];
    for (var k = 0; k < s.atoms.length; k++)
      ors.push({{chain: s.atoms[k][0], resi: s.atoms[k][1], atom: s.atoms[k][2]}});
    if (ors.length) {{
      viewer.addStyle({{or: ors}}, {{sphere: {{scale: 0.34, colorscheme: "greenCarbon"}}}});
      viewer.addStyle({{or: ors}}, {{stick: {{colorscheme: "greenCarbon", radius: 0.20}}}});
    }}
    viewer.zoomTo(); viewer.render();
    document.getElementById("{uid}_note").innerHTML =
      "aa.mask(<b>" + s.query + "</b>) &rarr; " + s.n + " atoms";
    for (var i = 0; i < data.length; i++) {{
      var b = document.getElementById("{uid}_b" + i);
      if (b) {{ b.style.background = (i === cur) ? "#cfe3cf" : "#f4f4f4";
                b.style.fontWeight = (i === cur) ? "bold" : "normal"; }}
    }}
  }}
  window["{uid}_show"] = function(i) {{ cur = i; draw(); }};
  function build() {{
    var el = document.getElementById("{uid}");
    if (!el || typeof $3Dmol === "undefined") {{ setTimeout(build, 120); return; }}
    viewer = $3Dmol.createViewer(el, {{backgroundColor: "white"}});
    draw();
  }}
  build();
}})();
</script>
""")

# --- one selection-driven RMSD used in section 4 to reproduce the metrics ------------------


def sc_rmsd(gt_coords, pred_coords, sel, group=None):
    """RMSD after superposing ``gt_coords`` onto ``pred_coords`` using the atoms in the boolean
    ``sel`` mask (independently within each ``group`` when given), atom-weighted across groups.
    ``sel`` is what ``aa.mask("<query>")`` returns, so the selection drives everything."""
    from biotite.structure import superimpose, rmsd
    ok = ~np.isnan(gt_coords).any(1) & ~np.isnan(pred_coords).any(1)
    groups = np.zeros(len(sel), dtype=int) if group is None else group
    w = t = 0.0
    for g in np.unique(groups):
        a = (groups == g) & sel & ok
        n = int(a.sum())
        if n == 0:
            continue
        fitted, _ = superimpose(gt_coords[a], pred_coords[a])
        w += float(rmsd(gt_coords[a], fitted)) ** 2 * n
        t += n
    return float((w / t) ** 0.5)

1. Inspect the input#

Before you design anything, read exactly what the input has locked down and what it has left open.

The theozyme fixes the catalytic geometry and leaves the rest to the model. Each catalytic residue keeps only the atoms the benchmark names, the ligand is pinned in part, the catalytic residues are unindexed so the model chooses where they sit in the chain, and an expandable segment stands in for the scaffold still to be designed.

New to this? What do “unindexed” and “expandable segment” mean?

Unindexed means the model isn’t told “this catalytic piece must sit at sequence position 39.” It’s told “these atoms must keep their exact 3D arrangement relative to each other and the substrate, but you (the model) decide where in the new protein’s sequence each one falls.” The geometry is fixed; the position in the chain is not.

Expandable segment is a placeholder for “there’s more protein needed here, but not a fixed amount yet.” Rather than committing to an exact length upfront, it tells the model “fill this gap with however much new scaffold it takes to physically connect and support the theozyme pieces.”

The tutorial ships one worked example, `M0732_1xs1 <../guides/tutorials/enzyme_design/M0732_1xs1_theozyme.cif>`__. It is an entry from M-CSA (Mechanism and Catalytic Site Atlas); a free, expert-curated database that records exactly which residues, atoms, and cofactors form the working chemistry of real enzymes, including step-by-step reaction mechanisms for many of them. This example’s theozyme carries the catalytic residues and the ligand with the benchmark’s own conditions already applied. Point THEOZYME at your own annotated CIF to run this workflow on a different constellation. The Python API tutorial covers building an annotated CIF file using the Python API.

[ ]:

# ── Settings — input ─────────────────────────────────────────────
THEOZYME = TUTORIAL_DIR / "M0732_1xs1_theozyme.cif"   # or your own annotated CIF
# ─────────────────────────────────────────────────────────────────

from pathlib import Path
theo = Path(THEOZYME)
if not theo.exists():
    raise FileNotFoundError(f"{theo} does not exist.")

print("theozyme:", theo)

Read back what the input fixes: coordinates through C_CRD, identity through C_SEQ, and the size of the region left to design.

[ ]:

import warnings
with warnings.catch_warnings():
    warnings.filterwarnings("ignore", message=".*Optimal rotation is not uniquely or poorly defined.*")
    aa = describe_conditions(theo, "theozyme")
motif = aa.mask("sa_mask('C_CRD') & ~hetero")

New to this? What does this printout mean?

describe_conditions summarizes exactly what the theozyme file locks down, atom by atom.

``atoms 224 | protein 196 ligand 28`` — the whole input has 224 atoms total: 196 belong to the 5 catalytic protein fragments, 28 belong to the ligand (substrate).

``C_CRD coordinates fixed: 28 (protein 16 + ligand 12)`` — 28 atoms have their exact 3D position pinned in space. Only 16 of the protein’s 196 atoms need this (the actual catalytic tips, e.g. arginine’s NH1/NH2/NE/CZ); only 12 of the ligand’s 28 atoms do too — meaning part of the substrate’s pose is left loose, not rigidly fixed.

``C_SEQ identity fixed: 44 (protein 16 + ligand 28)`` — 44 atoms have their chemical identity locked (what atom/element it must be). All 28 ligand atoms are here, since the substrate is one specific molecule — its identity is never in question, even where its position is loose. Only the same 16 protein atoms as above have both position and identity fixed together.

``designable residues: 180`` — the finished protein must total 180 residues; the 5 catalytic fragments (below) get placed somewhere within that, and everything else is invented from scratch.

The 5 ``res`` lines — each is one catalytic fragment: e.g. res 182 GLU ['CD','CG','OE1','OE2'] means a glutamate at placeholder position 182, with only its carboxylate group pinned. Two lines (180 ALA, 181 ARG) fix only generic backbone atoms (C,O) — meaning the real requirement there is just “a backbone carbonyl here,” not a specific side chain.

The catalytic constellation#

Green marks the coordinate conditioned motif atoms and orange marks the substrate. There is no scaffold yet. This is everything the model gets.

[ ]:

import py3Dmol, biotite.structure.io.pdb as biopdb, io as _io, warnings

def atomarray_to_pdb(sub):
    f = biopdb.PDBFile()
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*chararray class is deprecated.*")
        f.set_structure(sub)
    buf = _io.StringIO(); f.write(buf); return buf.getvalue()

show = (motif | aa.mask("hetero")) & aa.mask("~has_nan_coord()")   # skip the placeholder
v = py3Dmol.view(width=720, height=520)
v.addModel(atomarray_to_pdb(aa[show]), "pdb")
v.setStyle({"stick": {"colorscheme": "greenCarbon", "radius": 0.18}})
v.setStyle({"hetflag": True}, {"stick": {"colorscheme": "orangeCarbon", "radius": 0.22}})
v.addStyle({"not": {"hetflag": True}}, {"sphere": {"scale": 0.25}})
v.setBackgroundColor("0xffffff"); v.zoomTo()
print("green: conditioned motif atoms   orange: substrate")
v.show()

2. Scaffold design#

Sampling starts from pure noise. Backbone and sequence get generated together, constrained throughout by the motif geometry. This is the de novo step. It returned the lowest pass rate of the three sampling steps in the runs recorded here, on one catalytic site under one checkpoint.

Settings for scaffold generation#

Setting	Config key	Default	Effect
`NUM_REPLICATES`	`num_replicates`	1	Scaffolds to sample. Only a fraction pass, so sample more than you need
`BATCH_SIZE`	`batch_size`	1	Designs denoised together on one GPU. A step costs nearly the same whatever the batch holds, so keep batches full
`NSTEPS`	`nsteps`	60	Model evaluations
`GUIDANCE_W`	`guidance_w`	1.0	1.0 turns guidance off. Any other value enables classifier free guidance, which evaluates the network a second time per step and roughly doubles sampling time
`BB_CA_TSR_K`	`bb_ca.tsr_k`	16	TSR k for the backbone; higher values lower temperature
`LOCAL_LATENTS_TSR_K`	`local_latents.tsr_k`	4	TSR k for the sequence; higher values lower temperature

None leaves the configuration default in place. Every setting here except NUM_REPLICATES and BATCH_SIZE also applies to steps 5 and 6, where it appears again under a step-specific name.

[ ]:

# ── Settings ─────────────────────────────────────────────────────
NUM_REPLICATES              = 25    # num_replicates: scaffolds to sample
BATCH_SIZE                  = 5     # batch_size: designs denoised together on one GPU
NSTEPS                      = None  # generation.args.nsteps; None keeps the default, 60
GUIDANCE_W                  = None  # generation.args.guidance_w; None keeps the default, 1.0
BB_CA_TSR_K        = None  # bb_ca tsr_k; None keeps the default, 16
LOCAL_LATENTS_TSR_K = None # local_latents tsr_k; None keeps the default, 4
# ─────────────────────────────────────────────────────────────────

Now submit the batch and see how many designs land the motif accurately enough to pass.

[ ]:

ov2, name2 = [], f"n{NUM_REPLICATES}"
name2 = override(ov2, name2, "generation.args.nsteps", NSTEPS, "s")
name2 = override(ov2, name2, "generation.args.guidance_w", GUIDANCE_W, "g")
name2 = override(ov2, name2, BB_CA_TSR_K_KEY, BB_CA_TSR_K, "bb")
name2 = override(ov2, name2, LOCAL_LATENTS_TSR_K_KEY, LOCAL_LATENTS_TSR_K, "ll")

OUT2 = str(OUT_ROOT / f"scaffold_{name2}")
csv2 = launch(OUT2, "s2design", theo, NUM_REPLICATES, ov2, BATCH_SIZE)

Results#

One row per design. design_unindexed_motif_allatom_rmsd measures how accurately the motif got placed, and it runs without the refolding oracle, so it is deterministic. The refolded columns come from AF3 and carry the pass criteria.

New to this? How do I actually use this table?

Each row is one of your sampled designs. Look at the ``PASS /3`` column first — 3 means it cleared every AME gate (green across the row), fewer means it failed at least one. You don’t need to read every number: scan for the green rows, then use the viewer below to actually look at those designs in 3D. If nothing passed, the model didn’t find a working scaffold for this catalytic site in this batch — sampling more replicates, or adjusting a setting above, are the usual next moves.

[ ]:

styled2, df2 = ame_matrix(csv2)
styled2

3. Select the best design and inspect its saved conditions#

Generated CIFs carry the input’s saved conditioning targets forward. When the unindexed motif is inserted into the scaffold, its conditions are transferred onto the matched atoms. Coordinate targets retain the input geometry, which can differ from the generated coordinates; sequence targets retain the specified identities. The inserted residues now have definite scaffold positions, so they are indexed for subsequent rounds.

pick_best_design selects the best passing scaffold and copies its CIF unchanged. There is no need to re-apply conditions before inverse folding or partial diffusion. Rebuilding targets from generated coordinates would replace the original catalytic geometry with the model’s prediction.

The SA_CRDSRC annotation records which coordinate-conditioned input residue each matched atom came from. Use sa_mask('SA_CRDSRC') to select the motif atoms for inspection; the saved conditions themselves are available through C_CRD, C_SEQ, and C_CRDGRP (with C_INTCRD for internal targets).

[ ]:

best_cif = pick_best_design(OUT2, verbose=True)

Verify before chaining#

Read the selected CIF with the same function you used on the theozyme in section 1. For this absolute-coordinate theozyme, the coordinate and sequence conditions should still cover the catalytic atoms and ligand. The scaffold is now indexed, so the index-conditioned count changes. These counts check coverage; the saved target annotations, rather than the generated coordinates, specify the geometry and identities used in the next two steps.

[ ]:

_ = describe_conditions(best_cif, "saved conditions for refinement")

Inspect the designs#

Use the buttons to switch between designs. Catalytic residues are drawn in full, with the coordinate conditioned tip atoms in green and the rest of each residue in grey, which shows at a glance how much of the side chain the conditioning actually pins. Substrate in orange, scaffold as a grey cartoon.

[ ]:

switchable_view(samples_from_metrics(
    df2, n=6, sort_by="design_unindexed_motif_allatom_rmsd.rmsd",
    note_cols=("design_unindexed_motif_allatom_rmsd.rmsd",
               "refolded_motif_allatom_rmsd.rmsd", "oracle.overall_plddt")))

4. Query the design with the selection language#

Every output CIF carries the model’s per atom annotations, so you can pull out any region with the atomworks selection syntax. aa.mask("<query>") returns a boolean mask over the atoms. Pick a query with the buttons below to see exactly what it matches on the chosen design. The mask is evaluated in Python and baked into the viewer, so what you see is what the selection selects. That includes annotation and coordinate predicates like sa_mask('SA_CRDSRC') and z > …, which a viewer’s own selection language could not express. Drag to rotate, scroll to zoom.

For more on working with these files directly, see the atomworks guides on annotating and saving structures and loading and visualizing structures.

[ ]:

from atomworks.ml.conditions.annotator import ensure_annotations

# The design chosen in section 3 (top-ranked by ame.pass, then motif RMSD) -- same rule as
# pick_best_design -- loaded from its run path, whether it was just sampled or reused.
chosen = (pd.read_csv(csv2)
          .sort_values(["ame.pass", "design_unindexed_motif_allatom_rmsd.rmsd"], ascending=[False, True])
          .iloc[0])
chosen_cif = chosen["designed_structure_path"]
_aa = parse(chosen_cif, config=ANNOTATED_CIF)["assemblies"]["1"][0]
_zmid = float(np.median(_aa.coord[:, 2]))

selection_gallery(chosen_cif, [
    ("motif tip atoms",     "sa_mask('SA_CRDSRC') & ~hetero"),
    ("substrate",           "hetero"),
    ("backbone C\u03b1",   "atom_name == 'CA'"),
    ("aromatics",           "res_name in ['PHE','TYR','TRP','HIS']"),
    (f"z > {_zmid:.0f}",    f"z > {_zmid:.0f}"),
    ("sulfurs",             "element == 'S'"),
])

The gallery highlights a selection visually. The same query language also logs which design residue mapped onto which input motif position, which is what the next cells do.

Reproduce the metrics with the selection language#

Want to select the motif atoms yourself, to recompute a metric or to do your own analysis? Here is a showcase. The same aa.mask(...) selections rebuild the pipeline’s own scores. Below, the motif all atom RMSD (refolded_motif_allatom_rmsd) and the backbone RMSD (refolded_backbone_rmsd) for the chosen design get recomputed from the CIFs by superposition, and they land on the values already in the metrics table.

[ ]:

design = parse(chosen["designed_structure_path"], config=ANNOTATED_CIF)["assemblies"]["1"][0]
refold = parse(chosen["oracle.refolded_structure_path"], config=ANNOTATED_CIF)["assemblies"]["1"][0]
for _s in (design, refold):
    ensure_annotations(_s, "is_protein_backbone")

Before the RMSDs, read the conditioning back off the loaded design, meaning which design residue mapped onto which input motif position. The motif selection and the per residue filter are both aa.mask(...). The source label, A1 through A5, one per theozyme catalytic residue, is a data column you read with get_annotation.

[ ]:

from rfproteina.utils.standard_annotations import SA_CRDSRC

# The motif selection uses aa.mask(...); the input-motif position is a per-atom data column.
motif = design.mask("sa_mask('SA_CRDSRC') & ~hetero")
src   = np.asarray(SA_CRDSRC.annotation(design))
resids = sorted(set(design.res_id[motif].astype(int).tolist()))
print(f"catalytic motif: {int(motif.sum())} atoms over {len(resids)} residues")
for rid in resids:
    m = motif & (design.res_id == rid)
    print(f"  design res {rid:>4} {design.res_name[m][0]:<4}  <-  input motif {src[m][0]}"
          f"   {sorted(design.atom_name[m].tolist())}")
print(f"  substrate: {int(design.mask('hetero').sum())} ligand atoms")

This recomputes backbone RMSD by structurally overlaying the design and its refold (using only backbone atoms) and measuring the leftover distance. RMSD (root-mean-square deviation) is the standard way to measure how far apart, on average, two versions of the same structure are once optimally aligned. The manually reproduced value should exactly match the pipeline’s own reported number.

[ ]:

backbone = sc_rmsd(design.coord, refold.coord,
                   design.mask("is_protein_backbone"))
print(f"backbone RMSD         reproduced {backbone:.4f}   reported {chosen['refolded_backbone_rmsd.rmsd']:.4f}")

This does the same thing, but restricted to just the catalytic motif atoms, using a slightly more careful alignment method — sc_rmsd with per-atom groups, i.e. a grouped Kabsch superposition, the standard algorithm for best-fit overlay of two point sets. Again, the manual result should match the reported one exactly, confirming the catalytic geometry in the refolded structure stayed close to what was intended.

[ ]:

from rfproteina.conditions import C_CRDGRP
from rfproteina.utils.atom_permutation import permute_gt_atoms
from rfproteina.utils.standard_annotations import SA_CWI

sel    = design.mask("sa_mask('SA_CRDSRC') & is_protein")
groups = C_CRDGRP.annotation(design)
target = design.copy()
target.coord = SA_CWI.annotation(design).copy()  # input-motif geometry

# Match equivalent atoms (e.g. GLU OE1/OE2) in each group's own alignment frame.
for g in np.unique(groups[sel]):
    group_mask = groups == g
    group_sel = sel[group_mask]
    resolved = permute_gt_atoms(
        refold[group_mask], target[group_mask],
        atom_classes=np.where(group_sel, 0, -1), frame_mask=group_sel,
    )
    target.coord[group_mask] = resolved.coord[0]

motif = sc_rmsd(target.coord, refold.coord, sel, group=groups)
print(f"motif all-atom RMSD   reproduced {motif:.4f}   reported {chosen['refolded_motif_allatom_rmsd.rmsd']:.4f}")

5. Inverse folding#

Step 2 solved the hard part: RFD4-Proteina created a physically valid shape holding the catalytic geometry. That’s worth building on rather than re-risking, so this step explores cheaply around it. Here, the backbone gets held fixed and resampled to the input at every step, so only the sequence and side chains are drawn afresh, letting you try several sequences on one proven shape instead of gambling on a new one each time. (Mechanically, this is a special case of the partial diffusion technique covered in step 6, with backbone noise switched off entirely; its parameters can’t be combined with partial diffusion’s.)

Before working through it, here’s how steps 5 and 6 relate to step 2, for reference:

Step	Starts from	Backbone	Sequence	What selects it
Scaffold design	noise	generated	generated	neither switch below
Inverse folding	the design from step 3	fixed	resampled	`generation.args.inverse_folding=true`
Partial diffusion	the design from step 3	perturbed	perturbed	`generation.args.partial_trajectory_fraction_steps`

The two switches are mutually exclusive, and setting both raises an error.

Settings for sequence sampling temperature#

Setting	Config key	Default	Effect
`IF_LOCAL_LATENTS_TSR_K`	`local_latents.tsr_k`	4	TSR k for the sequence; lower values raise temperature
`IF_NUM_REPLICATES`	`num_replicates`	1	Sequences to sample
`IF_BB_CA_TSR_K`	`bb_ca.tsr_k`	16	TSR k for the backbone; higher values lower temperature
`IF_NSTEPS`	`nsteps`	60	Model evaluations
`IF_GUIDANCE_W`	`guidance_w`	1.0	1.0 turns guidance off

The generation=inverse_folding preset selects this step, applied in the execution cell below. generation.args.inverse_folding=true does the same thing.

We lower the sequence TSR k from 4 to 2 to explore a higher sampling temperature. Compare the resulting diversity and refolding metrics before choosing a setting. TSR sigma is a secondary schedule control.

[ ]:

# ── Settings ─────────────────────────────────────────────────────
IF_LOCAL_LATENTS_TSR_K = 2.0   # local_latents tsr_k; default 4
IF_NUM_REPLICATES               = 10    # num_replicates: sequences to sample
IF_BB_CA_TSR_K         = None  # bb_ca tsr_k; None keeps the default, 16
IF_NSTEPS                       = None  # generation.args.nsteps; None keeps the default, 60
IF_GUIDANCE_W                   = None  # generation.args.guidance_w; None keeps the default, 1.0
# ─────────────────────────────────────────────────────────────────

[ ]:

ov3, name3 = ["generation=inverse_folding"], f"n{IF_NUM_REPLICATES}"
name3 = override(ov3, name3, LOCAL_LATENTS_TSR_K_KEY, IF_LOCAL_LATENTS_TSR_K, "ll")
name3 = override(ov3, name3, BB_CA_TSR_K_KEY, IF_BB_CA_TSR_K, "bb")
name3 = override(ov3, name3, "generation.args.nsteps", IF_NSTEPS, "s")
name3 = override(ov3, name3, "generation.args.guidance_w", IF_GUIDANCE_W, "g")

OUT3 = str(OUT_ROOT / f"inverse_folding_{name3}")
# Two tasks of five sequences each fill BATCH_SIZE exactly.
csv3 = launch(OUT3, "s3invfold", best_cif, IF_NUM_REPLICATES, ov3, BATCH_SIZE, njobs=2)

Results#

Two further columns report how far each design has moved from the parent scaffold. Here the “parent scaffold” is the one design chosen in step 3, which every sequence here was generated from. ca_rmsd_to_parent covers Cα coordinates and seq_ident_to_parent covers the sequence. Inverse folding holds the backbone, so the first stays at zero and the whole change lands in the second.

New to this? What am I checking for here?

You are looking for the same pass/fail table as step 2, but every design here started from the one scaffold you picked in step 3, so a higher pass rate is expected, not a coincidence (see the Summary at the end for why). The two new columns say how much each sequence variant actually differs from that parent: ca_rmsd_to_parent should sit near zero, since the backbone was never allowed to move here, while seq_ident_to_parent is the fraction of the sequence that stayed the same — lower means a more different sequence on the same shape.

[ ]:

styled3, df3 = ame_matrix(csv3, reference=best_cif)
styled3

Inspect the designs#

[ ]:

switchable_view(samples_from_metrics(
    df3, n=6, sort_by="refolded_motif_allatom_rmsd.rmsd",
    note_cols=("refolded_motif_allatom_rmsd.rmsd", "ca_rmsd", "seq_ident", "oracle.overall_plddt")))

6. Partial diffusion#

The design gets re-noised to a point partway along the trajectory, then diffused back. Unlike inverse folding, both modalities get resampled, so the backbone moves too. Under the default settings the sequence is redrawn more freely while the backbone only takes a small nudge, making this a local search around an existing solution rather than a pure redesign.

Use this instead of inverse folding when you suspect the shape itself, not just the sequence, is holding back results (e.g. if step 5’s fixed-backbone sequences keep failing to refold correctly.)

Settings for how much to change#

PARTIAL_FRACTION is a fraction in (0, 1]. It re-runs the final part of the diffusion trajectory. The chosen design gets corrupted to a partly noised state and resampled from there, so larger values re-run more of the trajectory and move the design further. At 1 the whole trajectory re-runs and the enzyme scaffold is resampled from scratch, exactly like a de novo design. Because the backbone (bb_ca) is sampled slightly ahead of the latent carrying the sequence and side chains (local_latents), the sequence tends to change more noticeably while the backbone shifts only a little. At the fraction of 0.5 this tutorial uses, the design as a whole moves a small distance from the input scaffold.

Setting	Config key	Default	Effect
`PARTIAL_FRACTION`	`partial_trajectory_fraction_steps`	`null`	Fraction of the trajectory to re-run, in (0, 1]. Larger values change the design more
`PD_NUM_REPLICATES`	`num_replicates`	1	Refined variants to generate
`PD_BB_CA_TSR_K`	`bb_ca.tsr_k`	16	TSR k for the backbone; higher values lower temperature
`PD_LOCAL_LATENTS_TSR_K`	`local_latents.tsr_k`	4	TSR k for the sequence; higher values lower temperature
`PD_NSTEPS`	`nsteps`	60	Model evaluations
`PD_GUIDANCE_W`	`guidance_w`	1.0	1.0 turns guidance off

Concretely, PARTIAL_FRACTION sets the number of steps to run, as round(f * nsteps), and the noise level gets read from the schedule at that index. It indexes steps, not noise. At PARTIAL_FRACTION=0.5 with nsteps=60, that’s round(0.5*60)=30 model evaluations re-run — half the trajectory, not literally “50% noise” since the schedule is nonlinear. The bb_ca and local_latents schedules are nonlinear, so a small fraction leaves few steps to denoise.

You can set each modality’s starting point directly instead, with generation.args.partial_t_backbone and partial_t_latent, both in (0, 1). That lets you redraw the backbone further than the sequence, or the other way round.

Careful, the direction flips: PARTIAL_FRACTION near 1 means more change, but partial_t_backbone/partial_t_latent near 0 means more change — opposite conventions for a very similar-sounding idea.

Give both, or give one and set generation.args.couple_partial_t=true to derive the other from the schedule. They are mutually exclusive with PARTIAL_FRACTION.

[ ]:

# ── Settings ─────────────────────────────────────────────────────
PARTIAL_FRACTION                = 0.5   # generation.args.partial_trajectory_fraction_steps, (0, 1]
PD_NUM_REPLICATES               = 10    # num_replicates: refined variants to generate
PD_BB_CA_TSR_K         = None  # bb_ca tsr_k; None keeps the default, 16
PD_LOCAL_LATENTS_TSR_K = None  # local_latents tsr_k; None keeps the default, 4
PD_NSTEPS                       = None  # generation.args.nsteps; None keeps the default, 60
PD_GUIDANCE_W                   = None  # generation.args.guidance_w; None keeps the default, 1.0
# ─────────────────────────────────────────────────────────────────

[ ]:

ov4 = [f"generation.args.partial_trajectory_fraction_steps={PARTIAL_FRACTION}"]
name4 = f"frac{tag(PARTIAL_FRACTION)}_n{PD_NUM_REPLICATES}"
name4 = override(ov4, name4, BB_CA_TSR_K_KEY, PD_BB_CA_TSR_K, "bb")
name4 = override(ov4, name4, LOCAL_LATENTS_TSR_K_KEY, PD_LOCAL_LATENTS_TSR_K, "ll")
name4 = override(ov4, name4, "generation.args.nsteps", PD_NSTEPS, "s")
name4 = override(ov4, name4, "generation.args.guidance_w", PD_GUIDANCE_W, "g")

OUT4 = str(OUT_ROOT / f"partial_diffusion_{name4}")
# Two tasks of five variants each fill BATCH_SIZE exactly.
csv4 = launch(OUT4, "s4partial", best_cif, PD_NUM_REPLICATES, ov4, BATCH_SIZE, njobs=2)

Results#

The same deviation columns apply. Both modalities get resampled here, so ca_rmsd_to_parent is clearly non-zero. The backbone has moved, and that is what separates this step from inverse folding.

New to this? What am I checking for here?

Read this the same way as step 5’s table. The difference to watch is ca_rmsd_to_parent: here it’s genuinely greater than zero, because partial diffusion actually re-samples the backbone a little rather than freezing it. A larger value means the shape moved further from the parent scaffold — compare it against PARTIAL_FRACTION above to see how that setting controls how much drift you get.

[ ]:

styled4, df4 = ame_matrix(csv4, reference=best_cif)
styled4

Inspect the designs#

[ ]:

switchable_view(samples_from_metrics(
    df4, n=6, sort_by="refolded_motif_allatom_rmsd.rmsd",
    note_cols=("refolded_motif_allatom_rmsd.rmsd", "ca_rmsd", "seq_ident", "oracle.overall_plddt")))

Summary#

You went from a theozyme to a set of scored, inspectable enzyme designs.

You performed three design steps and analyzed their outputs:

Step	Starts from	Backbone	Sequence	What selects it
Scaffold design	noise	generated	generated	neither switch below
Inverse folding	the design from step 3	fixed	resampled	`generation.args.inverse_folding=true`
Partial diffusion	the design from step 3	perturbed	perturbed	`generation.args.partial_trajectory_fraction_steps`

The two switches (inverse folding and parital diffusion) are mutually exclusive, and setting both raises an error.

Every other parameter applies to all three steps: num_replicates, batch_size, generation.args.nsteps, generation.args.guidance_w, and the TSR temperature controls. generation.model.bb_ca.simulation_step_params.tsr_k defaults to 16 for the backbone, and the same key under local_latents defaults to 4 for the sequence. This tutorial changes the sequence k in step 5. Use tsr_sigma as a secondary control over when sharpening takes effect.

Steps 5 and 6 are two alternative refinement strategies, not sequential steps — both start independently from the design prepared in step 3, so run whichever fits (sequence-only vs. sequence-and-shape), or both, and compare. Because neither builds on step 2 directly, their pass rates shouldn’t be compared with step 2’s either.

To change a parameter, edit the settings block and re-run from that section. Output directories carry the parameter values, so your earlier results survive.

That’s the full loop: theozyme in, scored and inspectable enzyme designs out. From here: scale up num_replicates for a fuller search, swap in your own active site via the annotated CIF tutorial, or take the passing designs forward for experimental testing.

Download this notebook