From 1600842aeb6e7df34569aeb9325c69af597fbab0 Mon Sep 17 00:00:00 2001
From: mason h <9421505+drawadiagram@users.noreply.github.com>
Date: Thu, 23 Jul 2026 13:42:47 -0500
Subject: [PATCH] fix filename bug in fold-to-scaffoldd forwarding

---
 examples/small_molecule_binding/mock.py       | 10 ++++++----
 .../run_test_small_molecule_binding.py        | 19 +++++++++++++++++++
 .../small_molecule_binding.py                 |  2 +-
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/examples/small_molecule_binding/mock.py b/examples/small_molecule_binding/mock.py
index 5e1eef6..4e9ba76 100644
--- a/examples/small_molecule_binding/mock.py
+++ b/examples/small_molecule_binding/mock.py
@@ -177,15 +177,17 @@ async def af2(task_description=None, **kwargs):
         os.makedirs(f"{taskdir}/out", exist_ok=True)
 
         for rank in range(1, 6):
-            with open(f"{taskdir}/out/rank_{rank:03d}.pdb", "w") as fh:
+            stem = f"binder_scores_rank_{rank:03d}_alphafold2_ptm_model_1_seed_000"
+            with open(f"{taskdir}/out/{stem.replace('_scores_', '_unrelaxed_')}.pdb", "w") as fh:
                 fh.write(f"REMARK  mock af2 rank_{rank:03d}\nEND\n")
-            with open(f"{taskdir}/out/rank_{rank:03d}_scores.json", "w") as fh:
+            with open(f"{taskdir}/out/{stem}.json", "w") as fh:
                 json.dump({"plddt": [85.0 + rank] * 50, "max_pae": 5.0}, fh)
 
     @pipeline.auto_register_task(local_task=True)
     async def analysis_fold(task_description=None, **kwargs):
-        taskdir    = f"{pipeline.base_path}/{pipeline.taskcount}_alphafold"
-        best_model = f"{taskdir}/out/rank_005.pdb"
+        taskdir     = f"{pipeline.base_path}/{pipeline.taskcount}_alphafold"
+        best_scores = "binder_scores_rank_005_alphafold2_ptm_model_1_seed_000.json"
+        best_model  = f"{taskdir}/out/" + best_scores.replace('_scores_', '_unrelaxed_').replace('.json', '.pdb')
         pipeline.state['best_af2_model']        = best_model
         pipeline.state['last_analysis_step']    = 'fold'
         pipeline.state['last_analysis_metrics'] = {
diff --git a/examples/small_molecule_binding/run_test_small_molecule_binding.py b/examples/small_molecule_binding/run_test_small_molecule_binding.py
index 3c5ae0e..0a78c65 100644
--- a/examples/small_molecule_binding/run_test_small_molecule_binding.py
+++ b/examples/small_molecule_binding/run_test_small_molecule_binding.py
@@ -43,6 +43,24 @@ def setup_mock_inputs(pipeline_name: str) -> None:
                 fh.write(f"# mock placeholder: {fname}\n")
 
 
+def check_af2_filename_derivation() -> None:
+    """Regression check for analysis_fold()'s scores.json -> unrelaxed.pdb
+    filename derivation, against real ColabFold (colabfold_batch) naming."""
+    cases = [
+        (
+            "binder_scores_rank_001_alphafold2_model_3_seed_999.json",
+            "binder_unrelaxed_rank_001_alphafold2_model_3_seed_999.pdb",
+        ),
+        (
+            "binder_scores_rank_005_alphafold2_ptm_model_1_seed_000.json",
+            "binder_unrelaxed_rank_005_alphafold2_ptm_model_1_seed_000.pdb",
+        ),
+    ]
+    for sf, expected in cases:
+        derived = sf.replace('_scores_', '_unrelaxed_').replace('.json', '.pdb')
+        assert derived == expected, f"{sf!r} -> {derived!r}, expected {expected!r}"
+
+
 async def run_mock_test() -> None:
     pipeline_name = "p1"
     setup_mock_inputs(pipeline_name)
@@ -70,4 +88,5 @@ async def run_mock_test() -> None:
 
 
 if __name__ == "__main__":
+    check_af2_filename_derivation()
     asyncio.run(run_mock_test())
diff --git a/examples/small_molecule_binding/small_molecule_binding.py b/examples/small_molecule_binding/small_molecule_binding.py
index d4264e1..e8fc3f2 100644
--- a/examples/small_molecule_binding/small_molecule_binding.py
+++ b/examples/small_molecule_binding/small_molecule_binding.py
@@ -540,7 +540,7 @@ async def analysis_fold():
                     mean_plddt = sum(arr) / len(arr)
                     if mean_plddt > best_plddt:
                         best_plddt = mean_plddt
-                        best_model = sf.replace('_scores.json', '.pdb')
+                        best_model = sf.replace('_scores_', '_unrelaxed_').replace('.json', '.pdb')
 
             if best_model:
                 full_model_path = f"{out_dir}/{best_model}"