refactor: move AF3 ModelCIF and storage helpers beside the AF3 outputs

alphapulldown/folding_backend/alphafold3_backend.py

@@ -46,7 +46,7 @@
 from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache
 from alphapulldown.prediction.resume import finish_fold_record
 from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject
-from alphapulldown.utils.post_modelling import post_prediction_process_af3
+from alphapulldown.folding_backend.alphafold3_storage import post_prediction_process_af3
 
 
 @dataclasses.dataclass(frozen=True, slots=True, kw_only=True)

alphapulldown/folding_backend/alphafold3_outputs.py

@@ -31,17 +31,17 @@
     record_fold_complete,
     start_fold_record,
 )
-from alphapulldown.utils.af3_modelcif import augment_af3_modelcif_file
+from alphapulldown.folding_backend.alphafold3_modelcif import augment_af3_modelcif_file
+from alphapulldown.folding_backend.alphafold3_storage import (
+    AF3_RANKING_SCORES,
+    af3_best_sample_dir,
+    af3_ranking_scores,
+)
 from alphapulldown.utils.feature_metadata import (
     extract_metadata_from_fold_input,
     load_feature_metadata_sidecars,
 )
 from alphapulldown.utils.file_handling import open_atomic
-from alphapulldown.utils.post_modelling import (
-    AF3_RANKING_SCORES,
-    af3_best_sample_dir,
-    af3_ranking_scores,
-)
 
 if TYPE_CHECKING:
     from alphapulldown.folding_backend.alphafold3_runner import ResultsForSeed

alphapulldown/folding_backend/alphafold3_storage.py

@@ -0,0 +1,113 @@
+"""
+A finished AlphaFold 3 output directory: its ranking ledger (``ranking_scores.csv``),
+the best sample it names, and the ``--storage_mode`` presets that slim the directory.
+
+No AlphaFold 3 import, so tools without AF3 can read an AF3 output directory.
+"""
+
+import csv
+import logging
+import os
+
+from alphapulldown.utils.post_modelling import compress_file
+
+AF3_RANKING_SCORES = "ranking_scores.csv"
+
+
+def post_prediction_process_af3(output_path, job_name, storage_mode="vanilla"):
+    """Apply the AlphaFold3 storage preset to a single prediction directory.
+
+    Layout written by the AF3 backend (vanilla):
+        <output_path>/<job_name>_confidences.json      <- copy of best sample
+        <output_path>/<job_name>_data.json             <- copy of features input
+        <output_path>/<job_name>_model.cif             <- best model
+        <output_path>/<job_name>_summary_confidences.json
+        <output_path>/seed-*_sample-*/confidences.json <- per-sample (large)
+        <output_path>/seed-*_sample-*/model.cif
+        <output_path>/seed-*_sample-*/summary_confidences.json
+
+    The best sample's ``confidences.json`` is always left UNCOMPRESSED so
+    AlphaJudge (which reads the best model's PAE from it and has no xz support)
+    keeps working.
+
+    ``"slim"`` and ``"minimal"`` behave identically for AlphaFold3. Deleting a
+    non-best sample's ``confidences.json`` would not lose disk-only data: it is
+    the sole source of that sample's full token x token PAE matrix, and with it
+    gone AlphaJudge silently falls back to ``summary_confidences.json`` (a much
+    coarser per-chain-pair PAE), degrading PAE-derived scores for that sample
+    without any error. So both presets compress rather than delete it; the AF2
+    backend is where ``minimal`` additionally drops the (genuinely unused)
+    result pickles.
+
+    - ``"vanilla"``: no-op, keep byte-identical AlphaFold3 output.
+    - ``"slim"`` / ``"minimal"``: delete the two top-level duplicates
+      (``*_confidences.json`` is byte-identical to the best sample's,
+      ``*_data.json`` duplicates the saved features input) and xz-compress the
+      per-sample ``confidences.json`` of the NON-best samples. The best sample's
+      stays plain; every structure and summary is retained, and no PAE is lost.
+    """
+    if storage_mode == "vanilla":
+        return
+
+    try:
+        # Identify the best sample dir from ranking_scores.csv (highest score).
+        best_sample_dir = af3_best_sample_dir(output_path)
+
+        # Top-level duplicates: confidences.json mirrors the best sample,
+        # data.json mirrors features/<name>_af3_input.json. Remove them only if
+        # the best sample's own confidences.json is present (so the information
+        # is not lost), since AlphaJudge can fall back to either location.
+        best_conf_present = best_sample_dir is not None and os.path.isfile(
+            os.path.join(output_path, best_sample_dir, "confidences.json")
+        )
+        top_dups = [f"{job_name}_data.json"]
+        if best_conf_present:
+            top_dups.append(f"{job_name}_confidences.json")
+        for suffix in top_dups:
+            dup = os.path.join(output_path, suffix)
+            if os.path.isfile(dup):
+                logging.info(f"storage_mode={storage_mode}: removing top-level duplicate {dup}")
+                os.remove(dup)
+
+        sample_dirs = sorted(
+            d for d in os.listdir(output_path)
+            if os.path.isdir(os.path.join(output_path, d))
+            and d.startswith("seed-") and "_sample-" in d
+        )
+        for sample in sample_dirs:
+            conf = os.path.join(output_path, sample, "confidences.json")
+            if not os.path.isfile(conf):
+                continue
+            if sample == best_sample_dir:
+                # Keep plain so AlphaJudge reads best-model PAE directly.
+                continue
+            # slim and minimal both compress (never delete) non-best confidences
+            # so the full per-sample PAE matrix is never lost.
+            compress_file(conf, method="xz")
+
+    except FileNotFoundError as e:
+        logging.error(f"AF3 post-processing error: {e}.")
+    except Exception as e:
+        logging.error(f"Unexpected AF3 post-processing error: {e}")
+
+
+def af3_ranking_scores(output_path):
+    """The (seed, sample, ranking score) rows of an AF3 fold's ranking_scores.csv, in
+    order; empty if it is missing or unreadable."""
+    try:
+        with open(os.path.join(output_path, AF3_RANKING_SCORES), newline="") as f:
+            return [
+                (int(row["seed"]), int(row["sample"]), float(row["ranking_score"]))
+                for row in csv.DictReader(f)
+            ]
+    except (OSError, KeyError, ValueError):
+        return []
+
+
+def af3_best_sample_dir(output_path):
+    """The ``seed-*_sample-*`` dir with the highest ranking score; the first of equals."""
+    rows = af3_ranking_scores(output_path)
+    if not rows:
+        return None
+    seed, sample, _ = max(rows, key=lambda row: row[2])
+    return f"seed-{seed}_sample-{sample}"

alphapulldown/utils/post_modelling.py

@@ -1,4 +1,3 @@
-import csv
 import os
 import json
 import gzip
@@ -13,7 +12,6 @@
 # read), so it is pure redundancy inside the pickle.
 PAE_KEY = "predicted_aligned_error"
 MAX_PAE_KEY = "max_predicted_aligned_error"
-AF3_RANKING_SCORES = "ranking_scores.csv"
 
 
 def compress_file(file_path, *, method="gzip"):
@@ -151,102 +149,3 @@ def remove_irrelevant_pickles(output_path, best_pickle):
         if file_name.endswith('.pkl') and file_name != best_pickle:
             logging.info(f"Removing irrelevant pickle file: {file_path}")
             os.remove(file_path)
-
-
-def post_prediction_process_af3(output_path, job_name, storage_mode="vanilla"):
-    """Apply the AlphaFold3 storage preset to a single prediction directory.
-
-    Layout written by the AF3 backend (vanilla):
-        <output_path>/<job_name>_confidences.json      <- copy of best sample
-        <output_path>/<job_name>_data.json             <- copy of features input
-        <output_path>/<job_name>_model.cif             <- best model
-        <output_path>/<job_name>_summary_confidences.json
-        <output_path>/seed-*_sample-*/confidences.json <- per-sample (large)
-        <output_path>/seed-*_sample-*/model.cif
-        <output_path>/seed-*_sample-*/summary_confidences.json
-
-    The best sample's ``confidences.json`` is always left UNCOMPRESSED so
-    AlphaJudge (which reads the best model's PAE from it and has no xz support)
-    keeps working.
-
-    ``"slim"`` and ``"minimal"`` behave identically for AlphaFold3. Deleting a
-    non-best sample's ``confidences.json`` would not lose disk-only data: it is
-    the sole source of that sample's full token x token PAE matrix, and with it
-    gone AlphaJudge silently falls back to ``summary_confidences.json`` (a much
-    coarser per-chain-pair PAE), degrading PAE-derived scores for that sample
-    without any error. So both presets compress rather than delete it; the AF2
-    backend is where ``minimal`` additionally drops the (genuinely unused)
-    result pickles.
-
-    - ``"vanilla"``: no-op, keep byte-identical AlphaFold3 output.
-    - ``"slim"`` / ``"minimal"``: delete the two top-level duplicates
-      (``*_confidences.json`` is byte-identical to the best sample's,
-      ``*_data.json`` duplicates the saved features input) and xz-compress the
-      per-sample ``confidences.json`` of the NON-best samples. The best sample's
-      stays plain; every structure and summary is retained, and no PAE is lost.
-    """
-    if storage_mode == "vanilla":
-        return
-
-    try:
-        # Identify the best sample dir from ranking_scores.csv (highest score).
-        best_sample_dir = af3_best_sample_dir(output_path)
-
-        # Top-level duplicates: confidences.json mirrors the best sample,
-        # data.json mirrors features/<name>_af3_input.json. Remove them only if
-        # the best sample's own confidences.json is present (so the information
-        # is not lost), since AlphaJudge can fall back to either location.
-        best_conf_present = best_sample_dir is not None and os.path.isfile(
-            os.path.join(output_path, best_sample_dir, "confidences.json")
-        )
-        top_dups = [f"{job_name}_data.json"]
-        if best_conf_present:
-            top_dups.append(f"{job_name}_confidences.json")
-        for suffix in top_dups:
-            dup = os.path.join(output_path, suffix)
-            if os.path.isfile(dup):
-                logging.info(f"storage_mode={storage_mode}: removing top-level duplicate {dup}")
-                os.remove(dup)
-
-        sample_dirs = sorted(
-            d for d in os.listdir(output_path)
-            if os.path.isdir(os.path.join(output_path, d))
-            and d.startswith("seed-") and "_sample-" in d
-        )
-        for sample in sample_dirs:
-            conf = os.path.join(output_path, sample, "confidences.json")
-            if not os.path.isfile(conf):
-                continue
-            if sample == best_sample_dir:
-                # Keep plain so AlphaJudge reads best-model PAE directly.
-                continue
-            # slim and minimal both compress (never delete) non-best confidences
-            # so the full per-sample PAE matrix is never lost.
-            compress_file(conf, method="xz")
-
-    except FileNotFoundError as e:
-        logging.error(f"AF3 post-processing error: {e}.")
-    except Exception as e:
-        logging.error(f"Unexpected AF3 post-processing error: {e}")
-
-
-def af3_ranking_scores(output_path):
-    """The (seed, sample, ranking score) rows of an AF3 fold's ranking_scores.csv, in
-    order; empty if it is missing or unreadable."""
-    try:
-        with open(os.path.join(output_path, AF3_RANKING_SCORES), newline="") as f:
-            return [
-                (int(row["seed"]), int(row["sample"]), float(row["ranking_score"]))
-                for row in csv.DictReader(f)
-            ]
-    except (OSError, KeyError, ValueError):
-        return []
-
-
-def af3_best_sample_dir(output_path):
-    """The ``seed-*_sample-*`` dir with the highest ranking score; the first of equals."""
-    rows = af3_ranking_scores(output_path)
-    if not rows:
-        return None
-    seed, sample, _ = max(rows, key=lambda row: row[2])
-    return f"seed-{seed}_sample-{sample}"

docker/alphafold3.dockerfile

@@ -153,7 +153,7 @@ RUN uv pip install --no-cache "modelcif>=1.6" && \
       test/unit/test_feature_batch.py \
       test/unit/test_create_batch_features.py \
       test/unit/test_feature_metadata.py::test_embedded_json_is_accepted_and_round_tripped_by_vanilla_alphafold3 \
-      test/unit/test_af3_modelcif.py::test_augment_real_af3_modelcif_preserves_comments_and_is_modelcif_readable \
+      test/unit/test_alphafold3_modelcif.py::test_augment_real_af3_modelcif_preserves_comments_and_is_modelcif_readable \
       test/unit/test_convert_to_modelcif_helpers.py \
       test/unit/test_af3_fused_triangles.py \
       test/unit/test_alphafold3_backend_helpers.py \

test/cluster/code_changes.py

@@ -21,7 +21,7 @@
     "alphapulldown/features",
     "alphapulldown/scripts/create_individual_features.py",
     "alphapulldown/objects.py",
-    "alphapulldown/utils/af3_modelcif.py",
+    "alphapulldown/folding_backend/alphafold3_modelcif.py",
     "alphapulldown/utils/create_custom_template_db.py",
     "alphapulldown/utils/feature_metadata.py",
     "alphapulldown/utils/file_handling.py",

test/integration/test_backend_test_isolation.py

@@ -20,7 +20,7 @@ def test_af3_backend_stubs_do_not_leak_into_modelcif_tests():
             "-n", "0",
             "test/unit/test_alphafold3_backend_helpers.py::"
             "test_output_name_helpers_compact_and_normalise_fragments",
-            "test/unit/test_af3_modelcif.py::"
+            "test/unit/test_alphafold3_modelcif.py::"
             "test_augment_real_af3_modelcif_preserves_comments_and_is_modelcif_readable",
         ],
         cwd=root,

test/unit/test_af3_modelcif.py → test/unit/test_alphafold3_modelcif.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from alphapulldown.utils.af3_modelcif import (
+from alphapulldown.folding_backend.alphafold3_modelcif import (
     augment_af3_modelcif_file,
     build_alphapulldown_mmcif_updates,
     find_af3_modelcif_files,

test/unit/test_alphafold3_storage.py

@@ -0,0 +1,70 @@
+"""The --storage_mode presets for an AlphaFold 3 output directory."""
+
+import csv
+import json
+
+from alphapulldown.folding_backend import alphafold3_storage
+
+
+def _make_af3_dir(tmp_path, job="A_and_B"):
+    big = json.dumps({"pae": [[0.0]]})
+    (tmp_path / f"{job}_confidences.json").write_text(big, encoding="utf-8")
+    (tmp_path / f"{job}_data.json").write_text("{}", encoding="utf-8")
+    (tmp_path / f"{job}_model.cif").write_text("data_x\n", encoding="utf-8")
+    (tmp_path / f"{job}_summary_confidences.json").write_text("{}", encoding="utf-8")
+    rows = [("9", "0", 0.5), ("9", "1", 0.9), ("9", "2", 0.3)]  # best = sample-1
+    with open(tmp_path / "ranking_scores.csv", "w", newline="") as f:
+        w = csv.writer(f)
+        w.writerow(["seed", "sample", "ranking_score"])
+        w.writerows(rows)
+    for seed, sample, score in rows:
+        sd = tmp_path / f"seed-{seed}_sample-{sample}"
+        sd.mkdir()
+        (sd / "confidences.json").write_text(big, encoding="utf-8")
+        (sd / "model.cif").write_text("data_x\n", encoding="utf-8")
+        (sd / "summary_confidences.json").write_text(json.dumps({"ranking_score": score}), encoding="utf-8")
+    return job
+
+
+def test_af3_storage_mode_vanilla_is_noop(tmp_path):
+    job = _make_af3_dir(tmp_path)
+    alphafold3_storage.post_prediction_process_af3(str(tmp_path), job, storage_mode="vanilla")
+    assert (tmp_path / f"{job}_confidences.json").is_file()
+    assert (tmp_path / "seed-9_sample-0" / "confidences.json").is_file()
+
+
+def test_af3_storage_mode_slim_keeps_best_plain_xz_others(tmp_path):
+    job = _make_af3_dir(tmp_path)
+    alphafold3_storage.post_prediction_process_af3(str(tmp_path), job, storage_mode="slim")
+    # best (sample-1) plain for AlphaJudge; others xz
+    assert (tmp_path / "seed-9_sample-1" / "confidences.json").is_file()
+    assert not (tmp_path / "seed-9_sample-1" / "confidences.json.xz").exists()
+    for s in ("0", "2"):
+        assert (tmp_path / f"seed-9_sample-{s}" / "confidences.json.xz").is_file()
+        assert not (tmp_path / f"seed-9_sample-{s}" / "confidences.json").exists()
+    # top-level duplicates removed, structure kept
+    assert not (tmp_path / f"{job}_confidences.json").exists()
+    assert not (tmp_path / f"{job}_data.json").exists()
+    assert (tmp_path / f"{job}_model.cif").is_file()
+
+
+def test_af3_storage_mode_minimal_matches_slim_compresses_non_best(tmp_path):
+    # AF3 minimal == slim: non-best confidences are xz-compressed, never deleted,
+    # so the full per-sample PAE matrix is never lost (deleting it would silently
+    # degrade AlphaJudge to the coarser summary_confidences.json PAE).
+    job = _make_af3_dir(tmp_path)
+    alphafold3_storage.post_prediction_process_af3(str(tmp_path), job, storage_mode="minimal")
+    assert (tmp_path / "seed-9_sample-1" / "confidences.json").is_file()  # best plain
+    for s in ("0", "2"):
+        sd = tmp_path / f"seed-9_sample-{s}"
+        assert (sd / "confidences.json.xz").is_file()  # compressed, not deleted
+        assert not (sd / "confidences.json").exists()
+        assert (sd / "model.cif").is_file() and (sd / "summary_confidences.json").is_file()
+
+
+def test_af3_storage_mode_preserves_top_confidences_when_best_sample_lacks_it(tmp_path):
+    job = _make_af3_dir(tmp_path)
+    (tmp_path / "seed-9_sample-1" / "confidences.json").unlink()
+    alphafold3_storage.post_prediction_process_af3(str(tmp_path), job, storage_mode="slim")
+    # top-level copy is the only remaining source of best-model PAE -> must survive
+    assert (tmp_path / f"{job}_confidences.json").is_file()

test/unit/test_mmseqs2_gpu_dockerfiles.py

@@ -237,3 +237,17 @@ def test_ci_runs_the_tool_tests_with_the_images_mmseqs2():
         assert f"s/^ARG {arg}=//p" in text, arg
     assert "sha256sum -c -" in text
     assert "MMSEQS_INTEGRATION_BINARY=" in text
+
+
+@pytest.mark.parametrize(
+    "dockerfile", sorted((REPOSITORY / "docker").glob("*.dockerfile")), ids=lambda path: path.name
+)
+def test_every_test_an_image_runs_exists(dockerfile):
+    # A renamed test file or test fails the image build at its test stage, after an
+    # hour of building, and only in the container job.
+    for target in re.findall(r"\btest/[\w/]+\.py(?:::\w+)?", dockerfile.read_text()):
+        path, _, name = target.partition("::")
+        assert (REPOSITORY / path).is_file(), target
+        if name:
+            source = (REPOSITORY / path).read_text()
+            assert re.search(rf"^def {re.escape(name)}\(", source, re.MULTILINE), target