diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index f315bd3b..df3a8d6c 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/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) diff --git a/alphapulldown/utils/af3_modelcif.py b/alphapulldown/folding_backend/alphafold3_modelcif.py similarity index 100% rename from alphapulldown/utils/af3_modelcif.py rename to alphapulldown/folding_backend/alphafold3_modelcif.py diff --git a/alphapulldown/folding_backend/alphafold3_outputs.py b/alphapulldown/folding_backend/alphafold3_outputs.py index 9013c7f9..d27e357f 100644 --- a/alphapulldown/folding_backend/alphafold3_outputs.py +++ b/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 diff --git a/alphapulldown/folding_backend/alphafold3_storage.py b/alphapulldown/folding_backend/alphafold3_storage.py new file mode 100644 index 00000000..82e66b63 --- /dev/null +++ b/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): + /_confidences.json <- copy of best sample + /_data.json <- copy of features input + /_model.cif <- best model + /_summary_confidences.json + /seed-*_sample-*/confidences.json <- per-sample (large) + /seed-*_sample-*/model.cif + /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/_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}" diff --git a/alphapulldown/utils/post_modelling.py b/alphapulldown/utils/post_modelling.py index 2b59a717..c5abed61 100644 --- a/alphapulldown/utils/post_modelling.py +++ b/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): - /_confidences.json <- copy of best sample - /_data.json <- copy of features input - /_model.cif <- best model - /_summary_confidences.json - /seed-*_sample-*/confidences.json <- per-sample (large) - /seed-*_sample-*/model.cif - /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/_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}" diff --git a/docker/alphafold3.dockerfile b/docker/alphafold3.dockerfile index bd62edd4..8392697d 100644 --- a/docker/alphafold3.dockerfile +++ b/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 \ diff --git a/test/cluster/code_changes.py b/test/cluster/code_changes.py index d64a23d9..fe35526e 100644 --- a/test/cluster/code_changes.py +++ b/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", diff --git a/test/integration/test_backend_test_isolation.py b/test/integration/test_backend_test_isolation.py index 41400331..4eab791f 100644 --- a/test/integration/test_backend_test_isolation.py +++ b/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, diff --git a/test/unit/test_af3_modelcif.py b/test/unit/test_alphafold3_modelcif.py similarity index 99% rename from test/unit/test_af3_modelcif.py rename to test/unit/test_alphafold3_modelcif.py index b41db6b8..073c764c 100644 --- a/test/unit/test_af3_modelcif.py +++ b/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, diff --git a/test/unit/test_alphafold3_storage.py b/test/unit/test_alphafold3_storage.py new file mode 100644 index 00000000..c95b095a --- /dev/null +++ b/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() diff --git a/test/unit/test_mmseqs2_gpu_dockerfiles.py b/test/unit/test_mmseqs2_gpu_dockerfiles.py index 27f2cce9..36ec44d4 100644 --- a/test/unit/test_mmseqs2_gpu_dockerfiles.py +++ b/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 diff --git a/test/unit/test_post_modelling.py b/test/unit/test_post_modelling.py index 9facf404..55a45dd5 100644 --- a/test/unit/test_post_modelling.py +++ b/test/unit/test_post_modelling.py @@ -1,4 +1,3 @@ -import csv import gzip import json import lzma @@ -198,69 +197,3 @@ def test_storage_mode_minimal_deletes_all_pickles_keeps_sidecars(tmp_path): assert (tmp_path / "ranking_debug.json").is_file() for m in order: assert (tmp_path / f"pae_{m}.json").is_file() - - -# --- storage_mode presets (AF3) --- - -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) - post_modelling.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) - post_modelling.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) - post_modelling.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() - post_modelling.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()