diff --git a/alphapulldown/folding_backend/alphafold2_backend.py b/alphapulldown/folding_backend/alphafold2_backend.py index 6268f6e3..285d9cbb 100644 --- a/alphapulldown/folding_backend/alphafold2_backend.py +++ b/alphapulldown/folding_backend/alphafold2_backend.py @@ -867,6 +867,7 @@ def postprocess( compress_pickles: bool = False, remove_pickles: bool = False, remove_keys_from_pickles: bool = False, + storage_mode: str = "vanilla", convert_to_modelcif: bool = True, use_gpu_relax: bool = True, pae_plot_style: str = "red_blue", @@ -1088,5 +1089,6 @@ def postprocess( output_dir, compress_pickles=compress_pickles, remove_pickles=remove_pickles, - remove_keys=remove_keys_from_pickles + remove_keys=remove_keys_from_pickles, + storage_mode=storage_mode, ) diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index 0ac260aa..fe857807 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -2128,6 +2128,7 @@ def postprocess(**kwargs) -> None: prediction_results = kwargs.get('prediction_results') output_dir = kwargs.get('output_dir') fold_input_obj = kwargs.get('multimeric_object') + storage_mode = kwargs.get('storage_mode', 'vanilla') if prediction_results is None or output_dir is None or fold_input_obj is None: logging.warning('AF3 postprocess called with missing arguments; skipping.') @@ -2146,3 +2147,7 @@ def postprocess(**kwargs) -> None: output_dir=output_dir, job_name=job_name, ) + + # Apply the AF3 storage preset (no-op for the default 'vanilla'). + from alphapulldown.utils.post_modelling import post_prediction_process_af3 + post_prediction_process_af3(output_dir, job_name, storage_mode=storage_mode) diff --git a/alphapulldown/scripts/run_multimer_jobs.py b/alphapulldown/scripts/run_multimer_jobs.py index 8268fb7b..a8756d7c 100644 --- a/alphapulldown/scripts/run_multimer_jobs.py +++ b/alphapulldown/scripts/run_multimer_jobs.py @@ -158,6 +158,7 @@ def main(argv): "--compress_result_pickles": FLAGS.compress_result_pickles, "--remove_result_pickles": FLAGS.remove_result_pickles, "--remove_keys_from_pickles": FLAGS.remove_keys_from_pickles, + "--storage_mode": FLAGS.storage_mode, "--use_ap_style": True, "--use_gpu_relax": FLAGS.use_gpu_relax, "--protein_delimiter": FLAGS.protein_delimiter, diff --git a/alphapulldown/scripts/run_structure_prediction.py b/alphapulldown/scripts/run_structure_prediction.py index 38a190ab..1b1aae24 100644 --- a/alphapulldown/scripts/run_structure_prediction.py +++ b/alphapulldown/scripts/run_structure_prediction.py @@ -198,6 +198,16 @@ 'Whether the result pickles are going to be removed') flags.DEFINE_boolean('remove_keys_from_pickles',True, 'Whether to remove aligned_confidence_probs, distogram and masked_msa from pickles') +flags.DEFINE_enum( + 'storage_mode', 'vanilla', ['vanilla', 'slim', 'minimal'], + "High-level output storage preset (both backends). " + "'vanilla' (default): byte-identical to native AlphaFold2/3 output. " + "'slim': drop redundant duplicates (AF2 PAE inside pickles, AF3 top-level " + "confidences/data copies) and xz-compress the rest; output still has every " + "score/structure AlphaJudge and convert_to_modelcif need. " + "'minimal': slim plus drop the (unused) AF2 result pickles; for AF3 it is " + "identical to slim (non-best confidences are compressed, never deleted, to " + "avoid silently degrading per-sample PAE). No longer a vanilla AF layout.") flags.DEFINE_boolean('use_ap_style', False, 'Change output directory to include a description of the fold ' 'as seen in previous alphapulldown versions.') @@ -228,7 +238,7 @@ def _validate_flags_for_backend(backend_name: str) -> None: # Flags common to all backends common_flags = { 'input', 'output_directory', 'data_directory', 'features_directory', - 'protein_delimiter', 'fold_backend', 'random_seed', + 'protein_delimiter', 'fold_backend', 'random_seed', 'storage_mode', } # Backend-specific flags @@ -490,6 +500,7 @@ def main(argv): "compress_pickles": FLAGS.compress_result_pickles, "remove_pickles": FLAGS.remove_result_pickles, "remove_keys_from_pickles": FLAGS.remove_keys_from_pickles, + "storage_mode": FLAGS.storage_mode, "use_gpu_relax": FLAGS.use_gpu_relax, "models_to_relax": FLAGS.models_to_relax, "relax_best_score_threshold": FLAGS.relax_best_score_threshold, diff --git a/alphapulldown/utils/post_modelling.py b/alphapulldown/utils/post_modelling.py index 52669dda..2943d073 100644 --- a/alphapulldown/utils/post_modelling.py +++ b/alphapulldown/utils/post_modelling.py @@ -1,31 +1,48 @@ import os import json import gzip +import lzma import shutil import logging import pickle - -def compress_file(file_path): - """Compress a single file with gzip.""" - logging.info(f"Compressing file: {file_path}") - gz_path = file_path + '.gz' +# Keys whose arrays duplicate information already saved as standalone files in +# the AF2 output directory. ``predicted_aligned_error`` is byte-for-byte the same +# information as ``pae_.json`` (which AlphaJudge and convert_to_modelcif +# read), so it is pure redundancy inside the pickle. +PAE_KEY = "predicted_aligned_error" +MAX_PAE_KEY = "max_predicted_aligned_error" + + +def compress_file(file_path, *, method="gzip"): + """Compress a single file with gzip or xz, removing the original. + + ``method`` may be ``"gzip"`` (legacy default, ``.gz``) or ``"xz"`` (``.xz``, + ~2-3x smaller on AF2 result pickles / PAE JSON at the cost of CPU time). + """ + logging.info(f"Compressing file ({method}): {file_path}") + if method == "xz": + out_path = file_path + ".xz" + opener = lambda p: lzma.open(p, "wb", preset=6) + else: + out_path = file_path + ".gz" + opener = lambda p: gzip.open(p, "wb") try: - with open(file_path, 'rb') as f_in: - with gzip.open(gz_path, 'wb') as f_out: + with open(file_path, "rb") as f_in: + with opener(out_path) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(file_path) # Remove the original file after compression logging.info(f"File compressed and original removed: {file_path}") except Exception as e: logging.error(f"Failed to compress file: {file_path} with error: {e}") - return gz_path + return out_path -def compress_result_pickles(output_path): +def compress_result_pickles(output_path, *, method="gzip"): """Compress all .pkl files in the output directory.""" for file_name in os.listdir(output_path): if file_name.endswith('.pkl'): - compress_file(os.path.join(output_path, file_name)) + compress_file(os.path.join(output_path, file_name), method=method) def remove_keys_from_pickle(file_path, keys_to_remove): @@ -43,16 +60,39 @@ def remove_keys_from_pickle(file_path, keys_to_remove): # Save the modified data back to the pickle file with open(file_path, 'wb') as f: - pickle.dump(data, f) + pickle.dump(data, f, protocol=4) logging.info(f"Keys removed and file updated: {file_path}") except Exception as e: logging.error(f"Failed to remove keys from file: {file_path} with error: {e}") -def post_prediction_process(output_path, compress_pickles=False, remove_pickles=False, remove_keys=False): - """Process resulted files after the prediction.""" +def post_prediction_process( + output_path, + compress_pickles=False, + remove_pickles=False, + remove_keys=False, + storage_mode="vanilla", +): + """Process resulted files after the prediction. + + ``storage_mode`` is the high-level AlphaFold2 storage preset: + + - ``"vanilla"``: leave the result pickles untouched beyond the explicit + ``compress_pickles`` / ``remove_pickles`` / ``remove_keys`` flags. Output + stays a drop-in for tools expecting raw AlphaFold2 layout. + - ``"slim"``: additionally strip the redundant ``predicted_aligned_error`` + array from every pickle (it is preserved in ``pae_.json``, which is + what AlphaJudge and convert_to_modelcif actually read) and xz-compress the + remaining pickles. + - ``"minimal"``: drop all result pickles entirely. Neither AlphaJudge nor + convert_to_modelcif read pickle *contents*; scores come from the JSON + sidecars and structures from the PDBs. + """ keys_to_remove = ['aligned_confidence_probs', 'distogram', 'masked_msa'] + # In slim mode the PAE array inside the pickle duplicates pae_.json. + if storage_mode == "slim": + keys_to_remove = keys_to_remove + [PAE_KEY, MAX_PAE_KEY] try: # Get the best model from ranking_debug.json @@ -64,21 +104,34 @@ def post_prediction_process(output_path, compress_pickles=False, remove_pickles= logging.info(f"Identified best pickle file: {best_pickle}") - if remove_keys: + if storage_mode == "minimal": + # Pickle contents are unused downstream; delete them all. + logging.info("storage_mode=minimal: removing all result pickles.") + for file_name in os.listdir(output_path): + if file_name.endswith('.pkl'): + os.remove(os.path.join(output_path, file_name)) + return + + if remove_keys or storage_mode == "slim": logging.info(f"Removing specified keys from all pickle files in {output_path}") for file_name in os.listdir(output_path): if file_name.endswith('.pkl'): remove_keys_from_pickle(os.path.join(output_path, file_name), keys_to_remove) - if compress_pickles and remove_pickles: + # slim mode compresses with xz; explicit compress_pickles keeps gzip for + # backwards compatibility. + compress = compress_pickles or storage_mode == "slim" + method = "xz" if storage_mode == "slim" else "gzip" + + if compress and remove_pickles: # Compress only the best .pkl file and remove the others logging.info("Compressing and removing pickles based on conditions.") - compress_file(os.path.join(output_path, best_pickle)) + compress_file(os.path.join(output_path, best_pickle), method=method) remove_irrelevant_pickles(output_path, best_pickle) else: - if compress_pickles: + if compress: logging.info("Compressing all pickle files.") - compress_result_pickles(output_path) + compress_result_pickles(output_path, method=method) if remove_pickles: logging.info("Removing all non-best pickle files.") remove_irrelevant_pickles(output_path, best_pickle) @@ -96,3 +149,98 @@ 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_best_sample_dir(output_path): + """Return the ``seed-*_sample-*`` dir name with the highest ranking score.""" + import csv + ranking_csv = os.path.join(output_path, "ranking_scores.csv") + if not os.path.isfile(ranking_csv): + return None + best = (None, float("-inf")) + with open(ranking_csv, newline="") as f: + for row in csv.DictReader(f): + try: + score = float(row["ranking_score"]) + except (KeyError, ValueError): + continue + if score > best[1]: + best = (f"seed-{row['seed']}_sample-{row['sample']}", score) + return best[0] diff --git a/test/unit/test_post_modelling.py b/test/unit/test_post_modelling.py index 182908d6..9facf404 100644 --- a/test/unit/test_post_modelling.py +++ b/test/unit/test_post_modelling.py @@ -1,5 +1,7 @@ +import csv import gzip import json +import lzma import pickle import builtins from pathlib import Path @@ -146,3 +148,119 @@ def test_post_prediction_process_removes_non_best_pickles_without_compressing(tm assert sorted(path.name for path in tmp_path.glob("*.pkl")) == ["result_model_2.pkl"] assert not list(tmp_path.glob("*.gz")) + + +# --- storage_mode presets (AF2) --- + +def _make_af2_dir(tmp_path): + order = ["model_1", "model_2"] + (tmp_path / "ranking_debug.json").write_text(json.dumps({"order": order}), encoding="utf-8") + for m in order: + with open(tmp_path / f"result_{m}.pkl", "wb") as handle: + pickle.dump( + {"predicted_aligned_error": [[0.0]], "max_predicted_aligned_error": 30.0, + "iptm": 0.8, "plddt": [1.0], "distogram": 1, "masked_msa": 2}, + handle, + ) + (tmp_path / f"pae_{m}.json").write_text("[]", encoding="utf-8") + return order + + +def test_storage_mode_vanilla_leaves_pickles_untouched(tmp_path): + _make_af2_dir(tmp_path) + post_modelling.post_prediction_process(str(tmp_path), storage_mode="vanilla") + pkls = sorted(p.name for p in tmp_path.glob("*.pkl")) + assert pkls == ["result_model_1.pkl", "result_model_2.pkl"] + payload = pickle.load(open(tmp_path / "result_model_1.pkl", "rb")) + assert "predicted_aligned_error" in payload # vanilla keeps everything + + +def test_storage_mode_slim_strips_pae_and_xz_compresses(tmp_path): + order = _make_af2_dir(tmp_path) + post_modelling.post_prediction_process(str(tmp_path), storage_mode="slim") + assert not list(tmp_path.glob("*.pkl")) # all compressed + xz = sorted(tmp_path.glob("*.pkl.xz")) + assert len(xz) == 2 + payload = pickle.loads(lzma.open(xz[0]).read()) + assert "predicted_aligned_error" not in payload # PAE stripped (kept in sidecar) + assert "max_predicted_aligned_error" not in payload + assert "distogram" not in payload and "masked_msa" not in payload + assert "iptm" in payload and "plddt" in payload # scores retained + for m in order: # pae sidecars (what AlphaJudge reads) survive + assert (tmp_path / f"pae_{m}.json").is_file() + + +def test_storage_mode_minimal_deletes_all_pickles_keeps_sidecars(tmp_path): + order = _make_af2_dir(tmp_path) + post_modelling.post_prediction_process(str(tmp_path), storage_mode="minimal") + assert not list(tmp_path.glob("*.pkl")) + assert not list(tmp_path.glob("*.pkl.xz")) + 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() diff --git a/test/unit/test_script_entrypoints.py b/test/unit/test_script_entrypoints.py index 97aebf62..89de8635 100644 --- a/test/unit/test_script_entrypoints.py +++ b/test/unit/test_script_entrypoints.py @@ -331,6 +331,7 @@ def _load_run_multimer_jobs_module(): "compress_result_pickles": False, "remove_result_pickles": False, "remove_keys_from_pickles": True, + "storage_mode": "vanilla", "use_ap_style": False, "use_gpu_relax": True, "protein_delimiter": "+",