fix(rfo): limit MPNN to one highest-likelihood mutation

models/rfo/README.md

@@ -7,8 +7,10 @@
 RFOptimization (RFO) improves an initial biomolecular design through repeated
 sequence optimization and structure prediction. Each cycle selects RF3
 gradient-guided sequence optimization with probability `backprop_fraction`,
-or Boltz-2 structure prediction otherwise. MPNN then redesigns the unfixed
-residues of chain A, and the redesigned complex seeds the next cycle.
+or Boltz-2 structure prediction otherwise. MPNN then makes one amino-acid
+substitution among the unfixed residues of chain A, choosing the non-native
+amino acid with the highest conditional probability. The resulting complex
+seeds the next cycle.
 
 ## Installation
 
@@ -96,8 +98,12 @@ RF3 optimizer.
 `inputs/` with the redesigned complexes, and `<name>_records.json` with
 per-cycle scores, timing and designed sequences. Each cycle's `mpnn/`
 directory contains `mpnn_design.pdb`, `mpnn_result.json`, the design mask in
-`mpnn_spec.json`, and upstream samples. RFO carries the **first** MPNN sample
-forward when `mpnn_num_seqs > 1`; it does not rank the samples.
+`mpnn_spec.json`, and probabilities in `samples/scores.json`. The result JSON
+records the selected residue, old/new amino acid, and probability.
+`mpnn_num_seqs` controls the number of decoding orders whose probabilities
+are averaged before choosing one substitution. `mpnn_temperature` is retained
+for config compatibility but is not used in maximum-likelihood selection.
+Fully fixed chains are copied unchanged.
 
 RF3's `modelhub_pred/` output name is retained for compatibility with existing
 analysis scripts. Independent AlphaFold3 filtering is outside this pipeline.

models/rfo/docs/installation.md

@@ -15,7 +15,13 @@ foundry install rf3
 export RF3_CKPT="$HOME/.foundry/checkpoints/rf3_foundry_01_24_latest_remapped.ckpt"
 ```
 
-The Foundry/RF3 environment uses AtomWorks 3 (`atomworks[ml]>=3.0.1,<4`).
+The `rfo` extra pins AtomWorks to **3.0.1**. RFO uses the RF3 checkpoint
+transform pipeline, whose covalent-modification arguments are incompatible
+with AtomWorks 3.0.4. This temporary pin preserves the existing pipeline
+behavior without editing other models' code. Rerun the installation command above
+to apply the pin to an existing environment; avoid upgrading AtomWorks
+independently in that environment.
+
 Keep Boltz and upstream MPNN in their separate environments; their NumPy
 requirements need not match Foundry's.
 
@@ -68,8 +74,9 @@ such as `AAAA` that could otherwise be interpreted as RNA.
 
 Clone [dauparas/LigandMPNN](https://github.com/dauparas/LigandMPNN). This
 repository includes both ProteinMPNN and LigandMPNN inference and checkpoints;
-RFO expects its `run.py` interface. The separate ProteinMPNN repository's
-`protein_mpnn_run.py` interface is not compatible with this adapter.
+RFO runs a scoring worker in that environment using upstream `model_utils.py`
+and `data_utils.py`. The separate ProteinMPNN repository is not compatible
+with this adapter.
 
 For example, on a Linux CUDA machine, following the upstream environment setup:
 
@@ -116,7 +123,11 @@ RFO designs only unfixed residues in chain A and maps positional masks onto
 actual PDB residue numbers and insertion codes. Standard amino acids and
 complete N/CA/C/O backbone atoms are required in the designed chain. The
 adapter preserves other chains and ligands from the predicted complex,
-removes obsolete side-chain atoms at mutations, and selects the first sample.
+and changes only the highest-probability non-native amino acid at one unfixed
+position. Only that residue loses obsolete side-chain atoms. Probabilities are
+conditioned on the backbone and the other native residues, averaged over
+`mpnn_num_seqs` decoding orders. `mpnn_temperature` is accepted for backwards
+compatibility but is not used by this deterministic selection rule.
 When all residues are fixed, it copies the complex without calling MPNN.
 
 ## 3. Install Boltz-2
@@ -183,11 +194,11 @@ A completed run writes a records JSON and the redesigned complex under
 `inputs/`. Missing MPNN installations/checkpoints fail before model inference;
 failed or incomplete runs return a nonzero exit status.
 
-The portable adapter, CLI and gradient tests can be run from the repository root:
+The portable adapter, scoring-worker, CLI and gradient tests can be run from the repository root:
 
 ```bash
 PYTHONPATH=models/rfo/src python -m pytest models/rfo/tests
 ```
 
-These tests use a stand-in external CLI and small differentiable tensors;
+These tests use stand-in upstream model modules and small differentiable tensors;
 they do not replace a GPU smoke test with real model weights.

models/rfo/src/rfo/cli.py

@@ -4,8 +4,8 @@
     rfo --config config.yaml
 
 Runs the optimization: each cycle routes to RF3 backprop or Boltz with
-probability ``backprop_fraction`` (default 0.5), then redesigns the binder
-chain with MPNN and feeds the result into the next cycle. A YAML config sets
+probability ``backprop_fraction`` (default 0.5), then applies one MPNN-guided
+amino-acid substitution and feeds the result into the next cycle. A YAML config sets
 the inputs and output directory; any field can be overridden on the command
 line. The driver launches separate processes for Boltz, RF3 and external
 MPNN, optionally through Apptainer.
@@ -100,8 +100,16 @@ def build_parser() -> argparse.ArgumentParser:
     p.add_argument(
         "--num-samples", type=int, help="Boltz diffusion samples per prediction."
     )
-    p.add_argument("--mpnn-num-seqs", type=int)
-    p.add_argument("--mpnn-temperature", type=float)
+    p.add_argument(
+        "--mpnn-num-seqs",
+        type=int,
+        help="MPNN decoding orders to average for single-mutation scoring.",
+    )
+    p.add_argument(
+        "--mpnn-temperature",
+        type=float,
+        help="Compatibility option; unused by maximum-likelihood mutation selection.",
+    )
     p.add_argument("--template-plddt-threshold", type=int)
     p.add_argument(
         "--cyclic", action="store_true", help="Designed chain is a cyclic peptide."

models/rfo/src/rfo/cycling/README.md

@@ -2,8 +2,9 @@
 
 Each cycle chooses RF3 gradient/MCMC optimization with probability
 `backprop_fraction` (default 0.5), or Boltz-2 prediction otherwise. Both
-branches then redesign unfixed chain-A residues with an external MPNN
-installation and feed the redesigned complex into the next cycle.
+branches then apply one highest-probability amino-acid substitution among
+unfixed chain-A residues using an external MPNN installation, and feed the
+resulting complex into the next cycle.
 
 ```text
 RF3 gradient optimization ─┐
@@ -41,9 +42,16 @@ including insertion codes. Partner chains are never redesigned.
 
 `MPNN_MODEL_TYPE=protein_mpnn` selects ProteinMPNN and
 `MPNN_MODEL_TYPE=ligand_mpnn` selects LigandMPNN in the upstream checkout.
-The first generated sample is carried forward; all samples are retained.
-Fully fixed binders bypass sampling. The adapter calls upstream `run.py`
-and does not import the bundled Foundry MPNN engine.
+For each eligible position, MPNN scores amino acids conditioned on the
+backbone and the rest of the current sequence. RFO excludes the native amino
+acid and selects the single position/substitution with the highest probability
+across all unfixed positions. Equal scores use chain position, then alphabetical
+amino-acid order. Fully fixed binders bypass scoring.
+
+`mpnn_num_seqs` is the number of decoding orders to average;
+`mpnn_temperature` is accepted for compatibility but does not affect selection.
+The scoring worker imports the upstream model in `MPNN_PYTHON`, without
+importing the bundled Foundry MPNN engine.
 
 ## Outputs
 
@@ -54,9 +62,9 @@ recycle_<n>/
   rf3_config.yaml           # RF3 branch: exact per-cycle config
   mpnn/
     mpnn_design.pdb         # handoff complex
-    mpnn_result.json        # chain-A sequence and handoff path
-    mpnn_spec.json          # residue mask and sampling settings
-    samples/                # upstream FASTA and structures (if sampled)
+    mpnn_result.json        # chain-A sequence, handoff path, and mutation
+    mpnn_spec.json          # residue mask and scoring settings
+    samples/                # scoring spec.json and scores.json (if scored)
   <stem>/                   # RF3 history.json and modelhub_pred/<step>.cif
 <stem>_records.json
 ```

models/rfo/src/rfo/cycling/mpnn/mpnn_worker.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import json
+import math
 import os
 import shlex
 import shutil
@@ -16,7 +17,9 @@
 from pathlib import Path
 
 from Bio.PDB import PDBIO, PDBParser
-from Bio.SeqUtils import seq1
+from Bio.SeqUtils import seq1, seq3
+
+AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
 
 CHECKPOINTS = {
     "protein_mpnn": "proteinmpnn_v_48_020.pt",
@@ -35,6 +38,9 @@ def settings():
     repo = Path(repo).expanduser().resolve()
     if not (repo / "run.py").is_file():
         raise FileNotFoundError(f"MPNN_REPO must contain LigandMPNN run.py: {repo}")
+    for module in ("model_utils.py", "data_utils.py"):
+        if not (repo / module).is_file():
+            raise FileNotFoundError(f"MPNN_REPO must contain upstream {module}: {repo}")
     model = os.environ.get("MPNN_MODEL_TYPE", "protein_mpnn")
     if model not in CHECKPOINTS:
         raise ValueError(f"MPNN_MODEL_TYPE must be one of {tuple(CHECKPOINTS)}")
@@ -78,6 +84,37 @@ def chain_residues(structure, chain_id):
     return residues
 
 
+def select_mutation(scores, residues, positions, designed_chain):
+    """Choose one non-native AA by raw probability; ties follow chain/AA order."""
+    rows = {row["residue_id"]: row for row in scores}
+    if len(rows) != len(scores):
+        raise ValueError("MPNN scores contain duplicate residue IDs.")
+    best = None
+    for i in positions:
+        residue = residues[i]
+        residue_id = f"{designed_chain}{residue.id[1]}{residue.id[2].strip()}"
+        native = seq1(residue.resname)
+        row = rows.get(residue_id)
+        if row is None or row["native_aa"] != native:
+            raise ValueError(f"MPNN scores do not match input residue {residue_id}.")
+        probs = row["probabilities"]
+        for aa in AMINO_ACIDS:
+            value = probs.get(aa)
+            if value is None or not math.isfinite(value) or not 0 <= value <= 1:
+                raise ValueError(f"Invalid MPNN probability for {residue_id}/{aa}.")
+            if aa != native and (best is None or value > best["probability"]):
+                best = dict(
+                    position_0idx=i,
+                    residue_id=residue_id,
+                    from_aa=native,
+                    to_aa=aa,
+                    probability=value,
+                )
+    if best is None or best["probability"] <= 0:
+        raise ValueError("MPNN returned no positive-probability substitution.")
+    return best
+
+
 def redesign(
     structure_file,
     fixed_residues_0idx,
@@ -87,7 +124,7 @@ def redesign(
     seed,
     designed_chain="A",
 ):
-    """Return (chain sequence, PDB path), selecting the first sampled design.
+    """Return (chain sequence, PDB path) after one maximum-likelihood mutation.
 
     Masks are zero-based positions in the designed protein chain. Translate
     them to actual PDB residue IDs (including gaps and insertion codes). Keep
@@ -121,47 +158,28 @@ def redesign(
         seed=seed,
     )
 
-    # Upstream treats an empty redesign list as "design all". Explicitly skip
-    # sampling when the entire binder is fixed, retaining the original file.
+    mutation = None
+
+    # No eligible substitution: retain the original complex without scoring.
     if not positions:
         shutil.copy2(source, design_path)
     else:
         repo, model, checkpoint, python, sif = settings()
         spec.update(
             model_type=model, checkpoint_path=str(checkpoint), mpnn_repo=str(repo)
         )
-        # A fresh output directory prevents a failed/retried step from reading
-        # a previous sample. Keep the first sample as the stable handoff file.
+        # Score in the user's MPNN environment; the host needs no torch import.
+        # Fresh artifacts prevent failed/retried steps from reusing old scores.
+        worker = Path(__file__).with_name("score_worker.py").resolve()
         with tempfile.TemporaryDirectory(prefix="mpnn_", dir=out) as temp_dir:
-            command = [
-                python,
-                str(repo / "run.py"),
-                "--pdb_path",
-                str(source),
-                "--out_folder",
-                temp_dir,
-                "--model_type",
-                model,
-                f"--checkpoint_{model}",
-                str(checkpoint),
-                "--chains_to_design",
-                designed_chain,
-                "--redesigned_residues",
-                " ".join(ids),
-                "--batch_size",
-                "1",
-                "--number_of_batches",
-                str(num_seqs),
-                "--temperature",
-                str(temperature),
-                "--seed",
-                str(seed),
-                "--zero_indexed",
-                "0",
-            ]
+            score_path = Path(temp_dir) / "scores.json"
+            spec_path = Path(temp_dir) / "spec.json"
+            spec_path.write_text(json.dumps(spec))
+            command = [python, str(worker), str(spec_path), str(score_path)]
             if sif:
                 binds = [
                     str(repo),
+                    str(worker.parent),
                     str(source.parent),
                     str(out),
                     str(checkpoint.parent),
@@ -188,49 +206,35 @@ def redesign(
                 raise RuntimeError(
                     f"External MPNN failed with exit code {exc.returncode}"
                 ) from exc
-            # Upstream LigandMPNN parses --zero_indexed as `type=str`, so any
-            # non-empty string is truthy and the suffix is always "_0" regardless
-            # of the flag value. Accept the 0-indexed name and fall back to
-            # 1-indexed for older LigandMPNN builds.
-            backbones = Path(temp_dir) / "backbones"
-            candidate = backbones / f"{source.stem}_0.pdb"
-            if not candidate.is_file():
-                legacy = backbones / f"{source.stem}_1.pdb"
-                if legacy.is_file():
-                    candidate = legacy
-                else:
-                    raise RuntimeError(
-                        f"MPNN produced no first design: {candidate} (or legacy {legacy})"
-                    )
-            sampled = PDBParser(QUIET=True).get_structure("sample", str(candidate))
-            new_residues = chain_residues(sampled, designed_chain)
-            if [r.id for r in new_residues] != [r.id for r in residues]:
-                raise ValueError(
-                    "MPNN output changed the designed chain's residue IDs."
-                )
-            for i, (old, new) in enumerate(zip(residues, new_residues)):
-                if i in fixed and old.resname != new.resname:
-                    raise ValueError(
-                        f"MPNN changed a fixed residue: {designed_chain}{old.id[1]}"
-                    )
-                if old.resname != new.resname:
-                    old.resname = new.resname
-                    # Old side chains are chemically invalid after mutation.
-                    for atom in list(old):
-                        if atom.name not in {"N", "CA", "C", "O", "OXT"}:
-                            old.detach_child(atom.id)
+            if not score_path.is_file():
+                raise RuntimeError("MPNN produced no single-residue scores.")
+            scores = json.loads(score_path.read_text())["scores"]
+            mutation = select_mutation(scores, residues, positions, designed_chain)
+            residue = residues[mutation["position_0idx"]]
+            residue.resname = seq3(mutation["to_aa"]).upper()
+            # Only the mutated residue loses chemically invalid old side chains.
+            for atom in list(residue):
+                if atom.name not in {"N", "CA", "C", "O", "OXT"}:
+                    residue.detach_child(atom.id)
             writer = PDBIO()
             writer.set_structure(structure)
             writer.save(str(design_path))
-            # Retain upstream samples and FASTA for inspection / best-of-N work.
+            # Preserve scoring inputs/probabilities for inspection.
             samples = out / "samples"
             if samples.exists():
                 shutil.rmtree(samples)
             shutil.copytree(temp_dir, samples)
     sequence = "".join(seq1(res.resname) for res in residues)
     (out / "mpnn_spec.json").write_text(json.dumps(spec, indent=2) + "\n")
     result_path.write_text(
-        json.dumps({"sequence": sequence, "structure_pdb": str(design_path)}, indent=2)
+        json.dumps(
+            {
+                "sequence": sequence,
+                "structure_pdb": str(design_path),
+                "mutation": mutation,
+            },
+            indent=2,
+        )
         + "\n"
     )
     return sequence, str(design_path)