diff --git a/src/alphajudge/parsers/__init__.py b/src/alphajudge/parsers/__init__.py index af291aa3..5ee827f8 100644 --- a/src/alphajudge/parsers/__init__.py +++ b/src/alphajudge/parsers/__init__.py @@ -3,7 +3,9 @@ from pathlib import Path from typing import Any, Callable from abc import ABC, abstractmethod +import gzip import json +import lzma from Bio.PDB import PDBParser, MMCIFParser from ..confidence import Confidence from ..geometry import is_pae_token_residue, representative_atom @@ -23,11 +25,49 @@ def detect(d: Path) -> bool: ... @abstractmethod def parse_run(self, d: Path) -> Run: ... - @staticmethod - def _read_json(p: Path) -> dict: + # Compression magic numbers (file header bytes). + _MAGIC = ((b"\xfd7zXZ\x00", lzma.open), (b"\x1f\x8b", gzip.open)) + + @classmethod + def _open_maybe_compressed(cls, p: Path): + """Open ``p`` for text reading, detecting xz/gz by its magic bytes. + + Compression is identified from the file header rather than the + extension, so a compressed file is handled regardless of how it is + named. + """ + with p.open("rb") as fh: + head = fh.read(6) + for magic, opener in cls._MAGIC: + if head.startswith(magic): + return opener(p, "rt") + return p.open("rt") + + @classmethod + def _read_json(cls, p: Path) -> dict: + """Read a JSON file, transparently handling xz/gz compression. + + AlphaPulldown's ``--storage_mode slim/minimal`` may store large JSON + sidecars (e.g. AF3 per-sample ``confidences.json``) compressed. If the + plain path does not exist, fall back to a ``.xz`` or ``.gz`` sibling so + scoring is unaffected by the storage mode. Compression is detected from + the file's magic bytes, not its extension. + """ try: - with p.open() as fh: return json.load(fh) - except Exception: return {} + target = p + if not target.exists(): + # Plain path absent: try a compressed sibling from slim/minimal. + for ext in (".xz", ".gz"): + cp = p.with_name(p.name + ext) + if cp.exists(): + target = cp + break + else: + return {} + with cls._open_maybe_compressed(target) as fh: + return json.load(fh) + except Exception: + return {} @staticmethod def _load_structure(path: str): diff --git a/src/alphajudge/parsers/af3.py b/src/alphajudge/parsers/af3.py index 57fea381..389d9dc7 100644 --- a/src/alphajudge/parsers/af3.py +++ b/src/alphajudge/parsers/af3.py @@ -117,15 +117,44 @@ def _find_af3_json( is_best_model: bool, ) -> Path: model_dir = d / model - candidates = [model_dir / f"{kind}.json"] + + # Exact-name candidates, each also tried with .xz/.gz so that + # AlphaPulldown's slim/minimal storage modes (which compress per-sample + # confidences.json) are read transparently. + exact = [model_dir / f"{kind}.json"] if job_prefix: - candidates.append(model_dir / f"{job_prefix}_{model}_{kind}.json") + exact.append(model_dir / f"{job_prefix}_{model}_{kind}.json") if is_best_model: - candidates.append(d / f"{job_prefix}_{kind}.json") - if model_dir.is_dir(): - candidates.extend(sorted(model_dir.glob(f"*_{kind}.json"))) + exact.append(d / f"{job_prefix}_{kind}.json") if is_best_model: - candidates.append(d / f"ranked_0_{kind}.json") + exact.append(d / f"ranked_0_{kind}.json") + + candidates: list[Path] = [] + for base in exact: + candidates.append(base) + candidates.append(base.with_name(base.name + ".xz")) + candidates.append(base.with_name(base.name + ".gz")) + + # Wildcard fallback within the model dir. Exclude summary_confidences.json + # from a "confidences" search: it ends in "_confidences.json" but only + # carries coarse per-chain-pair PAE, not the full token x token matrix. + # Match the summary file by suffix so both "summary_confidences.json" and + # the job-prefixed "_summary_confidences.json" (official AF3 layout) + # are excluded. + def _is_summary(name: str) -> bool: + base = name + for ext in (".xz", ".gz"): + if base.endswith(ext): + base = base[: -len(ext)] + return base == "summary_confidences.json" or base.endswith("_summary_confidences.json") + + if model_dir.is_dir(): + for pattern in (f"*_{kind}.json", f"*_{kind}.json.xz", f"*_{kind}.json.gz"): + for hit in sorted(model_dir.glob(pattern)): + if kind == "confidences" and _is_summary(hit.name): + continue + candidates.append(hit) + return cls._find_existing(candidates) or candidates[0] @classmethod diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index b4e15bea..c899c9ac 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -750,3 +750,99 @@ def test_cli_recursive_single_directory_root(tmp_path: Path, af2_dir_src: Path): assert_has_expected_headers(r1, where=str(out_nonrec)) assert_has_expected_headers(r2, where=str(out_rec)) assert len(r1) == len(r2), "Recursive root with one run should match direct-run output row count" + + +# ------------------------- +# Compressed-confidences reading (AlphaPulldown slim/minimal storage modes) +# ------------------------- + +import gzip +import lzma + +from alphajudge.parsers import BaseParser + + +def test_read_json_reads_plain_xz_and_gz(tmp_path): + payload = {"a": 1, "pae": [[0.0, 1.0], [1.0, 0.0]]} + plain = tmp_path / "confidences.json" + plain.write_text(json.dumps(payload)) + xz = tmp_path / "x.json.xz" + with lzma.open(xz, "wt") as fh: + json.dump(payload, fh) + gz = tmp_path / "x.json.gz" + with gzip.open(gz, "wt") as fh: + json.dump(payload, fh) + + assert BaseParser._read_json(plain) == payload + assert BaseParser._read_json(xz) == payload + assert BaseParser._read_json(gz) == payload + + +def test_read_json_detects_compression_by_magic_not_extension(tmp_path): + # xz bytes stored under a plain .json name must still decode. + payload = {"k": "v"} + mislabeled = tmp_path / "confidences.json" + with lzma.open(mislabeled, "wt") as fh: + json.dump(payload, fh) + assert BaseParser._read_json(mislabeled) == payload + + +def test_read_json_falls_back_to_compressed_sibling(tmp_path): + payload = {"pae": [[0.0]]} + # Only the .xz sibling exists; the plain path is requested. + with lzma.open(tmp_path / "confidences.json.xz", "wt") as fh: + json.dump(payload, fh) + assert BaseParser._read_json(tmp_path / "confidences.json") == payload + + +def test_read_json_missing_returns_empty(tmp_path): + assert BaseParser._read_json(tmp_path / "nope.json") == {} + + +def _write_af3_sample(model_dir: Path, *, compress: bool, with_summary: bool = True): + model_dir.mkdir(parents=True, exist_ok=True) + conf = {"pae": [[0.0, 5.0], [5.0, 0.0]], "token_chain_ids": ["A", "B"]} + if compress: + with lzma.open(model_dir / "confidences.json.xz", "wt") as fh: + json.dump(conf, fh) + else: + (model_dir / "confidences.json").write_text(json.dumps(conf)) + if with_summary: + (model_dir / "summary_confidences.json").write_text(json.dumps({"iptm": 0.5, "ptm": 0.5})) + + +def test_find_af3_json_finds_compressed_confidences(tmp_path): + md = tmp_path / "seed-1_sample-0" + _write_af3_sample(md, compress=True) + found = AF3Parser._find_af3_json(tmp_path, "seed-1_sample-0", "confidences", None, False) + assert found.name == "confidences.json.xz" + assert found.exists() + + +def test_find_af3_json_does_not_shadow_confidences_with_summary(tmp_path): + # Even when only summary_confidences.json is present, the confidences search + # must not return it (it carries only coarse per-chain-pair PAE). The caller + # then falls back to summary explicitly. + md = tmp_path / "seed-1_sample-0" + md.mkdir(parents=True) + (md / "summary_confidences.json").write_text(json.dumps({"iptm": 0.5})) + found = AF3Parser._find_af3_json(tmp_path, "seed-1_sample-0", "confidences", None, False) + assert not found.name.startswith("summary_") + assert not found.exists() # genuinely no confidences file present + + +def test_find_af3_json_excludes_job_prefixed_summary(tmp_path): + # Official AF3 layout uses "_summary_confidences.json", which does NOT + # start with "summary_". The confidences search must still not return it, + # and must not let it shadow the real "__confidences.json". + job = "hello_fold" + model = "seed-1_sample-0" + md = tmp_path / model + md.mkdir(parents=True) + (md / f"{job}_{model}_confidences.json").write_text( + json.dumps({"pae": [[0.0]], "token_chain_ids": ["A"]}) + ) + (md / f"{job}_{model}_summary_confidences.json").write_text(json.dumps({"iptm": 0.5})) + found = AF3Parser._find_af3_json(tmp_path, model, "confidences", job, False) + assert found.name == f"{job}_{model}_confidences.json" + assert "summary" not in found.name