Python API tutorial: Create an Annotated CIF for a pocket conditioned design#

🚀 Try it live — no local setup needed
Launch on Brev

Goal. Start from a protein and ligand complex, and build one model ready annotated CIF using the Python API that asks RFD4-Proteina to design a new protein around the ligand. A handful of steps gets you an input with:

  • a flexible ligand with some atoms anchored,

  • one pocket residue fully fixed,

  • two more pocket residues with just their tip atoms coordinate conditioned,

  • H-bond donor and acceptor roles marked on the ligand’s polar atoms,

  • the CIF written out, plus a small sweep of variants generated in code.

No JSON. No contig strings. Just AtomWorks annotations on an atom array. See the Python API reference for more information about what you can do with the interface. See the annotations reference for the full list of annotations.

New to this?

This tutorial teaches you how to hand-build a minimal, fixed-atom motif from a real structure using Python. An annotated CIF is what is actually used as an input to RFD4-Proteina. It’s a normal CIF file with extra annotations that are used to condition the design. The Python API interfance is used to specify what we want those conditions to be.

The enzyme design tutorial starts from the same kind of minimal fixed-atom motif you built here, a theozyme, and takes it through scaffold generation, inverse folding, and partial diffusion.

New to this? What’s the example structure this tutorial actually uses?

The default input is a real, public structure: the second bromodomain of a human protein called PHIP, bound to its ligand — a small drug-discovery “fragment,” meaning a tiny candidate molecule caught during an early-stage screening experiment, not a finished drug. It comes from PLINDER, a benchmark dataset of real protein-ligand structures built for training machine learning models like this one. A bromodomain is a protein module that recognizes a specific chemical tag (acetylation) on other proteins, most often on histones (the spool-like proteins DNA winds around) and this recognition step helps control which nearby genes get switched on. PHIP itself has been linked to cancer and diabetes progression, which is part of why its bromodomain is of interest as a drug target.

[ ]:

from pathlib import Path

import os
import logging
logging.getLogger("atomworks.constants").setLevel(logging.ERROR)

# 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():          # e.g. run outside the checkout
    _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 has a `.project-root` marker), or set PROJECT_ROOT "
            "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"))

import numpy as np
import py3Dmol
from atomworks.io import parse
from atomworks.io.utils.annotator import ensure_annotations
from atomworks.io.utils.io_utils import to_cif_file
from atomworks.io.utils.visualize import view
from atomworks.ml.conditions import C_CRD, C_IDX, C_SEQ
from atomworks.ml.utils.expandable_segment import insert_expandable_segment

from rfproteina.conditions import C_HB, C_IDXGRP, HBondEnum
from rfproteina.utils.parse_presets import ANNOTATED_CIF, RCSB_WITH_METALS

# Point INPUT_CIF at your own protein-ligand complex. Default: a repo fixture.
INPUT_CIF = REPO_ROOT / "data/benchmarks/tip_atom_plinder_training_examples/5eni__1__1.A__1.B.cif"
OUT_DIR = REPO_ROOT / "runs/annotated_cif_tutorial"
OUT_DIR.mkdir(parents=True, exist_ok=True)

1. Parse and subselect to the pocket#

[ ]:

parse_config = RCSB_WITH_METALS.replace(hydrogen_policy="remove")

aa_full = parse(
    str(INPUT_CIF),
    config=parse_config,
)["assemblies"]["1"][0]

v = view(
    aa_full,
    show_surface=False,                  # skip the built-in faint surface
    polymer_sidechain_linewidth=0.001,   # effectively hides the side-chain wireframe
    show_cartoon=True,                   # keep the ribbon so you can still see the fold
)
v.addSurface(py3Dmol.VDW, {"opacity": 0.85, "color": "white"}, {"chain": "1.A"})
v.show()

Parse the complex, then keep only the ligand and the three pocket residues nearest to it. Here we fully fix TYR 41, and later pin one tip atom each on VAL 36 and THR 87.

[ ]:

# Pocket residues nearest the ligand (chain, res_id), found offline.
scaffold_residue_ids = [41, 36, 87]
TIPS = {("1.A", 36): "CG1", ("1.A", 87): "OG1"}  # pin one sidechain tip atom each

parse_config = RCSB_WITH_METALS.replace(hydrogen_policy="remove")

aa = parse(
    str(INPUT_CIF),
    config=parse_config,  # see from rfproteina.utils.parse_presets for common configs you might want for parse
)["assemblies"]["1"][0]
ensure_annotations(
    aa, "is_ligand"
)  # check out atomworks.io.utils.annotator and atomworks.ml.utils.annotator for easily addable annotations!

keep_mask = aa.is_ligand | ((aa.chain_id == "1.A") & np.isin(aa.res_id, scaffold_residue_ids))
aa = aa[keep_mask]

print(aa.array_length(), "atoms |", int(aa.is_ligand.sum()), "ligand atoms")

[ ]:

# Render inline: the ligand (sticks) surrounded by the three pocket residues.
view(aa)

New to this? What’s a “pocket,” and what do the chain/residue labels mean?

A pocket is the small cavity on a protein’s surface where a ligand physically sits — formed by whichever residues happen to be close enough to touch it. A chain is one distinct molecule within the file (here, 1.A is the protein, 1.B is the ligand). A residue ID is just that residue’s position number within its chain: residue 41 means the 41st position in this protein’s numbering, not the 41st atom.

2. Add an expandable segment#

[ ]:

aa = insert_expandable_segment(
    aa,
    seg_min=150,
    seg_max=200,
    new_chain=False,
    atom_idx=0,
    insert_before=True,
)

New to this? What’s an “expandable segment”?

After step 1, the kept residues are disconnected fragments — not a real, physically continuous protein. An expandable segment is a placeholder reserving space for new residues (here, 150–200 of them) that will later connect everything into one unbroken chain. Its exact final length isn’t decided yet, only the allowed range is.

3. Apply the pocket conditioning#

Three conditions define the motif:

  • Coordinates, C_CRD, fix the ligand chlorine atoms, all of TYR 41, and the selected tip atoms.

  • Sequence, C_SEQ, reveals identity for every scaffolded residue and the ligand.

  • Index, C_IDX and C_IDXGRP, marks the three motif residues unindexed, so their position in the designed sequence is free, and puts each in its own index group so they place independently.

[ ]:

coord_cond_mask = (aa.chain_id == "1.A") & (aa.res_id == 41)  # residue 41: all atoms
for (chain, res_id), atom in TIPS.items():                     # plus one pinned tip atom each
    coord_cond_mask |= (aa.chain_id == chain) & (aa.res_id == res_id) & (aa.atom_name == atom)
scaffold_residues = np.isin(aa.res_id, scaffold_residue_ids) & (aa.chain_id == "1.A")

# Coordinates: ligand chlorine atoms + fully-fixed residue + the pinned tip atoms.
coord_cond_mask |= aa.is_ligand & (np.char.upper(aa.element) == "CL")
coordinate = C_CRD.default_annotation(aa)
coordinate[coord_cond_mask] = C_CRD.ground_truth_value(aa)[coord_cond_mask]
C_CRD.set_annotation(aa, coordinate)

# Sequence: condition ligand and scaffolded residues (not expandable segment)
sequence_mask = scaffold_residues | aa.is_ligand
sequence = C_SEQ.default_annotation(aa)
sequence[sequence_mask] = C_SEQ.ground_truth_value(aa)[sequence_mask]
C_SEQ.set_annotation(aa, sequence)

# Index: unindex all three motif residues, then ensure each residue has it's own index group.
C_IDX.set_annotation(aa, ~scaffold_residues)

# The unindexed residues have non-contiguous residue IDs, so the default IDXGROUP would already
# put each in its own group. We reset it explicitly here anyway, just to show how - this step
# isn't strictly necessary given the default behavior.
groups = C_IDXGRP.annotation(aa, default="generate")
groups[aa.res_id == 36] = max(groups) + 1
groups[aa.res_id == 87] = max(groups) + 1
assert (
    len(np.unique(groups[scaffold_residues])) == 3
), "Expected three unique index groups for the three scaffold residues."
assert (
    len(np.unique(groups[~scaffold_residues])) == 2
), "Expected two unique index groups for the ligand and expandable segment."
C_IDXGRP.set_annotation(aa, groups)  # Again, not strictly necessary here since we aren't changing from default

print("coord-fixed atoms:", int(C_CRD.mask(aa).sum()), "| unindexed atoms:", (~C_IDX.mask(aa)).sum())

New to this? What do C_CRD, C_SEQ, C_IDX, and C_IDXGRP mean?

C_CRD stores target coordinates (NaN where free), and C_SEQ stores target residue identities (MASKED where free). Their masks are derived from those values. C_IDX records whether each residue has a fixed sequence position; C_IDXGRP groups residues whose relative sequence offsets must stay together.

These conditions don’t have to apply the same way to every residue. Residue 41 gets C_CRD applied to all of its atoms so it is fully fixed. Residues 36 and 87, by contrast, only get C_CRD applied to one atom each, via the TIPS dictionary so just the single side-chain atom that matters chemically, leaving the rest of those two residues free. That’s exactly why residue 41 never appears in TIPS: TIPS only covers the weaker case, pinning one atom, not the whole residue.

4. Hydrogen bond roles on the ligand#

Request donor and acceptor roles on the ligand’s polar atoms, so the designed pocket satisfies them. 5QA carries an amide and a hydroxyl. The amide N donates, the carbonyl O1 accepts, and the hydroxyl O accepts. Use HBondEnum.BOTH to request both donor and acceptor roles on an atom.

[ ]:

hbond = C_HB.default_annotation(aa).copy()
hbond[aa.is_ligand & (aa.atom_name == "N")] = HBondEnum.DONOR  # amide N
hbond[aa.is_ligand & (aa.atom_name == "O1")] = HBondEnum.ACCEPTOR  # carbonyl O
hbond[aa.is_ligand & (aa.atom_name == "O")] = HBondEnum.ACCEPTOR  # hydroxyl O
C_HB.set_annotation(aa, hbond.astype(int))
print("H-bond atoms set:", int(C_HB.mask(aa).sum()))

5. Write the annotated CIF#

Write with save_standard_annotations=True, which embeds every condition, then re-parse with the ANNOTATED_CIF preset. That is the exact path the model uses, so it confirms every condition survives the round trip.

[ ]:

out_path = OUT_DIR / "annotated_design.cif"
# save_standard_annotations defaults to True; shown explicitly since this is what saves all
# conditions to the file.
to_cif_file(
    aa,
    str(out_path),
    save_standard_annotations=True,
    chain_disambiguation="chain_iid",
)

# if you want to load again, you can do so like this:
reloaded = parse(
    str(out_path),
    config=ANNOTATED_CIF,  # this config ensures that all conditions are loaded from the file
)["assemblies"]["1"][0]

6. Emit many variants programmatically#

Coordinate conditioning stores explicit target coordinates. Moving an atom does not update its saved target: change the corresponding C_CRD value as well. Here we push the THR 87 tip atom, OG1, directly away from the ligand in 0.1 Ă… steps, writing one CIF per step. That gives you a batch probing how far the pocket tolerates that contact. Point rfd4-generate at the directory to run them all.

[ ]:

sweep = np.where((aa.chain_id == "1.A") & (aa.res_id == 87) & (aa.atom_name == "OG1"))[0][0]
away = aa.coord[sweep] - aa.coord[aa.is_ligand].mean(axis=0)
away /= np.linalg.norm(away)  # unit vector pointing from the ligand toward the tip atom

variants_dir = OUT_DIR / "variants"
variants_dir.mkdir(exist_ok=True)
base_xyz = aa.coord[sweep].copy()
for i in range(10):
    variant = aa.copy()
    variant.coord[sweep] = base_xyz + away * 0.1 * (i + 1)
    targets = C_CRD.annotation(variant).copy()
    targets[sweep] = variant.coord[sweep]
    C_CRD.set_annotation(variant, targets)
    to_cif_file(
        variant,
        str(variants_dir / f"variant_{i:02d}.cif"),
        save_standard_annotations=True,
        chain_disambiguation="chain_iid",
    )

New to this? What does “sweeping” a coordinate mean?

Sweeping means changing a saved coordinate target slightly and re-saving, over and over, to produce related inputs. Here we update both the displayed atom position and its C_CRD target, nudging one pinned atom further from the ligand each time while preserving all other targets.

7. Chekc your work: Visualize what you set#

Below is an interactive viewer over the annotated CIF, with a button for each thing this notebook fixed. Most buttons are AtomWorks selection queries, written as aa.mask("..."). The H-bond roles come through the condition object here, as C_HB.mask(check); the equivalent query is sa_mask('C_HB').

[ ]:

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 selection_gallery(aa, selections, width=760, height=520):
    """Interactive viewer: one button per selection. Each entry is ``(label, query_or_mask)`` --
    a string is evaluated with ``aa.mask(query)``; a boolean array is used as given (for
    conditions accessed directly, such as H-bond roles)."""
    import string
    aa = aa.copy()                                  # 3Dmol via PDB needs single-char chain ids
    _uniq = list(dict.fromkeys(np.asarray(aa.chain_id).astype(str).tolist()))
    _cmap = {c: string.ascii_uppercase[i % 26] for i, c in enumerate(_uniq)}
    aa.chain_id = np.array([_cmap[c] for c in np.asarray(aa.chain_id).astype(str)])
    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, sel in selections:
        if isinstance(sel, str):
            m = np.asarray(aa.mask(sel)); note = f'aa.mask("{sel}")'
        else:
            m = np.asarray(sel); note = 'C_HB.mask(aa)'
        atoms = [[str(chain[i]), int(resi[i]), str(aname[i])] for i in np.where(m)[0]]
        items.append({"label": label, "note": note, "n": int(m.sum()), "atoms": atoms})

    uid = "g" + uuid.uuid4().hex[:10]
    payload = ",\n".join(
        "{{label:{l},note:{q},n:{n},atoms:{a}}}".format(
            l=json.dumps(it["label"]), q=json.dumps(it["note"]),
            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 = structure &nbsp;|&nbsp; orange = ligand &nbsp;|&nbsp; green = selected atoms
  </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({{}}, {{stick: {{colorscheme: "grayCarbon", radius: 0.10}}}});
    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.18}}}});
    }}
    viewer.zoomTo(); viewer.render();
    document.getElementById("{uid}_note").innerHTML = s.note + " &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>
""")

[ ]:

check = parse(str(out_path), config=ANNOTATED_CIF)["assemblies"]["1"][0]
ensure_annotations(check, "is_ligand")
check = check[~np.isnan(check.coord).any(axis=1)]   # 3Dmol needs real coordinates

selection_gallery(check, [
    ("substrate",               "hetero"),
    ("coordinate-conditioned",  "sa_mask('C_CRD')"),
    ("sequence-fixed",          "sa_mask('C_SEQ')"),
    ("indexed residues",        "sa_mask('C_IDX')"),
    ("H-bond donors/acceptors", np.asarray(C_HB.mask(check))),
])

Summary#

In this tutorial, we learned how to hand-build an annotated CIF from a real protein-ligand structure: subselecting to the binding pocket, choosing exactly which atoms stay fixed versus which are left open to design, and marking hydrogen-bond requirements on the ligand. We also saw how to generate a batch of related variants by nudging a single conditioned atom, and how to verify the file, both visually and by round-tripping it through a save/reload, before spending compute on generation.

Download this notebook