boltzgen — designing a GLP-1 receptor binder

3IOL ─┬─ chain A (receptor ECD) ──► BoltzGen ──► designed peptides
      │                                               │
      └─ chain B (native GLP-1) ──► epitope ──────► compared: same groove?
             the answer key

super volume create boltzgen-weights     # one-time: holds 8.4 GB of checkpoints

# 1. fetch the target, find the epitope, write the design spec (local, stdlib)
python3 prep_target.py

# 2. warm a GPU and cache the checkpoints into the volume (one-time, ~4 min)
super create --sku gpu_1x_h100
super upload <sid> .
super exec <sid> --volume boltzgen-weights:/workspace/.cache -- \
    boltzgen download all --cache /workspace/.cache/boltzgen

# 3. design. 100 designs, keep the best 10.
super spawn <sid> --volume boltzgen-weights:/workspace/.cache -- bash -c \
    'cd /workspace/boltzgen && boltzgen run glp1r.yaml --output out_100 \
     --protocol peptide-anything --num_designs 100 --budget 10 \
     --devices 1 --cache /workspace/.cache/boltzgen'
super logs <sid> <pid> --follow

# 4. score the survivors against native GLP-1
super exec <sid> -- bash -c 'cd /workspace/boltzgen && python3 score.py --output out_100'

# Design a peptide binder for the GLP-1 receptor extracellular domain.
#
# Target: chain A of 3IOL, the GLP-1R ECD. Chain B of that structure is the
# native GLP-1 peptide, which we deliberately leave out — recovering something
# that binds where it binds is the point of the experiment.
entities:
  # The designed peptide. 26 residues of GLP-1 are resolved against the
  # ECD, so we let BoltzGen pick a length in a window around that.
  - protein:
      id: P
      sequence: 21..31

  - file:
      path: 3IOL.cif
      include:
        - chain:
            id: A
      # Hotspots: the 5 receptor residues most buried by native GLP-1.
      binding_types:
        - chain:
            id: A
            binding: 32,39,68,69,91
      structure_groups: "all"

{
  "pdb_id": "3IOL",
  "target_chain": "A",
  "native_chain": "B",
  "native_sequence": "GTFTSDVSSYLEGQAAKEFIAWLVKG",
  "native_length": 26,
  "target_sequence": "TVSLWETVQKWREYRRQCQRSLTEDPPPATDLFCNRTFDEYACWPDGEPGSFVNVSCPWYLPWASSVPQGHVYRFCTAEGLWLQKDNSSLPWRDLSECEE",
  "target_first_resnum": 29,
  "cutoff": 5.0,
  "epitope": [
    29,
    30,
    31,
    32,
    35,
    36,
    39,
    67,
    68,
    69,
    88,
    89,
    90,
    91,
    121,
    123,
    128
  ],
  "hotspots": [
    32,
    39,
    68,
    69,
    91
  ]
}

#!/usr/bin/env python3
"""Prepare the GLP-1 receptor target for BoltzGen and extract the answer key.

PDB 3IOL is GLP-1 (chain B) bound to the extracellular domain of the GLP-1
receptor (chain A). We hand BoltzGen chain A and ask it to design a peptide
that binds the same groove — which means the native peptide is a ground truth
we can score designs against.

Writes three files:
  3IOL.cif    the target, handed to BoltzGen unmodified (it selects the chain)
  glp1r.yaml  the design spec, with the hotspot residues filled in
  native.json the answer key: native GLP-1 sequence + the full native epitope

The .pdb of the same entry is downloaded purely because its fixed-column format
is trivial to parse; residue numbering is `auth` in both, so the numbers we
compute from the .pdb are the ones BoltzGen reads out of the .cif.
"""

import argparse
import json
import urllib.request

AA3TO1 = {
    "ALA": "A", "ARG": "R", "ASN": "N", "ASP": "D", "CYS": "C", "GLN": "Q",
    "GLU": "E", "GLY": "G", "HIS": "H", "ILE": "I", "LEU": "L", "LYS": "K",
    "MET": "M", "PHE": "F", "PRO": "P", "SER": "S", "THR": "T", "TRP": "W",
    "TYR": "Y", "VAL": "V",
}

SPEC_TEMPLATE = """\
# Design a peptide binder for the GLP-1 receptor extracellular domain.
#
# Target: chain A of 3IOL, the GLP-1R ECD. Chain B of that structure is the
# native GLP-1 peptide, which we deliberately leave out — recovering something
# that binds where it binds is the point of the experiment.
entities:
  # The designed peptide. {n_native} residues of GLP-1 are resolved against the
  # ECD, so we let BoltzGen pick a length in a window around that.
  - protein:
      id: P
      sequence: {lo}..{hi}

  - file:
      path: 3IOL.cif
      include:
        - chain:
            id: A
      # Hotspots: the {n_hot} receptor residues most buried by native GLP-1.
      binding_types:
        - chain:
            id: A
            binding: {hotspots}
      structure_groups: "all"
"""


def fetch(url):
    return urllib.request.urlopen(url, timeout=60).read().decode()


def parse_atoms(pdb_text, chain):
    """[(resseq, resname, x, y, z)] of heavy atoms in one chain, first altloc."""
    out = []
    for line in pdb_text.splitlines():
        if not line.startswith("ATOM") or line[21] != chain:
            continue
        if line[16] not in (" ", "A") or line[76:78].strip() == "H":
            continue
        out.append((
            int(line[22:26]), line[17:20].strip(),
            float(line[30:38]), float(line[38:46]), float(line[46:54]),
        ))
    if not out:
        raise SystemExit(f"no ATOM records for chain {chain}")
    return out


def sequence_of(atoms):
    """One-letter sequence in residue-number order."""
    seen = {}
    for resseq, resname, *_ in atoms:
        seen.setdefault(resseq, resname)
    return "".join(AA3TO1.get(seen[r], "X") for r in sorted(seen))


def contact_counts(target_atoms, ligand_atoms, cutoff):
    """{target resseq: number of ligand heavy atoms within cutoff}."""
    cut2 = cutoff * cutoff
    counts = {}
    # ponytail: brute force, ~1k x ~250 atoms. numpy if a target ever gets big.
    for tseq, _, tx, ty, tz in target_atoms:
        n = 0
        for _, _, lx, ly, lz in ligand_atoms:
            if (tx - lx) ** 2 + (ty - ly) ** 2 + (tz - lz) ** 2 <= cut2:
                n += 1
        if n:
            counts[tseq] = counts.get(tseq, 0) + n
    return counts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pdb-id", default="3IOL")
    ap.add_argument("--target-chain", default="A", help="receptor ECD")
    ap.add_argument("--native-chain", default="B", help="native GLP-1 peptide")
    ap.add_argument("--cutoff", type=float, default=5.0, help="contact distance (A)")
    ap.add_argument("--hotspots", type=int, default=5,
                    help="how many of the most-buried residues to hand BoltzGen")
    ap.add_argument("--length-window", type=int, default=5,
                    help="design lengths to allow either side of the native length")
    args = ap.parse_args()

    pid = args.pdb_id.upper()
    cif = fetch(f"https://files.rcsb.org/download/{pid}.cif")
    with open(f"{pid}.cif", "w") as fh:
        fh.write(cif)

    pdb = fetch(f"https://files.rcsb.org/download/{pid}.pdb")
    target = parse_atoms(pdb, args.target_chain)
    native = parse_atoms(pdb, args.native_chain)

    native_seq = sequence_of(native)
    target_seq = sequence_of(target)
    target_first = min(a[0] for a in target)
    counts = contact_counts(target, native, args.cutoff)
    epitope = sorted(counts)
    hotspots = sorted(sorted(counts, key=counts.get, reverse=True)[: args.hotspots])

    lo = max(4, len(native_seq) - args.length_window)
    hi = len(native_seq) + args.length_window
    with open("glp1r.yaml", "w") as fh:
        fh.write(SPEC_TEMPLATE.format(
            lo=lo, hi=hi, n_hot=len(hotspots), n_native=len(native_seq),
            hotspots=",".join(str(r) for r in hotspots),
        ))

    with open("native.json", "w") as fh:
        json.dump({
            "pdb_id": pid,
            "target_chain": args.target_chain,
            "native_chain": args.native_chain,
            "native_sequence": native_seq,
            "native_length": len(native_seq),
            # BoltzGen renumbers the target in its output; score.py realigns by
            # matching this sequence, so it needs the numbering it came from.
            "target_sequence": target_seq,
            "target_first_resnum": target_first,
            "cutoff": args.cutoff,
            "epitope": epitope,
            "hotspots": hotspots,
        }, fh, indent=2)

    print(f"{pid}: target chain {args.target_chain} = {len(set(a[0] for a in target))} residues")
    print(f"native {args.native_chain} = {native_seq} ({len(native_seq)} aa)")
    print(f"epitope ({args.cutoff} A, {len(epitope)} residues): "
          f"{','.join(str(r) for r in epitope)}")
    print(f"hotspots handed to BoltzGen: {','.join(str(r) for r in hotspots)}")
    print(f"wrote {pid}.cif, glp1r.yaml, native.json")


if __name__ == "__main__":
    main()

"""Ray-trace the GLP-1R binding site: native hormone vs designed peptide.

Run headless on a CPU box that has PyMOL:  pymol -cq render_site.py

Loads the native complex (3IOL: receptor chain A + GLP-1 chain B) and a designed
complex from `boltzgen run`, superimposes the two copies of the receptor, and
draws both peptides in the same groove. The epitope residues — the receptor
positions within 5 A of the native hormone — are drawn as sticks, so you can see
that the design contacts the same ones.

Writes results/binding_site.png and results/epitope_closeup.png.
"""

import json
import os

from pymol import cmd

HERE = os.getcwd()
RESULTS = os.path.join(HERE, "results")
os.makedirs(RESULTS, exist_ok=True)

with open(os.path.join(HERE, "native.json")) as fh:
    NATIVE = json.load(fh)
EPITOPE = "+".join(str(r) for r in NATIVE["epitope"])

cmd.load(os.path.join(HERE, "3IOL.cif"), "native")
cmd.load(os.path.join(RESULTS, "best_recovery.cif"), "designed")

# BoltzGen relabels chains: in its output the design is chain A and the target
# is chain B — the opposite of the crystal. Superimpose the two receptors so the
# peptides land in a shared frame.
cmd.create("receptor", "native and chain A and polymer")
cmd.create("hormone", "native and chain B and polymer")
cmd.create("design", "designed and chain A and polymer")
cmd.create("design_target", "designed and chain B and polymer")
cmd.super("design_target", "receptor")
cmd.matrix_copy("design_target", "design")
cmd.delete("native")
cmd.delete("designed")
cmd.delete("design_target")

cmd.hide("everything")
cmd.show("cartoon", "receptor or hormone or design")
cmd.dss()

cmd.color("grey70", "receptor")
cmd.color("skyblue", "hormone")
cmd.color("orange", "design")
cmd.set("cartoon_transparency", 0.15, "receptor")

# The answer key, drawn: receptor residues the real hormone touches.
cmd.show("sticks", f"receptor and resi {EPITOPE} and not name N+C+O")
cmd.color("firebrick", f"receptor and resi {EPITOPE} and elem C")
cmd.set("stick_radius", 0.18)

cmd.bg_color("white")
cmd.set("ray_shadows", 0)
cmd.set("antialias", 2)
cmd.set("ray_opaque_background", 1)
cmd.set("cartoon_fancy_helices", 1)

cmd.orient("receptor or hormone or design")
cmd.zoom("receptor or hormone or design", 3)
cmd.ray(1600, 1100)
cmd.png(os.path.join(RESULTS, "binding_site.png"), dpi=200)
print("wrote results/binding_site.png")

cmd.orient(f"receptor and resi {EPITOPE}")
cmd.zoom(f"receptor and resi {EPITOPE}", 4)
cmd.ray(1400, 1000)
cmd.png(os.path.join(RESULTS, "epitope_closeup.png"), dpi=200)
print("wrote results/epitope_closeup.png")

#!/usr/bin/env python3
"""Score BoltzGen designs against the native GLP-1 answer key.

BoltzGen ranks its own designs by confidence, but confidence is not the
question we asked. The question is whether a design lands where GLP-1 lands.
3IOL gives us the native complex, so for each design we recompute the epitope —
which receptor residues the design touches — and compare it to the epitope the
native peptide makes.

    recovery  = fraction of the native epitope the design also contacts
    precision = fraction of the design's contacts that are in the native epitope

A design can be confident and wrong (binds elsewhere on the receptor); these two
numbers catch that, and BoltzGen's own iptm does not.

Runs inside the session, next to the pipeline output:

    python3 score.py --output out_100
"""

import argparse
import csv
import glob
import json
import os
import sys

from prep_target import AA3TO1, contact_counts

CIF_COLS = ("group_PDB", "label_comp_id", "type_symbol",
            "Cartn_x", "Cartn_y", "Cartn_z", "auth_seq_id", "auth_asym_id")


def parse_cif_atoms(path):
    """{chain: [(resseq, resname, x, y, z)]} of heavy atoms, by auth numbering."""
    cols, rows = [], []
    in_loop = False
    with open(path) as fh:
        for line in fh:
            if line.startswith("_atom_site."):
                cols.append(line.strip().split(".", 1)[1])
                in_loop = True
            elif in_loop and line.startswith(("ATOM", "HETATM")):
                rows.append(line.split())
            elif in_loop and cols and rows:
                break
    idx = {c: cols.index(c) for c in CIF_COLS if c in cols}
    chains = {}
    for r in rows:
        if r[idx["type_symbol"]] == "H":
            continue
        chains.setdefault(r[idx["auth_asym_id"]], []).append((
            int(r[idx["auth_seq_id"]]), r[idx["label_comp_id"]],
            float(r[idx["Cartn_x"]]), float(r[idx["Cartn_y"]]), float(r[idx["Cartn_z"]]),
        ))
    return chains


def sequence_of(atoms):
    seen = {}
    for resseq, resname, *_ in atoms:
        seen.setdefault(resseq, resname)
    return "".join(AA3TO1.get(seen[r], "X") for r in sorted(seen))


def split_chains(chains, target_seq):
    """Return (design_atoms, target_atoms, renumber_offset).

    BoltzGen relabels and renumbers chains in its output, so identify the target
    by its sequence rather than trusting the chain id, and recover the offset
    back to the original PDB numbering from where it matches.
    """
    for cid, atoms in chains.items():
        seq = sequence_of(atoms)
        if seq and seq in target_seq:
            others = [a for c, a in chains.items() if c != cid]
            if not others:
                raise SystemExit(f"only found the target chain in {cid}")
            offset = target_seq.index(seq) - min(a[0] for a in atoms)
            return max(others, key=len), atoms, offset
    raise SystemExit("no chain matched the target sequence")


def score_design(path, native):
    chains = parse_cif_atoms(path)
    design, target, offset = split_chains(chains, native["target_sequence"])
    # offset maps output numbering -> index into target_sequence; add the first
    # resolved residue number to land back in original PDB numbering.
    base = offset + native["target_first_resnum"]
    contacts = contact_counts(target, design, native["cutoff"])
    epitope = {r + base for r in contacts}
    ref = set(native["epitope"])
    hit = epitope & ref
    return {
        "design": os.path.basename(path),
        "sequence": sequence_of(design),
        "length": len(sequence_of(design)),
        "epitope": sorted(epitope),
        "recovery": len(hit) / len(ref) if ref else 0.0,
        "precision": len(hit) / len(epitope) if epitope else 0.0,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--output", default="out_100", help="BoltzGen --output directory")
    ap.add_argument("--native", default="native.json")
    ap.add_argument("--out", default=None, help="where to write scores.json")
    args = ap.parse_args()

    with open(args.native) as fh:
        native = json.load(fh)

    ranked = os.path.join(args.output, "final_ranked_designs")
    cifs = sorted(glob.glob(os.path.join(ranked, "final_*_designs", "rank*.cif")))
    if not cifs:
        raise SystemExit(f"no ranked designs under {ranked}")

    metrics = {}
    for csv_path in glob.glob(os.path.join(ranked, "final_designs_metrics_*.csv")):
        with open(csv_path) as fh:
            for row in csv.DictReader(fh):
                metrics[row["designed_sequence"]] = row

    scored = []
    for path in cifs:
        s = score_design(path, native)
        m = metrics.get(s["sequence"], {})
        s["iptm"] = float(m["design_to_target_iptm"]) if m.get("design_to_target_iptm") else None
        s["helix"] = float(m["helix"]) if m.get("helix") else None
        scored.append(s)

    scored.sort(key=lambda s: s["recovery"], reverse=True)
    out_path = args.out or os.path.join(args.output, "scores.json")
    with open(out_path, "w") as fh:
        json.dump({"native": native["native_sequence"],
                   "native_epitope": native["epitope"],
                   "designs": scored}, fh, indent=2)

    print(f"native GLP-1 ({native['native_length']} aa): {native['native_sequence']}")
    print(f"native epitope: {len(native['epitope'])} receptor residues\n")
    print(f"{'design':22} {'len':>3} {'iptm':>5} {'helix':>5} {'recov':>6} {'prec':>6}  sequence")
    for s in scored:
        iptm = f"{s['iptm']:.2f}" if s["iptm"] is not None else "  -  "
        helix = f"{s['helix']:.2f}" if s["helix"] is not None else "  -  "
        print(f"{s['design'][:22]:22} {s['length']:3d} {iptm:>5} {helix:>5} "
              f"{s['recovery']:6.2f} {s['precision']:6.2f}  {s['sequence']}")
    best = scored[0]
    print(f"\nbest epitope recovery: {best['recovery']:.0%} ({best['design']})")
    print(f"wrote {out_path}")


def demo():
    """Self-check: a design placed on the native epitope scores recovery 1.0."""
    native = {"target_sequence": "ACDEFGHIKL", "target_first_resnum": 29,
              "cutoff": 5.0, "epitope": [30, 31]}
    # target residues 1..3 in output numbering -> 29..31 in PDB numbering.
    target = [(1, "ALA", 0.0, 0, 0), (2, "CYS", 4.0, 0, 0), (3, "ASP", 8.0, 0, 0)]
    design = [(1, "GLY", 6.0, 0, 0)]  # within 5 A of target residues 2 and 3
    chains = {"B": target, "A": design}
    d, t, off = split_chains(chains, native["target_sequence"])
    assert len(d) == 1 and len(t) == 3, (d, t)
    assert off == -1, off  # seq starts at index 0, output numbering starts at 1
    contacts = contact_counts(t, d, native["cutoff"])
    epitope = {r + off + native["target_first_resnum"] for r in contacts}
    assert epitope == {30, 31}, epitope
    print("demo ok")


if __name__ == "__main__":
    if "--demo" in sys.argv:
        demo()
    else:
        main()

{
    "sku": "gpu_1x_h100",
    "image": {
        "base": "pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel",
        "pip": [
            "boltzgen"
        ],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "HF_HOME": "/workspace/.cache/huggingface"
        }
    },
    "volumes": [
        {"name": "boltzgen-weights", "mount": "/workspace/.cache"}
    ]
}