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 64d78bb..a780ee7 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}"