VHH design tutorial: from a target to a set of designed nanobodies#

In this tutorial you will use RFD4-Proteina to design VHHs (single-domain antibodies, or nanobodies) against a target, look at what the model fixed and what it designed, and view the resulting antibody–target complexes.

A VHH has a conserved immunoglobulin framework and three variable loops, the complementarity-determining regions CDR1, CDR2 and CDR3, that make most of the contacts with the target. The inputs in examples/antibody_design/structures/vhh fix a cropped target and the framework of an existing VHH and leave the CDRs to the model: the framework keeps its internal geometry but is free to move relative to the target, CDR lengths are sampled anew for every candidate, and a patch of nearby target residues is chosen as hotspots for every candidate. CATH conditioning supplies the immunoglobulin fold context.

You will:

  1. inspect one input and read back which atoms it fixes;

  2. generate a few designs for it;

  3. split each design into target, framework and CDRs and view them;

  4. refine one design with partial diffusion.

This takes a few minutes on one GPU. Scoring the designs with an AlphaFold 3 or OpenFold3 refold, and scaling up to a campaign, are described in the campaign steps at the end.

Setup#

The next cell finds the repository, sets the paths and defines the small helpers used below: a function that runs rfd4-generate and viewer helpers. The code is folded by default.

[ ]:

import os, re, sys, shutil, subprocess, warnings, logging
from pathlib import Path

warnings.filterwarnings("ignore")
logging.getLogger("atomworks").setLevel(logging.ERROR)

# Locate the repository by walking up to the `.project-root` marker; $PROJECT_ROOT is the fallback.
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():
    if not os.environ.get("PROJECT_ROOT"):
        raise RuntimeError("Run this notebook inside the RFD4-Proteina checkout, or set PROJECT_ROOT.")
    REPO_ROOT = Path(os.environ["PROJECT_ROOT"])
os.environ["PROJECT_ROOT"] = str(REPO_ROOT)
os.environ.setdefault("CLUSTER", "local")
EXAMPLE_DIR = REPO_ROOT / "examples" / "antibody_design"
INPUT_DIR = EXAMPLE_DIR / "structures" / "vhh"
OUT_ROOT = Path(os.environ.get("ANTIBODY_TUTORIAL_RUNS", REPO_ROOT / "runs" / "antibody_design_tutorial"))
OUT_ROOT.mkdir(parents=True, exist_ok=True)

# Extra Hydra overrides applied to every run, for example "hf_ckpt_path=null" to use a native checkpoint.
EXTRA_OVERRIDES = os.environ.get("ANTIBODY_TUTORIAL_OVERRIDES", "").split()

# Subprocesses do not inherit sys.path: let them import the checkout's sources.
os.environ["PYTHONPATH"] = f"{REPO_ROOT / 'src'}:" + os.environ.get("PYTHONPATH", "")

import numpy as np, pandas as pd, torch
print("GPU:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none: generation needs a GPU")


def generate(config, out_dir, *overrides):
    """Run `generate.py` with one of the example configs and return the generated CIFs."""
    cmd = [sys.executable, str(REPO_ROOT / "src/rfproteina/generate.py"),
           "--config-path", str(EXAMPLE_DIR), "--config-name", config,
           f"out_dir={out_dir}", "~metrics", "compile_lengths=null", *overrides, *EXTRA_OVERRIDES]
    proc = subprocess.run(cmd, cwd=REPO_ROOT, capture_output=True, text=True)
    if proc.returncode != 0:
        print(proc.stdout[-1500:], proc.stderr[-3000:])
        raise RuntimeError("generation failed")
    return sorted(Path(out_dir).glob("inference_samples/*/*_generated.cif"))


from atomworks.io import parse
from rfproteina.utils.parse_presets import ANNOTATED_CIF
import rfproteina.conditions  # registers the condition names used in selections, such as C_HOT


def load(path):
    """Load an annotated CIF (also .cif.zst) as the first model of its first assembly."""
    return parse(path, config=ANNOTATED_CIF)["assemblies"]["1"][0]


def segments(res_ids):
    """Group sorted residue ids into runs of consecutive numbers: [(first, last), ...]."""
    res_ids = sorted(set(int(r) for r in res_ids))
    runs, start = [], None
    for prev, cur in zip([None] + res_ids[:-1], res_ids):
        if prev is None or cur != prev + 1:
            if start is not None:
                runs.append((start, prev))
            start = cur
    if start is not None:
        runs.append((start, res_ids[-1]))
    return runs


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

def to_pdb(atoms):
    atoms = atoms.copy()
    atoms.set_annotation("b_factor", np.zeros(atoms.array_length()))   # model B-factors can overflow the PDB columns
    atoms.set_annotation("occupancy", np.ones(atoms.array_length()))
    f = biopdb.PDBFile(); f.set_structure(atoms)
    buf = io.StringIO(); f.write(buf); return buf.getvalue()

TARGET, FRAMEWORK, CDR, HOTSPOT = "0xBBBBBB", "0x2A9D8F", "0xE76F51", "0x7B2D8E"

def style_complex(view, aa, model=-1):
    """Target in gray, VHH framework in teal, CDRs in orange sticks, hotspot atoms as purple spheres."""
    vhh = aa.chain_id == "H"
    cdr = vhh & ~aa.mask("sa_mask('C_SEQ')")
    hot = aa.mask("sa_mask('C_HOT')") & ~vhh
    sel = lambda m: {"serial": [int(i) + 1 for i in np.flatnonzero(m)]}
    view.setStyle({"model": model}, {})
    view.addStyle({"model": model, **sel(~vhh)}, {"cartoon": {"color": TARGET, "opacity": 0.9}})
    view.addStyle({"model": model, **sel(vhh & ~cdr)}, {"cartoon": {"color": FRAMEWORK}})
    view.addStyle({"model": model, **sel(cdr)}, {"cartoon": {"color": CDR, "thickness": 0.6}})
    view.addStyle({"model": model, **sel(cdr)}, {"stick": {"color": CDR, "radius": 0.12}})
    view.addStyle({"model": model, **sel(hot)}, {"sphere": {"color": HOTSPOT, "scale": 0.6}})

def view_complex(aa, width=720, height=520):
    v = py3Dmol.view(width=width, height=height)
    v.addModel(to_pdb(aa), "pdb")
    style_complex(v, aa)
    v.setBackgroundColor("0xffffff"); v.zoomTo()
    return v

1. Inspect the input#

The tutorial ships eight inputs: two cropped targets, PMVK (PDB 3CH4) and MZB1 (PDB 7AAH), each paired with four VHH frameworks. The file name reads <target>__<framework>_cdr123. Pick one:

[ ]:

# ── Settings — input ─────────────────────────────────────────────
INPUT_NAME = "pmvk_3CH4__VHH_ozoralizumab_cdr123"   # any file name in the folder, without ".cif.zst"
# ─────────────────────────────────────────────────────────────────

print("\n".join(sorted(p.name.removesuffix(".cif.zst") for p in INPUT_DIR.glob("*.cif.zst"))))
input_path = INPUT_DIR / f"{INPUT_NAME}.cif.zst"
aa_in = load(input_path)

An annotated CIF carries the conditions as per-atom and per-residue annotations (see Preparing inputs). Read back what this input fixes: chain A is the target and chain H the VHH framework.

[ ]:

rows = []
for chain in np.unique(aa_in.chain_id):
    m = aa_in.chain_id == chain
    rows.append({
        "chain": chain,
        "residues": len(set(aa_in.res_id[m])),
        "atoms": int(m.sum()),
        "coordinates fixed (C_CRD)": int((m & aa_in.mask("sa_mask('C_CRD')")).sum()),
        "sequence fixed (C_SEQ)": int((m & aa_in.mask("sa_mask('C_SEQ')")).sum()),
    })
pd.DataFrame(rows).set_index("chain")

The framework residues keep their identity and their internal geometry; the CDRs are not in the input at all. They are expandable segments: stretches of the chain whose length the model chooses within the bounds recorded in the input. Each candidate draws new lengths, so every design in the next section has different CDR lengths.

[ ]:

v = py3Dmol.view(width=720, height=420)
v.addModel(to_pdb(aa_in[~aa_in.mask("has_nan_coord()")]), "pdb")
v.setStyle({"chain": "A"}, {"cartoon": {"color": TARGET}})
v.setStyle({"chain": "H"}, {"cartoon": {"color": FRAMEWORK}})
v.setBackgroundColor("0xffffff"); v.zoomTo()
print("gray: target   teal: VHH framework (CDRs are still to be designed)")
v.show()

2. Generate designs#

vhh.yaml composes the base inference config with the VHH settings: a high, front-loaded recycling schedule and a sampling transform that picks two or three nearby hotspot residues on the target for every candidate. We restrict the run to the input chosen above and draw a few candidates.

Setting

Config key

Effect

NUM_REPLICATES

num_replicates

Candidates per input. The pilot in the example README uses 1; 100 or more is a production run.

BATCH_SIZE

batch_size

Candidates sampled together on the GPU.

SEED

seed

Makes the run repeatable.

 
[ ]:

# ── Settings ─────────────────────────────────────────────────────
NUM_REPLICATES = 4
BATCH_SIZE = 4
SEED = 42
# ─────────────────────────────────────────────────────────────────

OUT_DESIGN = OUT_ROOT / f"vhh_{INPUT_NAME}_seed{SEED}"
designs = generate("vhh", OUT_DESIGN, f"subset_to_examples={INPUT_NAME}",
                   f"num_replicates={NUM_REPLICATES}", f"batch_size={BATCH_SIZE}", f"seed={SEED}")
print(f"{len(designs)} designs in {OUT_DESIGN / 'inference_samples'}")
for p in designs: print(" ", p.relative_to(OUT_DESIGN))

3. Inspect the designs#

A generated CIF has the same annotations as the input, now for the completed complex. Two selections are enough to split a design into its parts:

  • the CDRs are the VHH atoms that are not sequence conditioned, chain_id == "H" and ~sa_mask('C_SEQ');

  • the hotspots are the target atoms with the hotspot condition, sa_mask('C_HOT').

The table lists, for every design, the three CDR lengths (a gap in the residue numbering separates one CDR from the next) and how many target residues were chosen as hotspots.

[ ]:

def describe(path):
    aa = load(path)
    vhh = aa.chain_id == "H"
    cdr_res = aa.res_id[vhh & ~aa.mask("sa_mask('C_SEQ')")]
    cdrs = segments(cdr_res)
    return aa, {
        "VHH residues": len(set(aa.res_id[vhh])),
        "CDR lengths": " / ".join(str(b - a + 1) for a, b in cdrs),
        "hotspot residues": len(set(aa.res_id[aa.mask("sa_mask('C_HOT')") & ~vhh])),
    }

loaded, table = [], []
for i, path in enumerate(designs):
    aa, row = describe(path); loaded.append(aa); table.append({"#": i, **row})
pd.DataFrame(table).set_index("#")

Each design has three stretches the model was free to design, with different lengths in every candidate. Target in gray, framework in teal, CDRs in orange and the hotspot atoms the model was asked to contact as purple spheres. The first design:

[ ]:

view_complex(loaded[0]).show()

All designs side by side, on the same target:

[ ]:

n = len(loaded)
cols = 2
grid = py3Dmol.view(width=900, height=330 * ((n + 1) // cols), viewergrid=((n + 1) // cols, cols))
for i, aa in enumerate(loaded):
    r, c = divmod(i, cols)
    grid.addModel(to_pdb(aa), "pdb", viewer=(r, c))
    grid.setStyle({"model": -1}, {}, viewer=(r, c))
    vhh = aa.chain_id == "H"; cdr = vhh & ~aa.mask("sa_mask('C_SEQ')")
    ser = lambda m: {"serial": [int(j) + 1 for j in np.flatnonzero(m)]}
    grid.setStyle({"model": -1, **ser(~vhh)}, {"cartoon": {"color": TARGET}}, viewer=(r, c))
    grid.addStyle({"model": -1, **ser(vhh & ~cdr)}, {"cartoon": {"color": FRAMEWORK}}, viewer=(r, c))
    grid.addStyle({"model": -1, **ser(cdr)}, {"cartoon": {"color": CDR}, "stick": {"color": CDR, "radius": 0.12}}, viewer=(r, c))
    grid.setBackgroundColor("0xffffff", viewer=(r, c)); grid.zoomTo(viewer=(r, c))
grid.show()

4. Refine a design with partial diffusion#

Sampling more candidates looks for distinct solutions. Once a design is promising, partial diffusion explores its neighborhood instead: it re-noises the design partway along the trajectory and denoises it again, so children stay close to the parent while the CDR loops and the contacts they make can still change.

partial.yaml keeps the parent’s sampled hotspots, target, framework conditioning and realized CDR lengths, and re-runs the last 42% of the schedule (partial_trajectory_fraction_steps). Use the generated annotated complex as the input, never an oracle-refolded structure.

[ ]:

# ── Settings ─────────────────────────────────────────────────────
PARENT = 0              # index of the design from section 2 to refine
NUM_CHILDREN = 4
# ─────────────────────────────────────────────────────────────────

parent_path = designs[PARENT]
OUT_REFINE = OUT_ROOT / f"refine_{INPUT_NAME}_parent{PARENT}"
children = generate("partial", OUT_REFINE, f"inputs={parent_path}",
                    f"num_replicates={NUM_CHILDREN}", f"batch_size={NUM_CHILDREN}", f"seed={SEED}")
print(f"{len(children)} files")
for p in children: print(" ", p.relative_to(OUT_REFINE))

Superpose each child on its parent over the target and compare the CDR loops: gray target, the parent’s CDRs in orange and the child’s in teal.

[ ]:

from biotite.structure import superimpose, rmsd

parent = loaded[PARENT]
child_aa = [load(p) for p in children]
tgt = lambda aa: aa.chain_id == "A"
ca = lambda aa: aa.atom_name == "CA"
ref = parent[tgt(parent) & ca(parent)]

rows, fitted = [], []
for i, aa in enumerate(child_aa):
    c_ca = aa[tgt(aa) & ca(aa)]
    n = min(len(ref), len(c_ca))
    _, transform = superimpose(ref[:n], c_ca[:n])
    moved = transform.apply(aa); fitted.append(moved)
    vh_p, vh_c = parent[(parent.chain_id == "H") & ca(parent)], moved[(moved.chain_id == "H") & ca(moved)]
    m = min(len(vh_p), len(vh_c))
    rows.append({"child": i, "VHH Cα RMSD to parent (Å)": round(float(rmsd(vh_p[:m], vh_c[:m])), 2)})
display(pd.DataFrame(rows).set_index("child"))   # how far each child moved from the parent

v = py3Dmol.view(width=720, height=520)
v.addModel(to_pdb(parent), "pdb")
v.setStyle({"model": 0}, {})
vhh = parent.chain_id == "H"; cdr_p = vhh & ~parent.mask("sa_mask('C_SEQ')")
ser = lambda m: {"serial": [int(j) + 1 for j in np.flatnonzero(m)]}
v.setStyle({"model": 0, **ser(~vhh)}, {"cartoon": {"color": TARGET}})
v.addStyle({"model": 0, **ser(vhh)}, {"cartoon": {"color": FRAMEWORK, "opacity": 0.5}})
v.addStyle({"model": 0, **ser(cdr_p)}, {"stick": {"color": CDR, "radius": 0.14}})
for k, aa in enumerate(fitted[:2], start=1):
    v.addModel(to_pdb(aa[aa.chain_id == "H"]), "pdb")
    v.setStyle({"model": k}, {})
    hv = aa[aa.chain_id == "H"]; cd = ~hv.mask("sa_mask('C_SEQ')")
    v.addStyle({"model": k, **ser(cd)}, {"stick": {"color": "0x1D6F8A", "radius": 0.1}})
v.setBackgroundColor("0xffffff"); v.zoomTo()
print("orange: parent CDRs   blue: CDRs of two children")
v.show()

Next steps#

You designed VHHs for one input, split the designs into target, framework and CDRs, and refined one of them. A design campaign repeats this at scale:

  1. Discover productive combinations with best-of-N. Run vhh.yaml on the whole input folder with a larger num_replicates (the count is per input: 100 across eight inputs is 800 attempts) and score every design with the antibody metrics and an AlphaFold 3 or OpenFold3 refold. See Evaluate with metrics and Refolding; setup is in the installation guide.

  2. Compare within each target. Which frameworks, CDR lengths and hotspot patches yield passers? Keep failed evaluations in the attempt count, cluster the passers, and keep sampling where diversity is still low. A passer meets computational criteria; binding still needs experimental validation.

  3. Refine within clusters. Choose generated parents from selected clusters for partial diffusion, as in section 4, track each child’s parent and score the children independently.

The configs used here (vhh.yaml, partial.yaml, common.yaml) and the eight inputs are in `examples/antibody_design <RosettaCommons/RFD4-Proteina>`__.

Download this notebook