PPI specificity tutorial using a YAML spec file: binders that pick their target#
This notebook designs minibinders against the Trk neurotrophin receptors and scores what they do not bind. This tutorial goes through two design challenges:
Specificity. Bind human TrkA, not its relatives TrkB and TrkC.
Multi-specificity. Bind TrkA of human, mouse and rabbit.
Both use one mechanism, a pre-oracle transform that keeps the designed binder, swaps in each structure you name, and refolds. Whether a strong off-target score passes or fails comes down to the threshold you write.
We turn the filter on in a basic inference run, point it the other way for multi-specificity, then wrap the same run in a reward guided search.
Sampling goes to Slurm, and the notebook submits and waits. The AF3 oracle needs the apptainer setup from Running the model.
Two mechanisms this rests on#
Nothing below needs new Python. Both knobs come from the metrics reference.
``metrics_tags`` is a run level list of strings. Every metric, oracle and pass criterion declares the tags it wants, through required_tags_any, required_tags_all and prohibited_tags, and runs only when the run’s tags satisfy it. An entry with no gate always runs. ppi_refolding brings the target templated AF3 refold and the binder RMSD metrics that read it. apo_refolding brings the binder alone fold. A metric whose tag you never pass gets skipped silently, which is nearly
always why a column is missing from per_example_metrics.csv.
Pre-oracle transforms decide what gets folded. One hangs off an oracle’s pre_oracle_transform, rewrites the structure on its way to the oracle, and can fan one design out into several outputs, each folded separately and reported as <oracle_key>.<save_id>.<metric>. ApoOnlyTransform is what the apo_refolding tag runs. OffTargetSubstitutionTransform, which keeps the binder and swaps its partner for a structure off disk, is what this tutorial is built on. Details in
pre-oracle transforms.
Setup#
Paths and Slurm resources, plus the helpers used throughout. run_config writes a run config from a shared preamble and the block that makes one run different from the next. launch submits it and waits for its metrics table. verdict_table, specificity_chart and complex_browser are then three views of the same result: the numbers against their thresholds, the spread across receptors, and the structures themselves.
[ ]:
import warnings
warnings.filterwarnings("ignore", message=".*NVML.*")
import os
import re
import shutil
import subprocess
import time
from pathlib import Path
import pandas as pd
import yaml
# 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():
REPO_ROOT = Path(os.environ["PROJECT_ROOT"])
# The Trk structures ship with the repository; inputs and run configs are written under OUT_ROOT.
TARGETS = REPO_ROOT / "examples/trka_crossreactivity/targets"
OUT_ROOT = Path(os.environ.get("PPI_TUTORIAL_RUNS", REPO_ROOT / "runs/ppi_specificity_tutorial"))
CFG_DIR = OUT_ROOT / "configs"
CFG_DIR.mkdir(parents=True, exist_ok=True)
# Slurm settings, all overridable by environment variable so the notebook moves between clusters.
CLUSTER = os.environ.get("CLUSTER", "local")
PARTITION = os.environ.get("SLURM_PARTITION", "gpu")
GRES = os.environ.get("SLURM_GRES", "gpu:small:1")
MEM = os.environ.get("SLURM_MEM", "32g")
TIME = os.environ.get("SLURM_TIME", "4:00:00")
OVERWRITE = False # False reuses existing outputs and skips sampling; True re-runs
os.environ.setdefault("PROJECT_ROOT", str(REPO_ROOT))
def launch(config: Path, njobs: int = 4, timeout: int = 14400) -> Path:
"""Submit a run config to Slurm and wait until its metrics table lands.
The config is a complete inference experiment living outside ``configs/`` — the
``--config-path`` pattern documented in ``configs/inference_experiment/TEMPLATE.yaml``.
Replicates are spread over ``njobs`` array tasks, one GPU each, and merged by the
concatenation job the launch script queues behind them.
"""
out_dir = Path(yaml.safe_load(config.read_text())["out_dir"])
csv = out_dir / "inference_metrics/per_example_metrics.csv"
if csv.exists() and not OVERWRITE:
print(f"[skip] {config.stem}: reusing {len(pd.read_csv(csv))} designs (OVERWRITE=True re-runs)")
return csv
shutil.rmtree(out_dir, ignore_errors=True)
# Create the sample directory up front: array tasks racing to mkdir it on NFS can fail.
(out_dir / "inference_samples").mkdir(parents=True, exist_ok=True)
t0 = time.time()
cmd = (f'bash "{REPO_ROOT}/scripts/launch_generate.sh" {CLUSTER} --njobs {njobs} '
f'--partition {PARTITION} --gres {GRES} --mem {MEM} --time {TIME} '
f'--code-dir "{REPO_ROOT}" --config-path "{config.parent}" --config-name {config.stem}')
out = subprocess.run(cmd, shell=True, capture_output=True, text=True, cwd=REPO_ROOT).stdout
job = re.search(r"Submitted batch job (\d+)", out)
print(f"[submit] {config.stem} over {njobs} tasks" + (f" (job {job.group(1)})" if job else ""))
deadline = time.time() + timeout
while time.time() < deadline:
if csv.exists():
print(f"[done] {config.stem}: {len(pd.read_csv(csv))} designs in {(time.time() - t0) / 60:.0f} min")
return csv
time.sleep(60)
raise TimeoutError(f"{config.stem} produced no metrics before the timeout")
# Every run below samples the same design job, from the same spec, under the same tags; what
# changes between them is the scoring. That invariant half lives here so each section can show
# only the block that makes its run different.
CONFIG_PREAMBLE = """
defaults:
- inference
- _self_
hydra:
searchpath: # lets this out-of-tree config find configs/
- file://${{oc.env:PROJECT_ROOT}}/configs
inputs: {spec}
out_dir: {out_dir}
run_name_: {name}
num_replicates: {replicates}
batch_size: {batch_size}
metrics_tags: [design, hotspot_conditioning, ppi_refolding, apo_refolding, off_target_refolding]
"""
def run_config(name: str, body: str, replicates: int = 12, batch_size: int = 3) -> Path:
"""Write a run config -- the preamble above, plus the ``body`` that makes this run different.
Returns the path, which is all :func:`launch` and the result views need: every one of them
reads what it needs back out of the file, so the config is the single source of truth for
where the run wrote, what it assayed, and what it had to clear.
"""
path = CFG_DIR / f"{name}.yaml"
path.write_text(CONFIG_PREAMBLE.format(
spec=SPEC, out_dir=OUT_ROOT / name.removeprefix("trka_"), name=name,
replicates=replicates, batch_size=batch_size).lstrip() + body)
# Echo the bars the file actually carries. The YAML is shown as prose above each of these
# cells, and this is what keeps that copy honest: the numbers here come off disk.
for criterion, spec in (yaml.safe_load(path.read_text())["metrics"].get("pass_criteria") or {}).items():
bars = ", ".join(_label(k, b) for k, b in spec["metric_thresholds"].items())
print(f"[config] {path.stem} -> {criterion}: {bars}")
return path
# The scoring block each run appends to that preamble. They live here, out of the way, because
# sections 2-4 show and explain them as YAML; keeping one copy each means the notebook runs the
# configs it prints rather than a paraphrase of them.
SPEC_BODY = f"""
metrics:
refolding_oracles:
# Specificity assay: keep the binder, swap in each off-target, refold once apiece.
oracle_off_target:
_target_: rfproteina.metrics.oracles.af3_oracle.AF3Oracle
oracle_name: af3_off_target
template_query: "is_loaded_target & is_protein"
aggregation: [median]
pre_oracle_transform:
_target_: rfproteina.metrics.oracles.pre_oracle_transform.OffTargetSubstitutionTransform
off_target_paths:
htrkb: {TARGETS}/hTrkB.cif
htrkc: {TARGETS}/hTrkC.cif
required_tags_any: [off_target_refolding]
pass_criteria:
trka_specific:
_target_: rfproteina.metrics.oracles.pass_criterion_metric.PassCriterionMetric
metric_thresholds:
oracle_ppi.min_ipae_coi: {{max: 1.5}} # binds hTrkA
oracle_ppi.iptm_coi: {{min: 80.0}}
refolded_target_aligned_binder_rmsd.rmsd: {{max: 2.0}} # at the epitope it was aimed at
refolded_apo_holo_rmsd.apo.rmsd: {{max: 1.5}} # and holds its fold apo
oracle_off_target.htrkb.min_ipae_coi: {{min: 5.0}} # does NOT bind hTrkB
oracle_off_target.htrkc.min_ipae_coi: {{min: 5.0}} # does NOT bind hTrkC
required_tags_any: [off_target_refolding]
"""
MULTI_BODY = f"""
metrics:
refolding_oracles:
oracle_off_target:
_target_: rfproteina.metrics.oracles.af3_oracle.AF3Oracle
oracle_name: af3_off_target
template_query: "is_loaded_target & is_protein"
aggregation: [median]
pre_oracle_transform:
_target_: rfproteina.metrics.oracles.pre_oracle_transform.OffTargetSubstitutionTransform
off_target_paths:
mtrka: {TARGETS}/mTrkA.cif # wanted, not avoided
rtrka: {TARGETS}/rTrkA.cif
required_tags_any: [off_target_refolding]
pass_criteria:
trka_crossreactive:
_target_: rfproteina.metrics.oracles.pass_criterion_metric.PassCriterionMetric
metric_thresholds:
oracle_ppi.min_ipae_coi: {{max: 1.5}} # binds the receptor it was designed on
oracle_ppi.iptm_coi: {{min: 80.0}}
refolded_target_aligned_binder_rmsd.rmsd: {{max: 2.0}}
oracle_off_target.mtrka.min_ipae_coi: {{max: 1.5}} # and the mouse orthologue, at the same bar
oracle_off_target.mtrka.iptm_coi: {{min: 80.0}}
oracle_off_target.rtrka.min_ipae_coi: {{max: 1.5}} # and the rabbit one, likewise
oracle_off_target.rtrka.iptm_coi: {{min: 80.0}}
required_tags_any: [off_target_refolding]
"""
SEARCH_BLOCK = """
generation:
args:
nsteps: 60
search:
n_branch: 2
beam_width: 2
save_rollouts: true
save_criteria:
oracle_ppi.min_ipae_coi: {max: 1.5}
oracle_ppi.iptm_coi: {min: 80.0}
refolded_target_aligned_binder_rmsd.rmsd: {max: 2.0}
refolded_apo_holo_rmsd.apo.rmsd: {max: 1.5}
oracle_off_target.htrkb.min_ipae_coi: {min: 5.0}
oracle_off_target.htrkc.min_ipae_coi: {min: 5.0}
checkpoints:
"30": &scoring_checkpoint
metrics:
- oracle_ppi
- oracle_apo
- oracle_off_target
- refolded_target_aligned_binder_rmsd
- refolded_apo_holo_rmsd
rewards:
oracle_ppi.iptm_coi: {weight: 1.0}
oracle_ppi.min_ipae_coi: {weight: -20.0, clip: {min: 0.5, max: 5.0}}
oracle_off_target.htrkb.min_ipae_coi: {weight: 10.0, clip: {max: 5.0}}
oracle_off_target.htrkc.min_ipae_coi: {weight: 10.0, clip: {max: 5.0}}
"60": *scoring_checkpoint
"""
def _thresholds(config: Path, criterion: str) -> dict[str, dict]:
"""The ``{metric: {min|max: value}}`` block one pass criterion is built from.
Every view below reads its bars from here rather than restating them, so a table, a chart and
the ``<criterion>.pass`` column the pipeline computed can never disagree about what was asked.
"""
spec = yaml.safe_load(Path(config).read_text())
return spec["metrics"]["pass_criteria"][criterion]["metric_thresholds"]
def _labels(d: pd.DataFrame) -> pd.Series:
"""Short per-design labels, as ``<what distinguishes it>.<array task>``.
Whatever every ``example_id`` in the table shares is dropped, so a plain run labels its designs
``0.1``, ``2.3``, while a searched run — whose ids also carry the checkpoint a sample was kept
at — keeps that: ``bsfinal_0.1`` is the final beam, ``bsc1_0.1`` was saved at checkpoint 1.
Every view names a design the same way, so a row in the table is a row in the chart.
"""
stem, job = (d["example_id"].astype(str).str.rsplit("_job", n=1).str[i] for i in (0, -1))
shared = os.path.commonprefix(list(stem))
return stem.str.removeprefix(shared[: shared.rfind("_") + 1]) + "." + job
def _label(key: str, bound: dict) -> str:
"""Column header for one threshold, e.g. ``htrkb.min_ipae >=2.5``."""
stem = (key.replace("oracle_ppi.", "on-target ")
.replace("oracle_off_target.", "")
.replace("refolded_target_aligned_binder_rmsd.rmsd", "dock rmsd")
.replace("refolded_apo_holo_rmsd.apo.rmsd", "apo-holo rmsd")
.replace("_coi", ""))
op, value = ("<=", bound["max"]) if "max" in bound else (">=", bound["min"])
return f"{stem} {op}{value:g}"
def verdict_table(csv: Path, config: Path, criterion: str):
"""One row per design, green where the config's pass criterion is met and red where it is not.
Thresholds are read back out of the run config, so the table and the ``<criterion>.pass``
column computed by the pipeline can never drift apart.
"""
thresholds = _thresholds(config, criterion)
d = pd.read_csv(csv)
ids = _labels(d)
met = pd.DataFrame({_label(k, b): (d[k] <= b["max"] if "max" in b else d[k] >= b["min"])
for k, b in thresholds.items()}).set_index(ids)
show = pd.DataFrame({_label(k, b): d[k].round(2) for k, b in thresholds.items()}).set_index(ids)
show["PASS"] = d[f"{criterion}.pass"].astype(int).to_numpy()
order = show.sort_values(["PASS", show.columns[0]], ascending=[False, True]).index
# The ink is set alongside the fill: the docs theme has a dark mode, and a pale wash keeping
# the theme's light text would leave these cells unreadable on half the site's readers.
colours = met.map(lambda ok: "background-color:#d8f0d8;color:#0b0b0b" if ok
else "background-color:#f7d7d7;color:#0b0b0b")
colours["PASS"] = ""
print(f"{int(show['PASS'].sum())}/{len(show)} designs pass `{criterion}`")
return (show.loc[order].style
.apply(lambda _: colours.loc[order], axis=None)
.format(precision=2)) # Styler prints the full float unless told not to
The other two views of a run. specificity_chart puts every design on one interface PAE axis, with a dot per receptor it was folded against. complex_browser shows the folds themselves, with every complex of a design superimposed on one receptor, so switching between them moves the binder and nothing else. Both take their thresholds and their target list from the same run config the table uses.
[ ]:
import io
import json
import uuid
import biotite.structure as struc
import biotite.structure.io.pdb as biopdb
import numpy as np
import py3Dmol
from atomworks.io import parse
from IPython.display import HTML
VIEWER_JS = "https://3Dmol.org/build/3Dmol-min.js"
BINDER_CHAIN = "B" # the designed chain; the run's chain_of_interest
# Categorical slots 1-3 of the docs palette as (light, dark) steps, assigned in the order the
# receptors are listed, so the on-target is always blue and one colour means one receptor in chart
# and viewer alike. The 3Dmol canvas is white in both themes, so it always takes the light step.
SERIES = [("#2a78d6", "#3987e5"), ("#eb6834", "#d95926"), ("#1baf7a", "#199e70")]
def targets_of(config: Path, on_target_label: str) -> list[tuple[str, str]]:
"""``[(label, metric prefix), …]`` for the receptors a run assayed, on-target first.
The off-targets are read out of the oracle's ``off_target_paths``, so adding a receptor to a
run config adds it to the chart and the viewer without touching either.
"""
spec = yaml.safe_load(Path(config).read_text())
off = spec["metrics"]["refolding_oracles"]["oracle_off_target"]["pre_oracle_transform"]["off_target_paths"]
return [(on_target_label, "oracle_ppi")] + [
(Path(path).stem, f"oracle_off_target.{name}") for name, path in off.items()
]
def specificity_chart(csv: Path, config: Path, criterion: str, targets, width=760, row_h=22):
"""One row per design, one dot per receptor, on a shared minimum-interface-PAE axis.
Rows are ordered by the on-target interface, so reading down the chart walks from the designs
that bound best to the ones that barely bound at all, and the off-target dots say what each of
them cost. Distinct thresholds are drawn once each as a labelled reference line.
"""
thresholds = _thresholds(config, criterion)
d = pd.read_csv(csv)
d = d.assign(label=_labels(d)).sort_values(f"{targets[0][1]}.min_ipae_coi")
rows = [{"label": r["label"], "pass": bool(r[f"{criterion}.pass"]),
"vals": [float(r[f"{p}.min_ipae_coi"]) for _, p in targets]}
for _, r in d.iterrows()]
hi = max(max(r["vals"]) for r in rows)
pad_l = 14 + 7 * max(len(r["label"]) for r in rows) # the label gutter, in monospace ems
pad_r, pad_t, pad_b = 26, 48, 28
plot_w, bottom = width - pad_l - pad_r, pad_t + row_h * len(rows)
height = bottom + pad_b
x = lambda v: pad_l + plot_w * min(v, hi) / hi # noqa: E731
# One line per distinct bar: "avoid >=5" is a single rule two receptors share, not two rules.
bars: dict[tuple[float, str], list[str]] = {}
for label, prefix in targets:
(op, value), = thresholds[f"{prefix}.min_ipae_coi"].items()
bars.setdefault((float(value), op), []).append(label)
# Drop any gridline a threshold already sits on -- they coincide at 5 A, and a dashed rule
# drawn over a solid gridline reads as one line rather than as two.
ruled = {v for v, _ in bars}
grid = "".join(
f'<line x1="{x(t):.1f}" y1="{pad_t - 6}" x2="{x(t):.1f}" y2="{bottom}" class="g"/>'
f'<text x="{x(t):.1f}" y="{pad_t - 12}" class="ax" text-anchor="middle">{t}</text>'
for t in range(0, int(hi) + 1, 5) if t not in ruled)
rules = "".join(
f'<line x1="{x(v):.1f}" y1="{pad_t - 6}" x2="{x(v):.1f}" y2="{bottom}" class="r"/>'
f'<text x="{x(v):.1f}" y="{pad_t - 28}" class="ax" text-anchor="middle">'
f'{"≤" if op == "max" else "≥"}{v:g} {" ".join(labels)}</text>'
for (v, op), labels in bars.items())
body = ""
for i, r in enumerate(rows):
y = pad_t + row_h * i + row_h / 2
body += (f'<line x1="{x(min(r["vals"])):.1f}" y1="{y:.1f}" x2="{x(max(r["vals"])):.1f}" '
f'y2="{y:.1f}" class="conn"/>')
body += "".join(
# The tooltip hangs off a transparent circle twice the radius of the dot: a 5px
# target is a hard thing to land on, and the dot itself stays out of the way.
f'<circle cx="{x(v):.1f}" cy="{y:.1f}" r="11" fill="transparent">'
f'<title>{targets[j][0]} min_ipae_coi {v:.2f} A</title></circle>'
f'<circle cx="{x(v):.1f}" cy="{y:.1f}" r="5" style="fill:var(--s{j % 3 + 1})" class="dot"/>'
for j, v in enumerate(r["vals"]))
body += (f'<text x="{pad_l - 8}" y="{y + 4:.1f}" class="lab" text-anchor="end">{r["label"]}</text>'
f'<text x="{width - pad_r + 6}" y="{y + 4:.1f}" class="{"ok" if r["pass"] else "no"}">'
f'{"✓" if r["pass"] else "✗"}</text>')
legend = "".join(f'<span style="margin-right:14px"><span style="color:var(--s{j % 3 + 1})">●</span>'
f' {t[0]}</span>' for j, t in enumerate(targets))
uid = "c" + uuid.uuid4().hex[:8]
return HTML(f"""
<div class="{uid}" style="font:13px/1.5 system-ui,sans-serif">
<style>
/* The light values live here rather than in a style= attribute: an inline custom property
outranks any selector, and would leave the dark overrides below with nothing to do. */
.{uid} {{ --ink2:#52514e; --muted:#898781; --grid:#e1e0d9; --conn:#c3c2b7; --surf:#fcfcfb;
--ok:#006300; --no:#a02020; --rule:#807e77;
--s1:{SERIES[0][0]}; --s2:{SERIES[1][0]}; --s3:{SERIES[2][0]}; }}
/* Dark is its own set of steps against the dark surface, not an automatic flip of the light
one. The media query follows the OS; the data-theme scope follows the site's toggle. */
@media (prefers-color-scheme: dark) {{ :root:not([data-theme="light"]) .{uid} {{
--ink2:#c3c2b7; --grid:#2c2c2a; --conn:#383835; --surf:#1a1a19; --ok:#0ca30c; --no:#e85d5d; --rule:#a5a39b;
--s1:{SERIES[0][1]}; --s2:{SERIES[1][1]}; --s3:{SERIES[2][1]}; }} }}
:root[data-theme="dark"] .{uid} {{
--ink2:#c3c2b7; --grid:#2c2c2a; --conn:#383835; --surf:#1a1a19; --ok:#0ca30c; --no:#e85d5d; --rule:#a5a39b;
--s1:{SERIES[0][1]}; --s2:{SERIES[1][1]}; --s3:{SERIES[2][1]}; }}
.{uid} svg {{ max-width:100%; height:auto; overflow:visible; }}
.{uid} .g {{ stroke:var(--grid); stroke-width:1; }}
.{uid} .r {{ stroke:var(--rule); stroke-width:1.5; stroke-dasharray:6 4; }}
.{uid} .conn {{ stroke:var(--conn); stroke-width:2; }}
.{uid} .dot {{ stroke:var(--surf); stroke-width:2; pointer-events:none; }}
.{uid} .ax {{ fill:var(--muted); font:11px ui-monospace,monospace; }}
.{uid} .lab {{ fill:var(--muted); font:11px ui-monospace,monospace; }}
.{uid} .ok {{ fill:var(--ok); font:13px system-ui; }}
.{uid} .no {{ fill:var(--no); font:13px system-ui; }}
.{uid} .cap {{ color:var(--muted); font-size:12px; }}
.{uid} .cap code {{ color:inherit; background:none; border:0; padding:0; font-size:inherit; }}
</style>
<div style="margin-bottom:2px">{legend}</div>
<svg viewBox="0 0 {width} {height}" width="{width}" role="img"
aria-label="Minimum interface PAE of every design against each receptor it was folded against">
{grid}{rules}{body}
<text x="{pad_l + plot_w / 2:.0f}" y="{height - 6}" class="ax" text-anchor="middle">
minimum interface PAE (Å) — lower is a more confident interface</text>
</svg>
<div class="cap">dashed line = the bar that receptor has to clear |
✓ / ✗ = <code>{criterion}</code>, over all of its columns</div>
</div>
""")
def to_pdb(atom_array: struc.AtomArray) -> str:
"""Render an AtomArray as a PDB string, which is what the 3Dmol viewers take."""
f = biopdb.PDBFile()
with warnings.catch_warnings():
# biotite's writer trips a numpy DeprecationWarning on every call. Suppressed here rather
# than globally: importing atomworks resets the filters, so a filter set up top is dropped.
warnings.simplefilter("ignore", DeprecationWarning)
f.set_structure(atom_array)
buf = io.StringIO()
f.write(buf)
return buf.getvalue()
def _backbone_pdb(path: str, reference: struc.AtomArray | None = None):
"""Read a complex, keep its backbone, and superimpose its receptor onto ``reference``.
Returns the complex, the binder alone, and the receptor — the last so the first complex read
can become the frame every later one is put into. Backbone atoms are enough for a cartoon and
keep the notebook a few megabytes rather than a few tens of them.
"""
aa = parse(path)["assemblies"]["1"][0]
aa = aa[np.isin(aa.atom_name, ["N", "CA", "C", "O"]) & ~np.isnan(aa.coord).any(-1)]
receptor = aa[np.asarray(aa.chain_id).astype(str) != BINDER_CHAIN]
if reference is not None:
# Homolog superposition: the receptors are relatives, not copies, so the correspondence
# has to be found rather than assumed.
_, transform, _, _ = struc.superimpose_homologs(reference, receptor)
aa, receptor = transform.apply(aa), transform.apply(receptor)
binder = aa[np.asarray(aa.chain_id).astype(str) == BINDER_CHAIN]
return to_pdb(aa), to_pdb(binder), receptor
def complex_browser(csv: Path, targets, n=4, width=760, height=520):
"""Pick a design, then pick which receptor it was folded against, and look at the fold.
Every complex is put into one receptor's frame and the camera is held across switches, so the
off-target views differ from the on-target one only in what the binder did. The on-target
binder stays on screen as a pale ghost to compare against: a design that fails an off-target
by sliding off the epitope looks different from one that stays put and merely scores worse.
"""
d = pd.read_csv(csv)
d = d.assign(label=_labels(d)).sort_values(f"{targets[0][1]}.min_ipae_coi").head(n)
designs, ref = [], None
for _, row in d.iterrows():
views = []
for label, prefix in targets:
pdb, binder, receptor = _backbone_pdb(row[f"{prefix}.refolded_structure_path"], ref)
ref = receptor if ref is None else ref
views.append({"label": label, "pdb": pdb, "binder": binder,
"note": f"min_ipae {row[f'{prefix}.min_ipae_coi']:.2f} A"
f" ipTM {row[f'{prefix}.iptm_coi']:.0f}"})
designs.append({"label": row["label"], "views": views})
uid = "b" + uuid.uuid4().hex[:10]
dbtn = "".join(f'<button id="{uid}_d{i}" onclick="{uid}_pick({i},null)" class="{uid}_b">'
f'{s["label"]}</button>' for i, s in enumerate(designs))
vbtn = "".join(f'<button id="{uid}_v{j}" onclick="{uid}_pick(null,{j})" class="{uid}_b">'
f'<span style="color:{SERIES[j % 3][0]}">●</span> {t[0]}</button>'
for j, t in enumerate(targets))
return HTML(f"""
<style>
.{uid}_b {{ margin:2px 4px 2px 0; padding:4px 10px; border:1px solid #bbb; border-radius:4px;
background:#f4f4f4; color:#0b0b0b; cursor:pointer; font:12px ui-monospace,monospace; }}
</style>
<div style="font:13px/1.5 system-ui,sans-serif">
<div style="margin-bottom:2px;color:#898781;font-size:11px">design</div>
<div>{dbtn}</div>
<div style="margin:6px 0 2px;color:#898781;font-size:11px">folded against</div>
<div>{vbtn}</div>
<div id="{uid}_note" style="margin:8px 0 6px;color:#898781;font:12px ui-monospace,monospace"> </div>
<div id="{uid}" style="width:{width}px;max-width:100%;height:{height}px;position:relative"></div>
<div style="margin-top:6px;color:#898781;font-size:12px">
grey cartoon = receptor | coloured cartoon = the designed binder |
pale grey = where that same binder sits on the on-target receptor
</div>
</div>
<script src="{VIEWER_JS}"></script>
<script>
(function() {{
var data = {json.dumps(designs)}, colours = {json.dumps([c for c, _ in SERIES])};
var di = 0, vi = 0, viewer = null, view = null;
function draw() {{
var views = data[di].views;
// Carry the camera -- including whatever the reader has rotated it to -- across every switch.
// The complexes share a frame, so a binder that moved off its epitope reads as movement
// rather than as an unrelated picture.
if (view !== null) view = viewer.getView();
viewer.clear();
viewer.addModel(views[vi].pdb, "pdb"); // model 0: the selected complex
viewer.setStyle({{model: 0}}, {{cartoon: {{color: "#c3c2b7"}}}});
viewer.setStyle({{model: 0, chain: "{BINDER_CHAIN}"}},
{{cartoon: {{color: colours[vi % colours.length]}}}});
if (vi !== 0) {{ // model 1: the on-target binder, as the reference
viewer.addModel(views[0].binder, "pdb");
viewer.setStyle({{model: 1}}, {{cartoon: {{color: "#8f8d86", opacity: 0.45}}}});
}}
if (view === null) {{ viewer.zoomTo(); view = viewer.getView(); }} else {{ viewer.setView(view); }}
viewer.render();
document.getElementById("{uid}_note").textContent =
data[di].label + " vs " + views[vi].label + " \\u2014 " + views[vi].note;
for (var i = 0; i < data.length; i++) mark("{uid}_d" + i, i === di);
for (var j = 0; j < views.length; j++) mark("{uid}_v" + j, j === vi);
}}
function mark(id, on) {{
var b = document.getElementById(id);
if (b) {{ b.style.background = on ? "#cde2fb" : "#f4f4f4";
b.style.fontWeight = on ? "bold" : "normal"; }}
}}
window["{uid}_pick"] = function(d, v) {{
if (d !== null) di = d;
if (v !== null) vi = v;
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>
""")
1. The design job#
One target, one epitope, one de novo binder. The same design job feeds both halves (specificity and multi-specificity) of this tutorial, and only the scoring differs.
Structure |
Chain |
Residues |
Role |
|---|---|---|---|
hTrkA |
A |
278 to 382 |
on-target in both cases, with hotspots 294, 296 and 333 |
hTrkB |
A |
278 to 382 |
off-target to avoid (section 2) |
hTrkC |
A |
295 to 399 |
off-target to avoid (section 2) |
mTrkA |
A |
282 to 385 |
orthologue to hit (section 3) |
rTrkA |
D |
853 to 950 |
orthologue to hit (section 3) |
The input is a condition spec. The target backbone is coordinate conditioned with C_CRD, so the receptor stays rigid while its side chains repack. Its sequence is revealed with C_SEQ. The RFdiffusion paper’s TrkA epitope is marked as hotspots with C_HOT, to steer the binder onto the right face. The binder chain carries no conditions at all, so the model generates it from noise, at a length it picks between 55 and 75 residues.
[ ]:
HOTSPOTS = [294, 296, 333] # the RFdiffusion-paper TrkA epitope, in hTrkA numbering
SPEC = CFG_DIR / "trka_binder.yaml"
SPEC.write_text(f"""
global_args:
input_as_annotated_cif: false # Starting from a normal (not annotated) CIF file
trka_binder:
input: {TARGETS}/hTrkA.cif
contig: "A278-382, /0, BINDER:55-75" # fixed target, then a de novo 55-75 residue binder
conditions:
C_CRD: # hold the target backbone; side chains stay free
true:
- select: "sel('A') and is_protein_backbone"
C_SEQ: # reveal the target sequence
true:
- select: "sel('A')"
C_HOT: # steer the binder onto the epitope
true:
- select: "sel('A/*/{HOTSPOTS}')"
""".lstrip())
The epitope, in orange, on the hTrkA fragment the binder is designed against.
[ ]:
target = parse(f"{TARGETS}/hTrkA.cif")["assemblies"]["1"][0]
target = target[target.mask("~has_nan_coord()")] # unresolved side-chain atoms
view = py3Dmol.view(width=640, height=420)
view.addModel(to_pdb(target), "pdb")
view.setStyle({"cartoon": {"color": "0xb9c0c9"}})
view.addStyle({"resi": HOTSPOTS}, {"stick": {"colorscheme": "orangeCarbon", "radius": 0.3}})
view.setBackgroundColor("0xffffff")
view.zoomTo()
view.show()
2. Specificity: bind hTrkA, avoid hTrkB and hTrkC#
Specificity is a scoring question, so it lives in the run config rather than the design spec. Three blocks turn it on.
metrics_tags.ppi_refoldingfolds the complex,apo_refoldingfolds the binder alone, andoff_target_refoldingenables the oracle below. All three runs share these, so they sit inCONFIG_PREAMBLEup in Setup, and each cell below shows only what differs.oracle_off_target, an AF3 oracle whosepre_oracle_transformswaps each structure inoff_target_pathsfor the target and refolds, reporting underoracle_off_target.<name>.*.template_queryhands AF3 the off-target’s known structure, so the fold is about the interface rather than the receptor.pass_criteria, one criterion covering both halves, the on-target bars and the off-target ones together.
Cost is linear in what you ask for. One AF3 fold per design per entry, on top of the holo and apo folds, which comes to four per design here.
The bars the criterion sets:
Bar |
Config key |
Value |
What it asks for |
|---|---|---|---|
binds |
|
|
AF3 is confident about the designed interface |
binds |
|
|
and confident about it as a whole |
docks as designed |
|
|
AF3’s own fold puts the binder where the design put it |
folds alone |
|
|
and keeps that fold without the target |
avoids hTrkB |
|
|
AF3 declines to model an interface |
avoids hTrkC |
|
|
likewise |
The exclusion bar sits far above the binding bar, at 5 Ã… rather than just outside 1.5, because what we want from an off-target is AF3 declining to model an interface at all. A binder scoring 2.5 Ã… against hTrkB is not specific. It is just worse at hTrkB.
Written out, the block is SPEC_BODY, appended to the preamble from Setup. {TARGETS} is the repository’s examples/trka_crossreactivity/targets.
metrics:
refolding_oracles:
# Specificity assay: keep the binder, swap in each off-target, refold once apiece.
oracle_off_target:
_target_: rfproteina.metrics.oracles.af3_oracle.AF3Oracle
oracle_name: af3_off_target
template_query: "is_loaded_target & is_protein"
aggregation: [median]
pre_oracle_transform:
_target_: rfproteina.metrics.oracles.pre_oracle_transform.OffTargetSubstitutionTransform
off_target_paths:
htrkb: {TARGETS}/hTrkB.cif
htrkc: {TARGETS}/hTrkC.cif
required_tags_any: [off_target_refolding]
pass_criteria:
trka_specific:
_target_: rfproteina.metrics.oracles.pass_criterion_metric.PassCriterionMetric
metric_thresholds:
oracle_ppi.min_ipae_coi: {max: 1.5} # binds hTrkA
oracle_ppi.iptm_coi: {min: 80.0}
refolded_target_aligned_binder_rmsd.rmsd: {max: 2.0} # AF3 agrees with the designed dock
refolded_apo_holo_rmsd.apo.rmsd: {max: 1.5} # and holds its fold apo
oracle_off_target.htrkb.min_ipae_coi: {min: 5.0} # does NOT bind hTrkB
oracle_off_target.htrkc.min_ipae_coi: {min: 5.0} # does NOT bind hTrkC
required_tags_any: [off_target_refolding]
[ ]:
SPEC_CFG = run_config("trka_specificity", SPEC_BODY)
spec_csv = launch(SPEC_CFG)
Results#
min_ipae_coi is AF3’s minimum interface PAE in Å, where low means a confident interface. The on-target columns want it low and the off-target columns want it high. Each cell turns green where the config’s threshold is met.
[ ]:
verdict_table(spec_csv, SPEC_CFG, "trka_specific")
The same run on one axis. Each row is a design, ordered by how well it bound hTrkA. The three dots are the same binder folded against the three receptors, and the dashed lines are the two bars the criterion sets. Specificity is the horizontal gap between the blue dot and the other two. A design whose dots sit together is a family binder, whatever its on-target number says.
[ ]:
SPEC_TARGETS = targets_of(SPEC_CFG, "hTrkA")
specificity_chart(spec_csv, SPEC_CFG, "trka_specific", SPEC_TARGETS)
The two halves are independent. A design can bind hTrkA beautifully and still light up hTrkB. Bundling them into one criterion is what answers both at once, and here the twelve designs fail in three different places. Seven clear both binding bars. Of those, two (0.1, 1.3) light up hTrkC at 2.55 and 3.52 Ã…, three (1.0, 2.1, 0.3) miss the 2 Ã… dock bar at 2.09 to 2.28 Ã…, meaning AF3 refolds them somewhere other than where they were designed, and two survive everything.
Which off-target bites is itself a result. hTrkB never gets close, with a lowest interface PAE of 9.74 Ã… across all twelve, while hTrkC accounts for both specificity failures. A run that assayed hTrkB alone would have reported perfect specificity and been wrong.
Below are the four best on-target designs. Pick one, then switch receptors. The binder’s on-target pose stays on screen in pale grey, so an off-target view shows what AF3 did with the same 70 residues once the receptor changed. Usually it shifts them off the epitope, which is the structural form of a high interface PAE.
[ ]:
complex_browser(spec_csv, SPEC_TARGETS)
3. Multi-specificity: bind human, mouse and rabbit TrkA#
Same design job, same oracle. All that changes is which structures go into off_target_paths and what you ask of them. Mouse and rabbit TrkA are wanted here, so they get the on-target bars, value for value, and nothing gets excluded.
Grading the inherited interfaces on a curve is tempting, since the model never got to optimise against them. But a binder that only half works in the animal model is not usable in the animal model.
The only difference from section 2 is which structures are named and what is asked of them.
metrics:
refolding_oracles:
oracle_off_target:
_target_: rfproteina.metrics.oracles.af3_oracle.AF3Oracle
oracle_name: af3_off_target
template_query: "is_loaded_target & is_protein"
aggregation: [median]
pre_oracle_transform:
_target_: rfproteina.metrics.oracles.pre_oracle_transform.OffTargetSubstitutionTransform
off_target_paths:
mtrka: {TARGETS}/mTrkA.cif # wanted, not avoided
rtrka: {TARGETS}/rTrkA.cif
required_tags_any: [off_target_refolding]
pass_criteria:
trka_crossreactive:
_target_: rfproteina.metrics.oracles.pass_criterion_metric.PassCriterionMetric
metric_thresholds:
oracle_ppi.min_ipae_coi: {max: 1.5} # binds the receptor it was designed on
oracle_ppi.iptm_coi: {min: 80.0}
refolded_target_aligned_binder_rmsd.rmsd: {max: 2.0}
oracle_off_target.mtrka.min_ipae_coi: {max: 1.5} # and the mouse orthologue, at the same bar
oracle_off_target.mtrka.iptm_coi: {min: 80.0}
oracle_off_target.rtrka.min_ipae_coi: {max: 1.5} # and the rabbit one, likewise
oracle_off_target.rtrka.iptm_coi: {min: 80.0}
required_tags_any: [off_target_refolding]
[ ]:
MULTI_CFG = run_config("trka_multispecificity", MULTI_BODY)
multi_csv = launch(MULTI_CFG)
Results#
[ ]:
verdict_table(multi_csv, MULTI_CFG, "trka_crossreactive")
The chart is the same one as before, and the ask has simply turned around. All three bars now sit at 1.5 Ã…, and a design passes by pulling its dots together on the left rather than spreading them apart. Where cross-reactivity breaks is a per species question, and the columns answer it directly. An orthologue whose dots track the on-target one is inherited for free, while one sitting out to the right is the species the epitope actually disagrees with.
[ ]:
MULTI_TARGETS = targets_of(MULTI_CFG, "hTrkA")
specificity_chart(multi_csv, MULTI_CFG, "trka_crossreactive", MULTI_TARGETS)
And the same folds, to see whether a refused orthologue is a binder thrown off the epitope or one still sitting on it with less confidence. Those two failure modes want different fixes, and only the structures tell them apart.
[ ]:
complex_browser(multi_csv, MULTI_TARGETS)
The two orthologues come apart, and not evenly. Of the seven designs that bind hTrkA at 1.5 Ã… or better, rabbit TrkA is inherited by six at a median 1.18 Ã…, essentially the human number, while mouse TrkA is inherited by two, at a median 2.17 Ã….
No design clears the full criterion, which also wants ipTM ≥ 80 on all three species. At twelve replicates that is a yield to sample past rather than a result. The assay is working, and it is telling you where the cost sits. Rabbit comes for free. Mouse is what will set the yield when this runs at scale.
4. Searching for it#
Both runs above sample blind and score afterwards. Reward guided search scores partial trajectories at checkpoints instead, branching from the survivors and dropping the rest, so the compute follows the interfaces that are already forming. It is one added search: block, and everything from section 2 stays as it is.
The checkpoint explicitly lists the on-target, apo, and off-target oracles plus the two structural comparison metrics, so a candidate gets judged on its on-target interface and on both off-targets: iptm_coi up and min_ipae_coi down on hTrkA, and min_ipae_coi up on hTrkB and hTrkC. The interface PAE terms are clipped where their pass bars sit, so nothing is earned past a bar already cleared. iptm_coi is left unclipped, because it is what still separates two candidates once both
are over every bar.
The save criteria repeat the raw thresholds used in the final table. Their metric providers are explicitly listed at both checkpoints; the search does not compute pass/fail metrics or infer dependencies. The YAML anchor reuses the same scoring configuration at steps 30 and 60.
The bill is four AF3 folds per candidate per checkpoint rather than one. To cut it, drop oracle_off_target from the checkpoint metric lists along with its two reward terms. Search then steers on the on-target interface alone, and replace the save criteria with a bar on search.reward.
SPEC_BODY from section 2 is reused unchanged. This block gets appended to it, and the replicate count drops from 12 to 4, because each one now costs a tree rather than a trajectory.
generation:
args:
nsteps: 60
search:
n_branch: 2
beam_width: 2
save_rollouts: true
save_criteria:
oracle_ppi.min_ipae_coi: {max: 1.5}
oracle_ppi.iptm_coi: {min: 80.0}
refolded_target_aligned_binder_rmsd.rmsd: {max: 2.0}
refolded_apo_holo_rmsd.apo.rmsd: {max: 1.5}
oracle_off_target.htrkb.min_ipae_coi: {min: 5.0}
oracle_off_target.htrkc.min_ipae_coi: {min: 5.0}
checkpoints:
"30": &scoring_checkpoint
metrics:
- oracle_ppi
- oracle_apo
- oracle_off_target
- refolded_target_aligned_binder_rmsd
- refolded_apo_holo_rmsd
rewards:
oracle_ppi.iptm_coi: {weight: 1.0}
oracle_ppi.min_ipae_coi: {weight: -20.0, clip: {min: 0.5, max: 5.0}}
oracle_off_target.htrkb.min_ipae_coi: {weight: 10.0, clip: {max: 5.0}}
oracle_off_target.htrkc.min_ipae_coi: {weight: 10.0, clip: {max: 5.0}}
"60": *scoring_checkpoint
[ ]:
SEARCH_CFG = run_config("trka_specificity_search", SPEC_BODY + SEARCH_BLOCK,
replicates=4, batch_size=1)
search_csv = launch(SEARCH_CFG)
Results#
Search can return more designs than replicates. You get the final beam, plus every candidate that already cleared trka_specific on its way out of it.
[ ]:
verdict_table(search_csv, SEARCH_CFG, "trka_specific")
The same chart as section 2, over the searched run. Both quantities the reward was built from show up at once. The blue dots want to be against the left edge, and the orange and green ones out past 5 Ã….
[ ]:
specificity_chart(search_csv, SEARCH_CFG, "trka_specific", SPEC_TARGETS)
And side by side with the unguided run of section 2, which sampled the same design job with the same criterion and no reward. worst off-target is the smaller of the two off-target interface PAEs, so it names the receptor each design came closest to also binding, which is the one that decides its specificity.
[ ]:
summary = pd.DataFrame(
[{"designs": len(d),
"median on-target min_ipae": d["oracle_ppi.min_ipae_coi"].median(),
"best on-target min_ipae": d["oracle_ppi.min_ipae_coi"].min(),
"median worst off-target min_ipae": d[[f"{p}.min_ipae_coi" for _, p in SPEC_TARGETS[1:]]].min(axis=1).median(),
"passing trka_specific": int(d["trka_specific.pass"].sum())}
for d in (pd.read_csv(spec_csv), pd.read_csv(search_csv))],
index=["unguided (section 2)", "beam search"],
).round(2)
summary.style.format(precision=2)
Search moved every column it was paid to move. The median on-target interface goes from 1.33 Ã… to 1.09 Ã…, the best from 0.96 Ã… to 0.80 Ã…, the median worst off-target out from 9.81 Ã… to 11.66 Ã…, and 8 of 14 designs pass where 2 of 12 did. Same design job, same criterion, a third of the replicates.
The off-target column improved rather than drifting because the reward pays for it. A candidate that starts engaging hTrkC gets pruned at the next checkpoint instead of surviving to be filtered at the end. The 14 rows are 8 final beam designs plus 6 that save_criteria kept mid trajectory, carrying the bsc0_* and bsc1_* labels, which is why a searched run can return more designs than it had replicates.
Summary#
To change |
Edit |
|---|---|
Which structures are assayed |
|
Avoid or hit an off-target |
|
How hard the bar is |
the threshold values under |
What search optimises |
|
What search keeps beyond the final beam |
|
How hard search looks |
|
The off-target oracle knows nothing about biology. It folds the binder against whatever you hand it. Anti-targets, orthologues and family members are one mechanism pointed in different directions, so a real campaign names all of them in one off_target_paths block and writes one criterion over the lot.
None of it is specific to AF3 either. pre_oracle_transform is a property of the oracle base class, so the same block works against the AF2 and OpenFold3 oracles. See configs/metrics/refolding_oracles/ for what each one exposes, plus the search guide and the metrics reference for the rest.