Read xz/gz-compressed AF3 confidences.json

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):

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 "<job>_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

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 "<job>_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 "<job>_<model>_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