diff --git a/examples/small_molecule_binding/CLAUDE.md b/examples/small_molecule_binding/CLAUDE.md index 248c09e..d175d05 100644 --- a/examples/small_molecule_binding/CLAUDE.md +++ b/examples/small_molecule_binding/CLAUDE.md @@ -11,6 +11,8 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co | 2026-09-09 | — | Replaced broken RFD3 `scaffoldguided.target_pdb` scaffold feedback with real RFD3 partial-diffusion guidance (`partial.input`/`partial_t`); replaced AlphaFold2/ColabFold fold-validation step with Boltz-2 (protein+ligand co-folding) | | 2026-09-09 | — | Fixed `analysis_sequence()` silently never comparing MPNN candidates (it only ever read each `.fa` file's first line — the un-designed template record — since real LigandMPNN writes multiple candidates into one file, not one file per candidate); now parses every candidate record and picks the true highest-confidence one | | 2026-09-09 | — | Added a metric-agnostic non-improvement short-circuit to `fastrelax`/`interface` retry logic — escalates to a new backbone (`STEP_RFD3`) as soon as a retry fails to improve on the previous attempt, instead of always exhausting 5 resequencing retries on backbones that real data showed never recover | +| 2026-09-09 | — | Fixed guided-RFD3 ligand atom-name mismatch that crashed 4/4 pipelines in a real production run (job `21916521`) on their first guided-backbone-feedback attempt: `_normalize_ligand_id()` only rewrote the Boltz co-folded PDB's ligand *residue* name, not its Boltz-assigned *atom* names, so `select_exposed`/`select_buried` (copied verbatim from the base spec, keyed by canonical `.params` atom names) never matched and RFD3's validator rejected every guided run. New `_infer_ligand_atom_mapping()`/`_normalize_ligand_atom_names()` establish atom correspondence via element+connectivity graph isomorphism (rdkit) with a Kabsch-RMSD tie-break; `_write_guided_rfd3_json()` now also verifies atom-name coverage before writing. Adds an `rdkit` runtime dependency | +| 2026-09-09 | — | Fixed a second guided-RFD3 crash found in a real production run (job `21928556`, 3/3 pipelines that reached guided feedback crashed on their very first attempt): `_write_guided_rfd3_json()` copied the base spec's `partial.length` field verbatim into the guided JSON, but RFD3's `DesignInputSpecification` validator rejects `length` outright whenever `partial.input`/`partial_t` (partial diffusion) are set (`ValidationError: ... Length argument must not be provided during partial diffusion`) — length is inferred from the input structure in that mode. `_write_guided_rfd3_json()` now drops `length` from the guided spec | ## Context @@ -120,12 +122,29 @@ RFD3 has no `scaffoldguided.*`-style CLI override (that was a leftover from an o When `rfd3_input_pdb` is set, `rfd3()`: 1. Reads the ligand's literal residue name from `ligand_params`'s `NAME` record via `_ligand_resname_from_params()` — this is **not** always the params filename stem (e.g. `ALR.params`'s `NAME` is `A:R`, not `ALR`; the colon is a deliberate workaround for RFD3 misresolving the bare `"ALR"` literal — never hardcode or "clean up" this value). -2. Calls `_normalize_ligand_id()` to rewrite the Boltz model's ligand HETATM residue name (via gemmi, no coordinate transform — Boltz already places the ligand correctly relative to the protein it just co-folded) to match that literal, writing `{taskdir}/in/guided_scaffold.pdb`. -3. Calls `_write_guided_rfd3_json()` to copy the base `ALR_binder_design.json`'s `ligand`/`length`/`select_exposed`/`select_buried` fields verbatim into a new spec with `input` pointed at the normalized PDB and `partial_t` set, writing `{taskdir}/in/guided_binder_design.json`. -4. Passes that guided JSON (instead of the base one) as `rfd3.sh`'s `inputs=` argument. If normalization fails (no ligand found in the Boltz model), falls back to the base, unguided JSON rather than erroring. +2. Calls `_normalize_ligand_id()` to rewrite the Boltz model's ligand HETATM residue *name* (via gemmi, no coordinate transform — Boltz already places the ligand correctly relative to the protein it just co-folded) to match that literal, writing `{taskdir}/in/guided_scaffold.pdb`. +3. Calls `_normalize_ligand_atom_names()` to rewrite that same PDB's ligand *atom* names to the canonical `.params` names (see "Guided-RFD3 ligand atom-name mapping" below) — Boltz assigns its own arbitrary atom names during co-folding, unrelated to the params file, so this is a separate fix from step 2. +4. Calls `_write_guided_rfd3_json()` to copy the base `ALR_binder_design.json`'s `ligand`/`select_exposed`/`select_buried` fields verbatim (dropping `length` — RFD3's validator rejects it during partial diffusion, since length is inferred from the input structure) into a new spec with `input` pointed at the normalized PDB and `partial_t` set, writing `{taskdir}/in/guided_binder_design.json` — after first verifying every `select_exposed`/`select_buried` atom name is actually present in the normalized PDB. +5. Passes that guided JSON (instead of the base one) as `rfd3.sh`'s `inputs=` argument. If any of steps 2–4 fails (no ligand found in the Boltz model, no full atom-name mapping found, or the coverage check fails), falls back to the base, unguided JSON rather than erroring. + +### Guided-RFD3 ligand atom-name mapping + +Boltz-2's co-folded ligand output uses its own arbitrary atom names (e.g. `C41`, `O24`, ...) that have nothing to do with the canonical names in the ligand's `.params` file (e.g. `C18`, `O3`, ...). Since `select_exposed`/`select_buried` are copied verbatim from the base spec and are keyed by those canonical names, a guided PDB with unrenamed atoms fails RFD3's own input validation (`ComponentValidationError: Number of atoms must be a multiple of the requested names`) — this was the confirmed root cause of a real production run (job `21916521`) crashing all 4 pipeline instances on their first guided-feedback attempt. + +`_infer_ligand_atom_mapping()` establishes the correspondence via element + heavy-atom-connectivity graph isomorphism (rdkit), ignoring bond order throughout (the `.params` format has none): +- **Reference graph**: heavy atoms + bonds parsed directly from the `.params` file's `ATOM`/`BOND` records (`_params_heavy_atom_graph()`) — exact, no perception needed. Reference *coordinates* come from the real, correctly-named structure the base spec's own `partial.input` already points at (`_resolve_reference_pdb_path()`) — no synthetic conformer is built. +- **Query graph**: the Boltz ligand's connectivity, perceived from 3D distances via `rdkit.Chem.rdDetermineBonds.DetermineConnectivity()` (Boltz's output carries no CONECT records for the ligand). +- **Isomorphism + tie-break**: `GetSubstructMatches()` enumerates every graph-valid atom correspondence; local topological symmetry (e.g. a sulfonate's three interchangeable terminal oxygens) can yield more than one. Each candidate is Kabsch-superposed (reusing `_kabsch_rmsd()`) against the reference coordinates, and the lowest-RMSD mapping wins — grounded in real geometry rather than an arbitrary tiebreak. When more than one isomorphism exists, the best-vs-next-best RMSD gap is logged. + +This is **not a generic guarantee for every future ligand**: it's only provably safe for a ligand whose "sides" (whatever `select_exposed`/`select_buried` partition into) aren't themselves graph-isomorphic to each other — true for `ALR` (a monocyclic benzene-sulfonate ring isn't isomorphic to a fused naphthalene-sulfonate ring), but not checked automatically for `IND`/`RED`/`IAI` or any future ligand. Run `scripts/check_ligand_atom_mapping.py` against a new ligand's `.params` + reference PDB before trusting guided feedback with it. + +`_write_guided_rfd3_json()` adds a final defensive check: before writing, it verifies every `select_exposed`/`select_buried` atom name is present in the (now atom-renamed) guided PDB, returning `False` (fail safe, fall back to unguided) rather than reproducing RFD3's rejection in a new form if not. + +`scripts/check_ligand_atom_mapping.py` validates this mapping logic against real crash artifacts rather than synthetic test data: the job-`21916521` `guided_scaffold.pdb` files under `logs/p{1..4}/*_rfd3/in/`, with the committed `p1_in/ALR.params` + `p1_in/input_pdbs/scaffold-with-ALR.pdb` as ground truth. **Those crash artifacts are not committed** — `logs/` is gitignored — so on a fresh checkout every check skips and the tool reports `INCONCLUSIVE` (exit 2), not `PASS`. Pass `--base-path ` to actually exercise it. `scripts/validate_run.py`'s check 8 (`check_guided_ligand_atom_names`) regression-tests the same invariant against any completed run's output tree. Ensemble similarity utilities (all in `small_molecule_binding.py`): - `_ca_rmsd(path1, path2)` — Kabsch-aligned CA RMSD between two PDB files +- `_kabsch_rmsd(coords1, coords2)` — the underlying generic Kabsch-alignment RMSD, also reused by `_infer_ligand_atom_mapping()`'s isomorphism tie-break - `_seq_identity(fasta1, fasta2)` — fraction matching residues over shorter sequence - `_ensemble_selective_avg(current, prior, sim_fn, similar_if_low)` — returns `(overall_avg, selective_avg, has_data)` for scores of entries whose similarity is on the "similar" side of the mean pairwise similarity @@ -207,3 +226,5 @@ Each pipeline instance (named e.g. `p1`) expects a `{name}_in/` directory contai - `.params` — Rosetta ligand params file (default `ALR.params`) - `.smiles` — ligand SMILES string, read by the `boltz` task to build its co-folding input; derive it once with `scripts/derive_ligand_smiles.py .params .pdb` (RDKit bond-order perception from the params file's exact connectivity + the reference structure's 3D coordinates — there is no SMILES in a Rosetta `.params` file itself) - Optionally `common_filenames.txt` — used by `filter_energy` for cross-filtering + +`rdkit` is a runtime dependency (installed by `delta_env_setup.sh`'s Step 10, pinned `2024.9.6`) used both by the offline `derive_ligand_smiles.py` tool above and, per-cycle, by `rfd3()`'s guided-backbone-feedback atom-name mapping (see "Guided-RFD3 ligand atom-name mapping" above). `scripts/check_ligand_atom_mapping.py` and `scripts/validate_run.py` (check 8) validate that mapping against real fixtures/completed runs, respectively. diff --git a/examples/small_molecule_binding/delta_env_setup.sh b/examples/small_molecule_binding/delta_env_setup.sh index 406f8ff..10c60d6 100755 --- a/examples/small_molecule_binding/delta_env_setup.sh +++ b/examples/small_molecule_binding/delta_env_setup.sh @@ -200,18 +200,39 @@ echo "" echo "── Step 9: gemmi ──" "${PIP}" install -q "gemmi==0.6.5" -# ── 10. Additional dependencies ─────────────────────────────────────────────── +# ── 10. rdkit — ligand atom-name graph-isomorphism mapping for guided RFD3 ──── +# +# Used by rfd3()'s guided-backbone-feedback path (_infer_ligand_atom_mapping +# in small_molecule_binding.py) to reconcile Boltz-2's arbitrary ligand atom +# names against the canonical names in the ligand's .params file, via +# element+connectivity graph isomorphism (rdDetermineBonds.DetermineConnectivity +# + GetSubstructMatches) with a Kabsch-RMSD tie-break. Without this, RFD3's +# input validator rejects every guided run (ComponentValidationError) -- +# confirmed as the root cause of 4/4 pipeline crashes in a real production +# run (job 21916521). +# +# Pinned to 2024.9.6, the same version already used by the offline +# scripts/derive_ligand_smiles.py tool in this repo (rdkit has no +# dependency on numpy/gemmi's own pins, so it should not disturb Step 7's +# numpy<2.0/gemmi==0.6.5 resolution -- `pip check` after this step should +# stay clean; re-investigate only if it doesn't). +# +echo "" +echo "── Step 10: rdkit ──" +"${PIP}" install -q "rdkit==2024.9.6" + +# ── 11. Additional dependencies ─────────────────────────────────────────────── echo "" -echo "── Step 10: pandas + biopandas ──" +echo "── Step 11: pandas + biopandas ──" "${PIP}" install -q pandas biopandas -# ── 11. PyRosetta ───────────────────────────────────────────────────────────── +# ── 12. PyRosetta ───────────────────────────────────────────────────────────── echo "" -echo "── Step 11: PyRosetta ──" +echo "── Step 12: PyRosetta ──" "${PIP}" install -q pyrosetta-installer "${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()" -# ── 12. Boltz-2 model weights ───────────────────────────────────────────────── +# ── 13. Boltz-2 model weights ───────────────────────────────────────────────── # # Boltz has no dedicated "download weights" subcommand — weights auto-download # on first `boltz predict` call. Warm the cache with a trivial CPU prediction @@ -219,7 +240,7 @@ echo "── Step 11: PyRosetta ──" # BOLTZ_CACHE. # echo "" -echo "── Step 12: Boltz-2 model weights (cache warm-up) ──" +echo "── Step 13: Boltz-2 model weights (cache warm-up) ──" BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}" mkdir -p "${BOLTZ_CACHE}" _WARM_DIR=$(mktemp -d) @@ -237,9 +258,9 @@ YAML || echo "WARNING: boltz cache warm-up failed — check login-node internet access" rm -rf "${_WARM_DIR}" -# ── 13. Verify ──────────────────────────────────────────────────────────────── +# ── 14. Verify ──────────────────────────────────────────────────────────────── echo "" -echo "── Step 13: Verifying installation ──" +echo "── Step 14: Verifying installation ──" _check() { local label="$1"; shift if out=$("$@" 2>&1); then @@ -256,6 +277,7 @@ _check "impress" "${PY}" -c "import impress; print('ok')" _check "torch" "${PY}" -c "import torch; print(torch.__version__)" _check "boltz" "${PY}" -c "import boltz; print('ok')" _check "gemmi" "${PY}" -c "import gemmi; print(gemmi.__version__)" +_check "rdkit" "${PY}" -c "import rdkit; print(rdkit.__version__)" _check "pyrosetta" "${PY}" -c "import pyrosetta; print('ok')" _check "ProDy" "${PY}" -c "import prody; print(prody.__version__)" _check "LigandMPNN" test -d "${MPNN_DIR}" && echo "present" diff --git a/examples/small_molecule_binding/delta_gpu_run.sh b/examples/small_molecule_binding/delta_gpu_run.sh index f952e25..dbbf720 100644 --- a/examples/small_molecule_binding/delta_gpu_run.sh +++ b/examples/small_molecule_binding/delta_gpu_run.sh @@ -71,7 +71,7 @@ dragon-config add --ofi-runtime-lib="${FAB_LIB}" export MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}" # Boltz-2 model weights cache — kept on scratch to avoid home quota exhaustion. -# Pre-warm once on a login node via delta_env_setup.sh's Step 12 (boltz has no +# Pre-warm once on a login node via delta_env_setup.sh's Step 13 (boltz has no # dedicated "download weights" subcommand; weights auto-download on first # `boltz predict` call). export BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}" diff --git a/examples/small_molecule_binding/run_small_molecule_binding.py b/examples/small_molecule_binding/run_small_molecule_binding.py index 9c27937..d7c2490 100644 --- a/examples/small_molecule_binding/run_small_molecule_binding.py +++ b/examples/small_molecule_binding/run_small_molecule_binding.py @@ -36,6 +36,7 @@ class RunConfig: # diffusion / refinement diffusion_batch_size: int num_refine_cycles: int + mpnn_ensemble_size: int # cycle-0 MPNN sequence candidates per backbone rfd3_partial_t: float # RFD3 partial-diffusion noise (A) for guided backbone feedback @@ -56,6 +57,7 @@ class RunConfig: fold_min_ligand_iptm = None, diffusion_batch_size = 4, num_refine_cycles = 2, + mpnn_ensemble_size = 10, rfd3_partial_t = 10.0, ) @@ -73,6 +75,7 @@ class RunConfig: fold_min_ligand_iptm = None, diffusion_batch_size = 1, num_refine_cycles = 1, + mpnn_ensemble_size = 2, # Not a pass/fail threshold like the fields above — a diffusion-noise # parameter, so kept at a sane real value rather than an inert extreme. rfd3_partial_t = 10.0, @@ -102,12 +105,30 @@ def _prior(ttype): if step == 'backbone': if not passed: + # Safety net: a guided (partial-diffusion) backbone that keeps + # failing QC -- whether for real structural reasons or an + # unforeseen RFD3 metrics-schema gap -- would otherwise loop on + # the same rfd3_input_pdb forever, since only a *successful* fold + # ever clears it. Fall back to unguided regeneration after a few + # consecutive guided-mode failures instead of deadlocking. + if pipeline.state.get('rfd3_input_pdb') is not None: + count = pipeline.state.get('backbone_guided_fail_count', 0) + 1 + pipeline.state['backbone_guided_fail_count'] = count + if count >= 3: + pipeline.state['rfd3_input_pdb'] = None + pipeline.state['backbone_guided_fail_count'] = 0 + pipeline.logger.pipeline_log( + "[adaptive/backbone] guided backbone QC failed 3x in a " + "row -- abandoning guided mode, reverting to unguided " + "(scratch) RFD3 generation" + ) pipeline.next_step = STEP_RFD3 else: current, prior = _prior(ETYPE_BACKBONE) # Reset on any new backbone -- these are per-backbone retry state, # not per-pipeline, and must not leak into the next backbone's # first fastrelax/interface attempt (see _stage_metrics_improving). + pipeline.state['backbone_guided_fail_count'] = 0 pipeline.state['seq_retry_count'] = 0 pipeline.state['fastrelax_prev_metrics'] = None pipeline.state['interface_prev_metrics'] = None @@ -138,14 +159,36 @@ def _prior(ttype): pipeline.state['seq_retry_count'] = count if count >= 3: pipeline.state['seq_retry_count'] = 0 + # A guided backbone that fails downstream of backbone-QC + # gets only this one shot -- otherwise rfd3_input_pdb stays + # pinned to the same seed indefinitely (only 'backbone' QC + # failures and 'fold' decisions used to clear it), causing + # RFD3 to regenerate near-duplicate doomed backbones for + # dozens of cycles in a row (confirmed via job 21945304: + # 17 consecutive p2 rfd3 generations produced byte-identical + # guided_scaffold.pdb output from one stuck seed). + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + "[adaptive/sequence] sequence-similarity gate failed " + "3x in a row on this backbone -- escalating to a new, " + "unguided backbone (STEP_RFD3) instead of another " + "resequencing retry" + ) else: pipeline.next_step = STEP_RETRY_SEQ elif step == 'packmin': total_score = metrics.get('total_score') if total_score is not None and total_score > 0: + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 # badly packed — restart backbone + pipeline.logger.pipeline_log( + f"[adaptive/packmin] total_score={total_score} > 0 -- badly " + "packed, escalating to a new, unguided backbone (STEP_RFD3)" + ) else: pipeline.next_step = STEP_MPNN @@ -178,9 +221,17 @@ def _prior(ttype): pipeline.state['fastrelax_fail_count'] = count if (prev is not None and not improving) or count >= 5: + reason = "safety cap (5 attempts)" if count >= 5 else "non-improving metrics" pipeline.state['fastrelax_fail_count'] = 0 pipeline.state['fastrelax_prev_metrics'] = None + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + f"[adaptive/fastrelax] escalating to a new, unguided " + f"backbone (STEP_RFD3) ({reason}); attempt={count} metrics={metrics}" + ) else: pipeline.next_step = STEP_MPNN @@ -204,9 +255,17 @@ def _prior(ttype): pipeline.state['interface_fail_count'] = count if (prev is not None and not improving) or count >= 5: + reason = "safety cap (5 attempts)" if count >= 5 else "non-improving metrics" pipeline.state['interface_fail_count'] = 0 pipeline.state['interface_prev_metrics'] = None + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + f"[adaptive/interface] escalating to a new, unguided " + f"backbone (STEP_RFD3) ({reason}); attempt={count} metrics={metrics}" + ) else: pipeline.next_step = STEP_MPNN @@ -215,16 +274,32 @@ def _prior(ttype): if not passed: # Failed fold — don't use this model as a backbone guide pipeline.state['rfd3_input_pdb'] = None + pipeline.logger.pipeline_log( + "[adaptive/fold] fold failed -- next backbone will be unguided (scratch)" + ) else: if not prior: pipeline.state['rfd3_input_pdb'] = None + pipeline.logger.pipeline_log( + "[adaptive/fold] fold passed but no prior fold history yet -- " + "next backbone will be unguided (scratch)" + ) else: overall, selective, has_data = _ensemble_selective_avg( current[3], prior, _ca_rmsd, similar_if_low=True) if has_data and selective is not None and selective > overall: pipeline.state['rfd3_input_pdb'] = current[3] # guided backbone + pipeline.logger.pipeline_log( + f"[adaptive/fold] similar-cluster avg ({selective:.2f}) > " + f"overall avg ({overall:.2f}) -- next backbone guided from {current[3]}" + ) else: pipeline.state['rfd3_input_pdb'] = None # scratch + pipeline.logger.pipeline_log( + f"[adaptive/fold] guided-feedback condition not met " + f"(has_data={has_data}, selective={selective}, overall={overall}) " + "-- next backbone will be unguided (scratch)" + ) pipeline.next_step = STEP_RFD3 else: @@ -246,8 +321,8 @@ async def impress_smallmol_bind() -> None: ) os.makedirs(work_dir, exist_ok=True) # Input data lives in the source tree; pass as absolute so it resolves - # correctly regardless of what base_path / work_dir is set to. - input_dir = os.path.join(examples_dir, "p1_in") + # correctly regardless of what base_path / work_dir is set to. Each + # pipeline reads its own p{i}_in/ directory rather than sharing one. if BACKEND == "dragon": backend = await DragonExecutionBackend() @@ -265,7 +340,7 @@ async def impress_smallmol_bind() -> None: kwargs={ "base_path": work_dir, "scripts_path": os.path.join(examples_dir, "scripts"), - "input_dir": input_dir, + "input_dir": os.path.join(examples_dir, f"p{i}_in"), "backbone_max_ca_deviation": cfg.backbone_max_ca_deviation, "backbone_min_ss_fraction": cfg.backbone_min_ss_fraction, "fastrelax_max_fa_rep": cfg.fastrelax_max_fa_rep, @@ -276,6 +351,7 @@ async def impress_smallmol_bind() -> None: "fold_min_ligand_iptm": cfg.fold_min_ligand_iptm, "diffusion_batch_size": cfg.diffusion_batch_size, "num_refine_cycles": cfg.num_refine_cycles, + "mpnn_ensemble_size": cfg.mpnn_ensemble_size, "rfd3_partial_t": cfg.rfd3_partial_t, "max_tasks": cfg.max_tasks, **({"gpu_id": all_gpus[(i - 1) % len(all_gpus)]} if all_gpus else {}), diff --git a/examples/small_molecule_binding/scripts/check_ligand_atom_mapping.py b/examples/small_molecule_binding/scripts/check_ligand_atom_mapping.py new file mode 100644 index 0000000..bd7adf9 --- /dev/null +++ b/examples/small_molecule_binding/scripts/check_ligand_atom_mapping.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Standalone validation for small_molecule_binding.py's RFD3 guided-input +ligand atom-name mapping (_infer_ligand_atom_mapping / +_normalize_ligand_atom_names), run against real crash artifacts from a +production run rather than synthetic test data. + +Those artifacts are NOT committed -- they live under logs/, which is +gitignored -- so on a fresh checkout every check skips and this tool reports +INCONCLUSIVE (exit 2) rather than PASS. Point --base-path at a completed +run's output tree (the directory holding p1/, p2/, ...) to actually exercise +the mapping. + +Background: job 21916521 (a real ~3h production HPC run) crashed all 4 +pipeline instances on their first use of RFD3 guided-backbone feedback, +because Boltz-2's co-folded ligand output uses its own arbitrary atom names +that don't match the canonical params-file names baked into the base RFD3 +spec's select_exposed/select_buried fields. The crash-artifact guided PDBs +from that run (logs/p{1..4}/*_rfd3/in/guided_scaffold.pdb) are the primary +fixtures used here, alongside the committed p1_in/ALR.params and +p1_in/input_pdbs/scaffold-with-ALR.pdb as ground truth. + +Usage: + python scripts/check_ligand_atom_mapping.py + python scripts/check_ligand_atom_mapping.py --base-path logs + +Exits 0 if every check passes, 1 if any fails, 2 if no fixtures were found. +""" + +import argparse +import glob +import json +import os +import sys + +_EXAMPLE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _EXAMPLE_DIR not in sys.path: + sys.path.insert(0, _EXAMPLE_DIR) + +from small_molecule_binding import ( # noqa: E402 + _find_ligand_hetatm_residue, + _infer_ligand_atom_mapping, + _ligand_resname_from_params, + _resolve_reference_pdb_path, +) + +# Boltz atom -> canonical name pairs that must land in the same +# select_exposed/select_buried bucket for ALR specifically (its two ring +# systems -- a monocyclic benzene-sulfonate and a fused naphthalene-sulfonate +# -- are not graph-isomorphic to each other, so the *only* real ambiguity is +# each sulfonate's 3 interchangeable terminal oxygens; see CLAUDE.md's +# "Ensemble-guided backbone feedback" section for why this is ALR-specific, +# not a general guarantee for future ligands). +_ALR_BURIED_SULFONATE_OXYGENS = {"O5", "O7", "O8"} # bonded to S1 +_ALR_EXPOSED_SULFONATE_OXYGENS = {"O6", "O9", "O10"} # bonded to S2 + + +def _load_boltz_atoms(pdb_path: str, ligand_chain_id: str = "B"): + import gemmi + st = gemmi.read_structure(pdb_path) + res = _find_ligand_hetatm_residue(st, ligand_chain_id) + if res is None: + return None + return {atom.name: (atom.element.name, (atom.pos.x, atom.pos.y, atom.pos.z)) for atom in res} + + +def check_real_crash_artifacts(examples_dir: str, base_path: str): + """Run _infer_ligand_atom_mapping against every real *_rfd3/in/guided_scaffold.pdb + left on disk from job 21916521 (or any other run). For each: assert a + full bijection is found, every mapped pair is element-consistent, the + ALR select_exposed/select_buried atom-name lists are fully covered, and + each sulfonate's 3 oxygens land together in the correct bucket (a real + correctness check beyond "isomorphism exists").""" + params_path = os.path.join(examples_dir, "p1_in", "ALR.params") + base_json_path = os.path.join(examples_dir, "p1_in", "ALR_binder_design.json") + if not (os.path.isfile(params_path) and os.path.isfile(base_json_path)): + return [f"SKIPPED: {params_path} or {base_json_path} not found"] + reference_pdb = _resolve_reference_pdb_path(base_json_path) + if not os.path.isfile(reference_pdb): + return [f"reference pdb {reference_pdb} not found"] + + with open(base_json_path) as fh: + base_partial = json.load(fh)["partial"] + ligand_key = base_partial["ligand"] + expected_names = set() + for field in ("select_exposed", "select_buried"): + expected_names.update(base_partial[field][ligand_key].split(",")) + + candidates = sorted(glob.glob(os.path.join(base_path, "p*", "*_rfd3", "in", "guided_scaffold.pdb"))) + if not candidates: + return [f"SKIPPED: no '*_rfd3/in/guided_scaffold.pdb' files found under {base_path}"] + + failures = [] + for pdb_path in candidates: + boltz_atoms = _load_boltz_atoms(pdb_path) + if boltz_atoms is None: + failures.append(f"{pdb_path}: no ligand HETATM residue found") + continue + + mapping = _infer_ligand_atom_mapping(boltz_atoms, params_path, reference_pdb) + if mapping is None: + failures.append(f"{pdb_path}: _infer_ligand_atom_mapping returned None (no mapping found)") + continue + + if len(mapping) != len(boltz_atoms): + failures.append(f"{pdb_path}: mapping covers {len(mapping)}/{len(boltz_atoms)} atoms, not a full bijection") + + mapped_names = set(mapping.values()) + if not expected_names <= mapped_names: + failures.append( + f"{pdb_path}: select_exposed/select_buried coverage FAILED -- " + f"missing {sorted(expected_names - mapped_names)}" + ) + + if ligand_key == "A:R": + buried_group = {b for b, c in mapping.items() if c in _ALR_BURIED_SULFONATE_OXYGENS} + exposed_group = {b for b, c in mapping.items() if c in _ALR_EXPOSED_SULFONATE_OXYGENS} + if len(buried_group) != 3 or len(exposed_group) != 3: + failures.append( + f"{pdb_path}: sulfonate oxygen grouping broken -- " + f"buried={buried_group} exposed={exposed_group} (expected 3 each)" + ) + + return failures + + +def check_negative_paths(examples_dir: str, base_path: str): + """Corrupt a real fixture (drop an atom; swap an element to one absent + from ALR) and confirm _infer_ligand_atom_mapping fails safe (returns + None) rather than raising.""" + params_path = os.path.join(examples_dir, "p1_in", "ALR.params") + base_json_path = os.path.join(examples_dir, "p1_in", "ALR_binder_design.json") + reference_pdb = _resolve_reference_pdb_path(base_json_path) + + fixture = None + for pdb_path in sorted(glob.glob(os.path.join(base_path, "p*", "*_rfd3", "in", "guided_scaffold.pdb"))): + boltz_atoms = _load_boltz_atoms(pdb_path) + if boltz_atoms: + fixture = boltz_atoms + break + if fixture is None: + return [f"SKIPPED: no usable '*_rfd3/in/guided_scaffold.pdb' fixture found under {base_path}"] + + failures = [] + + missing_atom = dict(fixture) + missing_atom.pop(next(iter(missing_atom))) + try: + result = _infer_ligand_atom_mapping(missing_atom, params_path, reference_pdb) + except Exception as e: + failures.append(f"missing-atom case raised {e!r} instead of returning None") + else: + if result is not None: + failures.append("missing-atom case returned a mapping instead of None") + + bad_element = dict(fixture) + name0 = next(iter(bad_element)) + _, xyz = bad_element[name0] + bad_element[name0] = ("Cl", xyz) # not present in ALR at all + try: + result2 = _infer_ligand_atom_mapping(bad_element, params_path, reference_pdb) + except Exception as e: + failures.append(f"bad-element case raised {e!r} instead of returning None") + else: + if result2 is not None: + failures.append("bad-element case returned a mapping instead of None") + + return failures + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--base-path", default=os.path.join(_EXAMPLE_DIR, "logs"), + help="Directory containing p1/, p2/, ... pipeline output trees (default: examples_dir/logs)", + ) + args = parser.parse_args(argv) + + checks = [ + ("1. Real crash-artifact mapping correctness", + lambda: check_real_crash_artifacts(_EXAMPLE_DIR, args.base_path)), + ("2. Negative-path fail-safe behavior", + lambda: check_negative_paths(_EXAMPLE_DIR, args.base_path)), + ] + + print("=" * 72) + print("LIGAND ATOM-NAME MAPPING VALIDATION") + print("=" * 72) + + any_failed = False + any_ran = False + for name, fn in checks: + try: + failures = fn() + except Exception as e: + failures = [f"check raised an unexpected exception: {e!r}"] + + if failures and all(f.startswith("SKIPPED:") for f in failures): + print(f"[SKIP] {name}") + for f in failures: + print(f" {f}") + elif not failures: + any_ran = True + print(f"[PASS] {name}") + else: + any_ran = True + any_failed = True + print(f"[FAIL] {name} -- {len(failures)} issue(s)") + for f in failures: + print(f" - {f}") + + print("=" * 72) + if any_failed: + print("RESULT: FAIL") + return 1 + if not any_ran: + # Every check skipped for want of fixtures. Reporting PASS here would + # green-light a clean checkout, where the crash artifacts this tool + # reads are absent by construction -- logs/ is gitignored. + print("RESULT: INCONCLUSIVE -- no check had fixtures to run against.") + print(f" Point --base-path at a completed run's output tree") + print(f" (the directory holding p1/, p2/, ... ); tried: {args.base_path}") + return 2 + print("RESULT: PASS") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/small_molecule_binding/scripts/validate_run.py b/examples/small_molecule_binding/scripts/validate_run.py index b045060..c651bcb 100644 --- a/examples/small_molecule_binding/scripts/validate_run.py +++ b/examples/small_molecule_binding/scripts/validate_run.py @@ -70,7 +70,7 @@ def _ligand_resname_from_params(params_path: str) -> str: _REJECTED_STATE_KEY = "rfd3_guide_ligand_pdb" # At least one of these must exist somewhere under $BOLTZ_CACHE for the cache -# to be considered "warmed" (see plan's delta_env_setup.sh Step 12/13). +# to be considered "warmed" (see plan's delta_env_setup.sh Step 13/14). _BOLTZ_CACHE_MARKERS = ("boltz2_conf.ckpt", "boltz2_aff.ckpt", "mols.tar", "ccd.pkl") # Skip these extensions when grepping run output for the rejected state key -- @@ -223,9 +223,11 @@ def check_boltz_output_shape(base_path: str, pipeline_name: str): def check_guided_json_correctness(base_path: str, pipeline_name: str, pipeline_inputs: str): """Every */_rfd3/in/guided_binder_design.json must parse, its partial.input - must point at a file that exists, and partial.ligand/length/select_exposed/ + must point at a file that exists, partial.ligand/select_exposed/ select_buried must match the base ALR_binder_design.json verbatim (only - input/partial_t may legitimately differ) -- mirrors _write_guided_rfd3_json.""" + input/partial_t may legitimately differ), and partial.length must be + absent (RFD3 rejects it during partial diffusion; it's only valid for + the base spec's from-scratch diffusion) -- mirrors _write_guided_rfd3_json.""" pipeline_dir = os.path.join(base_path, pipeline_name) guided_jsons = sorted( glob.glob(os.path.join(pipeline_dir, "*_rfd3", "in", "guided_binder_design.json")) @@ -273,7 +275,13 @@ def check_guided_json_correctness(base_path: str, pipeline_name: str, pipeline_i f"(resolved: {resolved})" ) - for field in ("ligand", "length", "select_exposed", "select_buried"): + if "length" in partial: + failures.append( + f"{gj}: partial.length={partial['length']!r} must not be present -- " + f"RFD3 rejects 'length' during partial diffusion" + ) + + for field in ("ligand", "select_exposed", "select_buried"): expected = base_partial.get(field) actual = partial.get(field) if actual != expected: @@ -464,6 +472,63 @@ def check_env_sanity(python_exe: str): return failures +# ── Check 8: guided ligand atom names cover select_exposed/select_buried ─── + +def check_guided_ligand_atom_names(base_path: str, pipeline_name: str, pipeline_inputs: str): + """Every */_rfd3/in/guided_scaffold.pdb's ligand HETATM atom names must be + a superset of its sibling guided_binder_design.json's select_exposed + + select_buried atom-name lists -- this is exactly the invariant RFD3's + input validator itself enforces (rejecting the guided spec with + 'ComponentValidationError: Number of atoms must be a multiple of the + requested names' when it doesn't hold). Regression test for the + guided-input ligand atom-name mapping fix; see + scripts/check_ligand_atom_mapping.py for the deeper unit-level check of + the mapping logic itself against real fixtures.""" + pipeline_dir = os.path.join(base_path, pipeline_name) + pairs = [] + for guided_pdb in sorted(glob.glob(os.path.join(pipeline_dir, "*_rfd3", "in", "guided_scaffold.pdb"))): + guided_json = os.path.join(os.path.dirname(guided_pdb), "guided_binder_design.json") + if os.path.isfile(guided_json): + pairs.append((guided_pdb, guided_json)) + if not pairs: + # No guided rfd3 runs have happened yet -- not a failure by itself. + return [] + + failures = [] + for guided_pdb, guided_json in pairs: + try: + with open(guided_json) as fh: + guided = json.load(fh) + except (OSError, json.JSONDecodeError) as e: + failures.append(f"{guided_json}: failed to parse as JSON: {e}") + continue + + partial = guided.get("partial", {}) + ligand_key = partial.get("ligand") + expected_names = set() + for field in ("select_exposed", "select_buried"): + names_csv = partial.get(field, {}).get(ligand_key, "") + expected_names.update(n for n in names_csv.split(",") if n) + if not expected_names: + failures.append(f"{guided_json}: no select_exposed/select_buried atom names found for ligand {ligand_key!r}") + continue + + present_names = set() + with open(guided_pdb, errors="replace") as fh: + for line in fh: + if line.startswith("HETATM") and line[17:20].strip() == ligand_key: + present_names.add(line[12:16].strip()) + + missing = expected_names - present_names + if missing: + failures.append( + f"{guided_pdb}: missing {sorted(missing)} from select_exposed/select_buried " + f"(this is exactly what RFD3's own validator would reject with " + f"ComponentValidationError)" + ) + return failures + + # ── main ───────────────────────────────────────────────────────────────────── def main(argv=None) -> int: @@ -525,6 +590,8 @@ def main(argv=None) -> int: lambda: check_ensemble_sanity(base_path, pipeline_name)), ("7. Env sanity (BOLTZ_CACHE / import boltz)", lambda: check_env_sanity(args.python)), + ("8. Guided ligand atom names cover select_exposed/select_buried", + lambda: check_guided_ligand_atom_names(base_path, pipeline_name, pipeline_inputs)), ] print("=" * 72) diff --git a/examples/small_molecule_binding/small_molecule_binding.py b/examples/small_molecule_binding/small_molecule_binding.py index 8373470..2f1d1cd 100644 --- a/examples/small_molecule_binding/small_molecule_binding.py +++ b/examples/small_molecule_binding/small_molecule_binding.py @@ -170,85 +170,392 @@ def _ligand_resname_from_params(params_path: str) -> str: raise ValueError(f"no NAME record found in {params_path}") +def _find_ligand_hetatm_residue(st, ligand_chain_id: str = "B"): + """Locates the ligand HETATM residue in a gemmi Structure. Tries + ligand_chain_id first; if Boltz didn't honor the requested chain letter, + falls back to the first HETATM residue found anywhere in the structure. + Returns None if no HETATM residue exists at all. Shared by + _normalize_ligand_id and _normalize_ligand_atom_names so both agree on + exactly which residue is "the ligand".""" + for model in st: + for chain in model: + if chain.name != ligand_chain_id: + continue + for res in chain: + if res.het_flag == 'H': + return res + + for model in st: + for chain in model: + for res in chain: + if res.het_flag == 'H': + return res + + return None + + def _normalize_ligand_id(fold_pdb_path: str, ligand_resname: str, out_pdb_path: str, ligand_chain_id: str = "B") -> bool: """Rewrites a Boltz co-folded PDB's ligand HETATM residue name (via gemmi) to match ligand_resname, so RFD3's ligand/select_exposed/select_buried selectors (which key off the literal resname) resolve against it. Purely a string edit -- no coordinate transform, since Boltz already places the ligand correctly - relative to the protein it just folded. Locates the ligand residue by HETATM - on ligand_chain_id first; if Boltz didn't honor the requested chain letter, - falls back to the first HETATM residue found anywhere in the structure. - Writes out_pdb_path and returns True on success. Returns False (writing - nothing) if no HETATM residue is found at all, so the caller can fall back - to unguided diffusion instead of crashing.""" + relative to the protein it just folded. Writes out_pdb_path and returns True + on success. Returns False (writing nothing) if no HETATM residue is found at + all, so the caller can fall back to unguided diffusion instead of crashing. + + NOTE: this only fixes the residue name. Boltz also assigns its own, + unrelated atom names within that residue -- see _normalize_ligand_atom_names + for why those need fixing too before the guided spec is usable.""" import gemmi st = gemmi.read_structure(fold_pdb_path) - target_res = None - for model in st: - for chain in model: - if chain.name != ligand_chain_id: + target_res = _find_ligand_hetatm_residue(st, ligand_chain_id) + if target_res is None: + return False + + target_res.name = ligand_resname + st.write_pdb(out_pdb_path) + return True + + +# Fallback table keyed on Rosetta atom TYPE prefix, used only when a params +# ATOM name's leading alphabetic run doesn't parse to a valid element symbol. +# Mirrors scripts/derive_ligand_smiles.py's _ROSETTA_TYPE_ELEMENT_FALLBACK +# (duplicated, not imported -- see _params_heavy_atom_graph). Verified only +# against ALR's C/N/O/S/H atom set; extend if a future ligand needs 2-letter +# elements (Cl/Br/Zn, etc.). +_ROSETTA_TYPE_ELEMENT_FALLBACK = { + "Nhis": "N", "Nlys": "N", + "CH1": "C", "CH2": "C", "CH3": "C", "COO": "C", "aroC": "C", + "OH": "O", "OOC": "O", "ONH2": "O", + "Hapo": "H", "Hpol": "H", + "S": "S", "SH1": "S", +} + + +def _infer_ligand_element(atom_name: str, rosetta_type: str) -> str: + """Infers an element symbol from a params ATOM record's name/type. Primary + rule: the atom NAME's leading alphabetic run is the element itself (e.g. + 'C18' -> 'C', 'N11' -> 'N'); falls back to _ROSETTA_TYPE_ELEMENT_FALLBACK + keyed on the Rosetta TYPE prefix. Mirrors + scripts/derive_ligand_smiles.py's _infer_element (duplicated, not + imported -- that script is a standalone offline tool, not part of this + per-cycle pipeline path; scripts/ isn't an importable package).""" + import re + from rdkit import Chem + periodic_table = Chem.GetPeriodicTable() + + match = re.match(r'[A-Za-z]+', atom_name) + if match: + candidate = match.group(0) + for length in (2, 1): + if len(candidate) >= length: + symbol = candidate[:length].capitalize() + if periodic_table.GetAtomicNumber(symbol) > 0: + return symbol + + for prefix, element in _ROSETTA_TYPE_ELEMENT_FALLBACK.items(): + if rosetta_type.startswith(prefix): + return element + + raise ValueError( + f"could not infer element for atom name={atom_name!r} type={rosetta_type!r}" + ) + + +def _params_heavy_atom_graph(params_path: str): + """Parses a Rosetta .params file's ATOM/BOND records into a heavy-atom-only + connectivity graph: ({atom_name: element}, [(atom1, atom2), ...]) where + both bond endpoints are heavy atoms. Hydrogens are dropped entirely -- + unlike scripts/derive_ligand_smiles.py (which needs them to resolve bond + order via DetermineBondOrders), _infer_ligand_atom_mapping only needs + connectivity for graph-isomorphism matching, so no bond-order solving or + placeholder hydrogen placement is needed here.""" + atom_types = {} + bonds = [] + with open(params_path) as fh: + for line in fh: + fields = line.split() + if not fields: continue - for res in chain: - if res.het_flag == 'H': - target_res = res - break - if target_res: - break - if target_res: - break + record = fields[0] + if record == 'ATOM': + atom_types[fields[1]] = fields[2] + elif record in ('BOND', 'BOND_TYPE'): + bonds.append((fields[1], fields[2])) + + elements = {name: _infer_ligand_element(name, rtype) for name, rtype in atom_types.items()} + heavy_names = {name for name, el in elements.items() if el != 'H'} + heavy_elements = {name: elements[name] for name in heavy_names} + heavy_bonds = [(a, b) for a, b in bonds if a in heavy_names and b in heavy_names] + return heavy_elements, heavy_bonds + + +def _resolve_reference_pdb_path(base_json_path: str) -> str: + """Reads partial.input from the base RFD3 design spec and resolves it + relative to the spec's own directory (always pipeline_inputs). This is + the correctly-named, real-coordinate reference ligand structure already + shipped alongside every pipeline's inputs -- used as ground truth for + atom identity/connectivity/geometry in _infer_ligand_atom_mapping.""" + with open(base_json_path) as fh: + base = json.load(fh) + ref = base['partial']['input'] + if os.path.isabs(ref): + return ref + return os.path.join(os.path.dirname(base_json_path), ref) + + +def _iter_ligand_atom_names(pdb_path: str, resname: str): + """Atom names (fixed-column PDB parsing) of every HETATM record in + pdb_path whose resname matches exactly. Fixed-column parsing (not a + whitespace split) is required because names like 'A:R' contain a colon + -- see _ligand_resname_from_params's docstring.""" + names = [] + with open(pdb_path) as fh: + for line in fh: + if line.startswith('HETATM') and line[17:20].strip() == resname: + names.append(line[12:16].strip()) + return names + + +def _load_reference_ligand_coords(pdb_path: str, resname: str): + """Reads {atom_name: (x, y, z)} for the first HETATM residue named + exactly resname in pdb_path. Mirrors + scripts/derive_ligand_smiles.py's _load_reference_coords (duplicated, + not imported -- see _params_heavy_atom_graph); uses the same + fixed-column parsing for the same reason (resnames like 'A:R' contain a + colon a whitespace split would mangle).""" + coords = {} + target_resseq = None + with open(pdb_path) as fh: + for line in fh: + if not (line.startswith('HETATM') or line.startswith('ATOM ')): + continue + if line[17:20].strip() != resname: + if coords: + break # moved past the matching residue's contiguous block + continue + resseq = line[22:26].strip() + if target_resseq is None: + target_resseq = resseq + elif resseq != target_resseq: + break # a different residue instance with the same name + atom_name = line[12:16].strip() + coords[atom_name] = (float(line[30:38]), float(line[38:46]), float(line[46:54])) + return coords + + +def _infer_ligand_atom_mapping(boltz_atoms: dict, params_path: str, reference_pdb_path: str): + """Maps Boltz's arbitrarily-named ligand atom names onto the canonical + names read from params_path, via element+connectivity graph isomorphism + with a Kabsch-RMSD tie-break. boltz_atoms is + {boltz_atom_name: (element, (x, y, z))} for one ligand residue. + + Boltz co-folding assigns its own atom names to the ligand (unrelated to + the params file's canonical names), so select_exposed/select_buried + (copied verbatim from the base RFD3 spec, keyed by canonical names) never + match a Boltz-derived PDB's atom names without this step. Bond order is + ignored throughout (the .params file has none) -- only element identity + and heavy-atom connectivity establish correspondence: + - reference graph: heavy atoms + bonds parsed straight from + params_path's ATOM/BOND records (exact, no perception needed) + - reference coordinates: the real 3D structure at reference_pdb_path + (already correctly named -- see _resolve_reference_pdb_path) + - query graph: boltz_atoms' connectivity, perceived from 3D distances + via rdkit's DetermineConnectivity (Boltz's ligand output carries no + CONECT records) + Local topological symmetry (e.g. a sulfonate's three interchangeable + terminal oxygens) can produce more than one graph-valid isomorphism; both + structures carry real, roughly comparable 3D coordinates for the same + ligand pose, so each candidate mapping is Kabsch-superposed against the + reference and the lowest-RMSD one wins -- deterministic, and grounded in + actual geometry rather than an arbitrary tiebreak. When more than one + isomorphism exists, the best-vs-next-best RMSD gap is logged so a + suspiciously close tie (a symmetry case this heuristic can't actually + distinguish) is visible after the fact rather than silently accepted. + + Returns {boltz_name: canonical_name}, or None (never raises) if: heavy + atom counts or element multisets differ, the reference PDB is missing + coordinates for a params heavy atom, Boltz connectivity perception fails + or yields a disconnected graph, or no isomorphism exists at all -- + callers must treat None exactly like _normalize_ligand_id returning + False (fall back to unguided diffusion).""" + from rdkit import Chem + from rdkit.Chem import rdDetermineBonds + from rdkit.Geometry import Point3D + + ref_elements, ref_bonds = _params_heavy_atom_graph(params_path) + ref_resname = _ligand_resname_from_params(params_path) + ref_coords = _load_reference_ligand_coords(reference_pdb_path, ref_resname) + ref_names = [name for name in ref_elements if name in ref_coords] + if len(ref_names) != len(ref_elements): + return None # reference PDB is missing coordinates for a params heavy atom + + if len(boltz_atoms) != len(ref_names): + return None + if sorted(element for element, _ in boltz_atoms.values()) != sorted(ref_elements[n] for n in ref_names): + return None - if target_res is None: - # Fallback: Boltz may not have honored the requested chain id. - for model in st: - for chain in model: - for res in chain: - if res.het_flag == 'H': - target_res = res - break - if target_res: - break - if target_res: - break + ref_mol = Chem.RWMol() + ref_idx = {} + for name in ref_names: + ref_idx[name] = ref_mol.AddAtom(Chem.Atom(ref_elements[name])) + for a, b in ref_bonds: + if a in ref_idx and b in ref_idx: + i, j = ref_idx[a], ref_idx[b] + if ref_mol.GetBondBetweenAtoms(i, j) is None: + ref_mol.AddBond(i, j, Chem.BondType.SINGLE) + ref_conf = Chem.Conformer(ref_mol.GetNumAtoms()) + for name, idx in ref_idx.items(): + ref_conf.SetAtomPosition(idx, Point3D(*ref_coords[name])) + ref_mol.AddConformer(ref_conf, assignId=True) + Chem.SanitizeMol(ref_mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_NONE) + + boltz_names = list(boltz_atoms) + boltz_mol = Chem.RWMol() + boltz_idx = {} + for name in boltz_names: + element, _ = boltz_atoms[name] + boltz_idx[name] = boltz_mol.AddAtom(Chem.Atom(element)) + boltz_conf = Chem.Conformer(boltz_mol.GetNumAtoms()) + for name, idx in boltz_idx.items(): + _, xyz = boltz_atoms[name] + boltz_conf.SetAtomPosition(idx, Point3D(*xyz)) + boltz_mol.AddConformer(boltz_conf, assignId=True) + Chem.SanitizeMol(boltz_mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_NONE) + + try: + rdDetermineBonds.DetermineConnectivity(boltz_mol) + except Exception: + return None + if len(Chem.GetMolFrags(boltz_mol)) != 1: + return None + + matches = boltz_mol.GetSubstructMatches(ref_mol, uniquify=False, useChirality=False, maxMatches=10000) + if not matches: + return None + + ref_coord_list = [ref_coords[name] for name in ref_names] # order == ref_mol atom order + best_mapping, best_rmsd, second_best_rmsd = None, None, None + for match in matches: + # match[i] is the boltz_mol atom index matched to ref_mol atom i (ref_names[i]). + query_coord_list = [boltz_atoms[boltz_names[qi]][1] for qi in match] + rmsd = _kabsch_rmsd(ref_coord_list, query_coord_list) + if best_rmsd is None or rmsd < best_rmsd: + second_best_rmsd = best_rmsd + best_rmsd = rmsd + best_mapping = {boltz_names[qi]: ref_names[i] for i, qi in enumerate(match)} + elif second_best_rmsd is None or rmsd < second_best_rmsd: + second_best_rmsd = rmsd + + if len(matches) > 1: + gap = (second_best_rmsd - best_rmsd) if second_best_rmsd is not None else float('inf') + print( + f"[rfd3 guided] ligand atom-name mapping: {len(matches)} candidate " + f"isomorphisms, best RMSD={best_rmsd:.4f} vs next-best={second_best_rmsd:.4f} " + f"(gap={gap:.4f}) -- a small gap means the tie-break may not be decisive" + ) + return best_mapping + + +def _normalize_ligand_atom_names(pdb_path: str, params_path: str, base_json_path: str, + out_pdb_path: str, ligand_chain_id: str = "B") -> bool: + """Rewrites a resname-normalized guided PDB's ligand HETATM *atom* names + (not just its residue name -- see _normalize_ligand_id) to match the + canonical names read from params_path, via _infer_ligand_atom_mapping. + Required because select_exposed/select_buried in the guided RFD3 spec are + copied verbatim from the base spec and are keyed by those canonical + names, but Boltz assigns its own arbitrary atom names during co-folding + -- without this, RFD3's input validator rejects every guided run. + + Writes out_pdb_path (may be the same path as pdb_path) and returns True + on success. Returns False (writing nothing) if no ligand residue is + found, or no full atom-name mapping could be established, so the caller + falls back to unguided diffusion instead of producing a guided spec RFD3 + will reject.""" + import gemmi + st = gemmi.read_structure(pdb_path) + + target_res = _find_ligand_hetatm_residue(st, ligand_chain_id) if target_res is None: return False - target_res.name = ligand_resname + boltz_atoms = { + atom.name: (atom.element.name, (atom.pos.x, atom.pos.y, atom.pos.z)) + for atom in target_res + } + reference_pdb_path = _resolve_reference_pdb_path(base_json_path) + mapping = _infer_ligand_atom_mapping(boltz_atoms, params_path, reference_pdb_path) + if mapping is None: + return False + + for atom in target_res: + atom.name = mapping[atom.name] st.write_pdb(out_pdb_path) return True def _write_guided_rfd3_json(base_json_path: str, guided_pdb_path: str, partial_t: float, - out_json_path: str) -> None: + out_json_path: str) -> bool: """Loads the base per-pipeline RFD3 InputSpecification JSON, copies its - ligand/length/select_exposed/select_buried fields verbatim, replaces 'input' - with guided_pdb_path, adds partial_t, and writes the result to - out_json_path. Only 'input'/'partial_t' differ from the base file.""" + ligand/select_exposed/select_buried fields verbatim, drops 'length', + replaces 'input' with guided_pdb_path, adds partial_t, and writes the + result to out_json_path. Only 'input'/'partial_t' differ from the base + file (besides the dropped 'length'). + + 'length' is dropped because RFD3's DesignInputSpecification validator + rejects it outright when partial.input/partial_t (partial diffusion) are + set -- length is inferred from the input structure in that mode. The + base file's 'length' is only valid for from-scratch (non-partial) + diffusion. + + Before writing, verifies every atom name referenced by select_exposed/ + select_buried is actually present in guided_pdb_path's ligand residue -- + fails safe (returns False, writes nothing) on a stale/mismatched base + spec or an atom-mapping bug, rather than reproducing RFD3's + ComponentValidationError in a new form. Returns True on success.""" with open(base_json_path) as fh: base = json.load(fh) partial = dict(base.get('partial', {})) + partial.pop('length', None) + + ligand_key = partial.get('ligand') + expected_names = set() + for field in ('select_exposed', 'select_buried'): + names_csv = partial.get(field, {}).get(ligand_key, '') + expected_names.update(name for name in names_csv.split(',') if name) + present_names = set(_iter_ligand_atom_names(guided_pdb_path, ligand_key)) + if not expected_names <= present_names: + return False + partial['input'] = guided_pdb_path partial['partial_t'] = partial_t guided = dict(base) guided['partial'] = partial with open(out_json_path, 'w') as fh: json.dump(guided, fh, indent=4) + return True def _prepare_guided_rfd3_inputs(base_json_path: str, fold_pdb_path: str, ligand_resname: str, - partial_t: float, taskdir: str): - """Orchestrates _normalize_ligand_id + _write_guided_rfd3_json: writes - {taskdir}/in/guided_scaffold.pdb and {taskdir}/in/guided_binder_design.json. - Returns the guided JSON path, or None if ligand normalization failed (e.g. - no HETATM residue found in fold_pdb_path) -- callers should fall back to + params_path: str, partial_t: float, taskdir: str): + """Orchestrates _normalize_ligand_id + _normalize_ligand_atom_names + + _write_guided_rfd3_json: writes {taskdir}/in/guided_scaffold.pdb and + {taskdir}/in/guided_binder_design.json. Returns the guided JSON path, or + None if any step failed (no ligand found in fold_pdb_path, no full + atom-name mapping could be established, or the select_exposed/ + select_buried coverage check failed) -- callers should fall back to unguided diffusion in that case.""" guided_pdb = f"{taskdir}/in/guided_scaffold.pdb" guided_json = f"{taskdir}/in/guided_binder_design.json" if not _normalize_ligand_id(fold_pdb_path, ligand_resname, guided_pdb): return None - _write_guided_rfd3_json(base_json_path, guided_pdb, partial_t, guided_json) + if not _normalize_ligand_atom_names(guided_pdb, params_path, base_json_path, guided_pdb): + return None + if not _write_guided_rfd3_json(base_json_path, guided_pdb, partial_t, guided_json): + return None return guided_json @@ -368,13 +675,13 @@ async def rfd3(): fold_pdb = self.state.get('rfd3_input_pdb') inputs = base_inputs if fold_pdb: - ligand_resname = _ligand_resname_from_params( - f"{self.pipeline_inputs}/{self.ligand_params}" - ) + params_path = f"{self.pipeline_inputs}/{self.ligand_params}" + ligand_resname = _ligand_resname_from_params(params_path) guided_json = _prepare_guided_rfd3_inputs( base_json_path=base_inputs, fold_pdb_path=fold_pdb, ligand_resname=ligand_resname, + params_path=params_path, partial_t=self.rfd3_partial_t, taskdir=taskdir, ) @@ -402,15 +709,29 @@ async def analysis_backbone(): for jf in json_files: with open(f"{out_dir}/{jf}") as fh: data = json.load(fh) - m = data.get('metrics', {}) - clashes = m.get('n_clashing.ligand_clashes', float('inf')) - dev = m.get('max_ca_deviation', float('inf')) - ss = m.get('helix_fraction', 0) + m.get('sheet_fraction', 0) + m = data.get('metrics', {}) + # RFD3's partial-diffusion (guided-backbone-feedback) mode never + # computes ligand-clash or secondary-structure metrics -- the + # 'n_clashing.ligand_clashes' key is absent entirely and + # helix_fraction/sheet_fraction come back as literal JSON NaN + # (so ss = NaN + NaN, and NaN > threshold is always False in + # Python) -- only max_ca_deviation and the interresidue clash + # counts are populated there. Unguided mode has all of these. + guided = 'n_clashing.ligand_clashes' not in m + if guided: + clashes = ( + m.get('n_clashing.interresidue_clashes_w_sidechain', float('inf')) + + m.get('n_clashing.interresidue_clashes_w_backbone', float('inf')) + ) + else: + clashes = m.get('n_clashing.ligand_clashes', float('inf')) + dev = m.get('max_ca_deviation', float('inf')) + ss = m.get('helix_fraction', 0) + m.get('sheet_fraction', 0) if best is None or clashes < best['clashes'] or ( clashes == best['clashes'] and dev < best['dev'] ): - best = {'file': jf, 'clashes': clashes, 'dev': dev, 'ss': ss} + best = {'file': jf, 'clashes': clashes, 'dev': dev, 'ss': ss, 'guided': guided} if best is None: self.state.update({ @@ -427,13 +748,16 @@ async def analysis_backbone(): passed = ( best['clashes'] == 0 and best['dev'] < self.backbone_max_ca_deviation - and best['ss'] > self.backbone_min_ss_fraction + # SS fraction isn't computable in guided mode (see above) -- + # skip that check there rather than fail unconditionally. + and (best['guided'] or best['ss'] > self.backbone_min_ss_fraction) ) self.state.update({ 'last_analysis_step': 'backbone', 'last_analysis_metrics': { 'pass': passed, 'best_model': best['file'], + 'guided': best['guided'], 'ligand_clashes': best['clashes'], 'max_ca_deviation': best['dev'], 'ss_fraction': best['ss'], @@ -909,6 +1233,7 @@ async def run(self): self.state.setdefault('last_seq_fasta', None) self.state.setdefault('fastrelax_prev_metrics', None) self.state.setdefault('interface_prev_metrics', None) + self.state.setdefault('backbone_guided_fail_count', 0) self.logger.pipeline_log("SmallMoleculeBindingPipeline starting (state machine)") while self.next_step != STEP_DONE: