Add interface contact probability scores

README.md

@@ -17,9 +17,9 @@ AlphaJudge parses AF2, AF3, and Boltz-2 outputs and summarizes per-model / per-i
 
 | category | metrics (examples) | notes |
 | --- | --- | --- |
-| **AlphaFold internal** | ipTM, pTM, iptm+ptm/confidence_score, avg interface PAE, avg interface pLDDT | unified for AF2/AF3 |
+| **AlphaFold internal** | ipTM, pTM, iptm+ptm/confidence_score, avg interface PAE, avg interface pLDDT, contact probabilities | unified for AF2/AF3 where available |
 | **physical & geometric** | buried area, contact pairs, H-bonds, salt bridges, interface composition, shape complementarity | self-contained |
-| **derived scores** | pDockQ, pDockQ2, mpDockQ, ipSAE, LIS, interface score | implemented here |
+| **derived scores** | pDockQ, pDockQ2, mpDockQ, ipSAE, LIS, cLIS, iLIS, interface score, contact-probability summaries | implemented here |
 
 Use cases: rank poses, sanity-check AF confidences, or export features for ML.
 
@@ -181,7 +181,7 @@ process_many(
 )
 ```
 
-Key outputs per interface include: `average_interface_pae`, `interface_average_plddt`, `interface_contact_pairs`, `interface_area`, `interface_hb`, `interface_sb`, `interface_sc`, `interface_solv_en`, `interface_ipSAE`, `interface_LIS`, `interface_pDockQ2`, and per-run `pDockQ/mpDockQ`.
+Key outputs per interface include: `average_interface_pae`, `interface_average_plddt`, `interface_contact_pairs`, `interface_contact_prob_max`, `interface_contact_prob_top10_mean`, `interface_area`, `interface_hb`, `interface_sb`, `interface_sc`, `interface_solv_en`, `interface_ipSAE`, `interface_LIS`, `interface_cLIS`, `interface_iLIS`, `interface_pDockQ2`, and per-run `pDockQ/mpDockQ`.
 
 ---
 
@@ -207,7 +207,8 @@ AlphaJudge writes `interfaces.csv` with one row per interface (and includes the
 - **iptm_ptm, iptm, ptm, confidence_score**: unified AF confidences
 - **pDockQ/mpDockQ**: global dockQ-like score (mpDockQ if multimer; pDockQ if dimer)
 - **average_interface_pae, interface_average_plddt, interface_num_intf_residues**
-- **interface_contact_pairs, interface_score, interface_pDockQ2, interface_ipSAE, interface_LIS**
+- **interface_contact_pairs, interface_score, interface_pDockQ2, interface_ipSAE, interface_LIS, interface_cLIS, interface_iLIS**
+- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean**: AF3 native `contact_probs`, or Humphreys-style AF2 distogram-derived contact probability (`d < 12 A`) when full result pickles retain the `distogram` key. If AF2 result pickles do not contain `distogram` (for example because AlphaPulldown removed it with `--remove_keys_from_pickles`), these columns are empty/NaN and AlphaJudge logs a warning.
 - **interface_hb, interface_sb, interface_sc, interface_area, interface_solv_en**
 
 Exact header is asserted in tests to be consistent across AF2 and AF3 runs.

scripts/extract_pisa_molref.py

@@ -112,9 +112,7 @@ def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
         "molref_dir",
-        nargs="?",
         type=Path,
-        default=Path("/g/kosinski/dima/PycharmProjects/pisa/molref"),
         help="Directory containing PISA molref.idx and molref.rdt",
     )
     args = parser.parse_args()

src/alphajudge/confidence.py

@@ -17,3 +17,5 @@ class Confidence:
     # AF3 only: per-chain-pair ipTM matrix (indexed by chain order).
     # When present, use this for per-interface iptm instead of global iptm.
     chain_pair_iptm: list[list[float]] | None = None
+    contact_prob_matrix: np.ndarray | None = None
+    contact_prob_source: str | None = None

src/alphajudge/contact_probs.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import numpy as np
+
+AF2_DISTOGRAM_CONTACT_CUTOFF = 12.0
+
+
+def contact_probs_from_distogram(
+    logits: np.ndarray,
+    bin_edges: np.ndarray,
+    contact_cutoff: float = AF2_DISTOGRAM_CONTACT_CUTOFF,
+) -> np.ndarray:
+    """
+    Convert AF2 distogram logits to P(distance < contact_cutoff).
+
+    AF2 distogram ``bin_edges`` are finite upper bin boundaries. For scoring,
+    use the lower-bound convention used by AF2 contact-probability analyses:
+    append a 0 A lower bound and include bins whose lower bound is below the
+    cutoff. This follows the Humphreys et al. AF2 contact-probability
+    convention and includes the bin that straddles 12 A in the standard AF2
+    distogram, instead of dropping it because its upper edge is slightly above
+    12 A.
+    """
+    logits_arr = np.asarray(logits)
+    if logits_arr.ndim != 3:
+        raise ValueError(f"distogram logits must be 3D, got shape {logits_arr.shape}")
+
+    edges = np.asarray(bin_edges, dtype=float).ravel()
+    if edges.size != logits_arr.shape[-1] - 1:
+        raise ValueError(
+            f"distogram bin_edges length {edges.size} does not match logits bins "
+            f"{logits_arr.shape[-1]}"
+        )
+
+    work = (
+        logits_arr
+        if np.issubdtype(logits_arr.dtype, np.floating)
+        else logits_arr.astype(np.float32)
+    )
+
+    lower_bounds = np.concatenate([[0.0], edges])
+    clipped_cutoff = float(np.clip(float(contact_cutoff), 3.0, 20.0))
+    n_contact_bins = int(np.count_nonzero(lower_bounds < clipped_cutoff))
+    n_contact_bins = max(1, min(n_contact_bins, work.shape[-1]))
+
+    max_logit = np.max(work, axis=-1)
+    dtype = np.result_type(work.dtype, np.float32)
+    numerator = np.zeros(max_logit.shape, dtype=dtype)
+    denominator = np.zeros(max_logit.shape, dtype=dtype)
+    for bin_idx in range(work.shape[-1]):
+        delta = (work[..., bin_idx] - max_logit).astype(dtype, copy=False)
+        bin_mass = np.exp(delta)
+        denominator += bin_mass
+        if bin_idx < n_contact_bins:
+            numerator += bin_mass
+    return numerator / denominator
+
+
+def summarize_contact_prob_block(
+    matrix: np.ndarray | None,
+    idx1: np.ndarray,
+    idx2: np.ndarray,
+    top_n: int = 10,
+) -> tuple[float, float]:
+    """Summarize one inter-chain block as max and top-N mean."""
+    if matrix is None or idx1.size == 0 or idx2.size == 0:
+        nan = float("nan")
+        return nan, nan
+
+    m = np.asarray(matrix, dtype=float)
+    if m.ndim != 2:
+        nan = float("nan")
+        return nan, nan
+    if idx1.max(initial=-1) >= m.shape[0] or idx2.max(initial=-1) >= m.shape[1]:
+        nan = float("nan")
+        return nan, nan
+
+    block = m[np.ix_(idx1, idx2)].ravel()
+    if (
+        m.shape[0] == m.shape[1]
+        and idx2.max(initial=-1) < m.shape[0]
+        and idx1.max(initial=-1) < m.shape[1]
+    ):
+        reverse = m[np.ix_(idx2, idx1)].T.ravel()
+        both = np.isfinite(block) & np.isfinite(reverse)
+        only_reverse = ~np.isfinite(block) & np.isfinite(reverse)
+        if np.any(both):
+            block[both] = 0.5 * (block[both] + reverse[both])
+        if np.any(only_reverse):
+            block[only_reverse] = reverse[only_reverse]
+
+    finite = block[np.isfinite(block)]
+    if finite.size == 0:
+        nan = float("nan")
+        return nan, nan
+
+    n = max(1, min(int(top_n), finite.size))
+    top = np.partition(finite, finite.size - n)[-n:]
+    return float(np.max(finite)), float(np.mean(top))

src/alphajudge/interface.py

@@ -16,6 +16,7 @@
     shape_complementarity as _scasa_sc,
 )
 from .docking_scores import D0, PDOCKQ, PDOCKQ2
+from .contact_probs import summarize_contact_prob_block
 from .geometry import (
     CHARGED_RES,
     HYDROPHOBIC_RES,
@@ -37,6 +38,7 @@ def __init__(self, chain1, chain2, complex_ctx: Complex):
         if not self.chain1 or not self.chain2:
             # pae_matrix is already a numpy array for memory efficiency
             self._pae = self.c.conf.pae_matrix
+            self._contact_prob = self.c.conf.contact_prob_matrix
             self._rim = self.c._res_index_map
             self._cid = self.c._chain_indices_by_id
             self._cid1_id = ""
@@ -51,6 +53,7 @@ def __init__(self, chain1, chain2, complex_ctx: Complex):
 
         # pae_matrix is already a numpy array for memory efficiency
         self._pae = self.c.conf.pae_matrix
+        self._contact_prob = self.c.conf.contact_prob_matrix
         self._rim = self.c._res_index_map
         self._cid = self.c._chain_indices_by_id
 
@@ -213,6 +216,18 @@ def ilis(self) -> float:
             return 0.0
         return float(math.sqrt(lis * clis))
 
+    @cached_property
+    def contact_probability_scores(self) -> tuple[float, float]:
+        return summarize_contact_prob_block(self._contact_prob, self._idx1, self._idx2)
+
+    @cached_property
+    def contact_prob_max(self) -> float:
+        return self.contact_probability_scores[0]
+
+    @cached_property
+    def contact_prob_top10_mean(self) -> float:
+        return self.contact_probability_scores[1]
+
     @property
     def polar(self) -> float:
         return self._frac(POLAR_RES)

src/alphajudge/meta_score.py

@@ -46,14 +46,13 @@
     1.0,
 )
 
-# Frozen deciles from the benchmark_26 final synchronized best-interface run
-# (final_sync_20260523_225722, after the missing pair-matched predictions were
-# back-filled). Calibrated on POSITIVE (interacting) pairs ONLY: 3,878 AF2/AF3
-# positive rows out of the 7,756-row balanced table. The database-negative
-# re-pairings are deliberately excluded so a new prediction is ranked against
-# the distribution of real interfaces, not against a 50% non-interacting decoy
-# population. Regenerate with
-# `python scripts/freeze_metascore_quantiles.py --label-filter positive`.
+# Frozen deciles from the AlphaJudge interacting reference set. Calibrated on
+# POSITIVE (interacting) pairs ONLY: 3,878 AF2/AF3 positive rows out of the
+# 7,756-row balanced table. The database-negative re-pairings are deliberately
+# excluded so a new prediction is ranked against the distribution of real
+# interfaces, not against a 50% non-interacting decoy population. Regenerate
+# manually with
+# `python test/manual/freeze_metascore_quantiles.py --input-csv ... --label-filter positive`.
 # Values are already oriented so larger is better; e.g. PAE and solvation
 # energy are stored after sign flip.
 BENCHMARK_QUANTILES = {

src/alphajudge/parsers/af2.py

@@ -1,11 +1,22 @@
 from __future__ import annotations
+import gzip
+import logging
+import lzma
+import pickle
 from pathlib import Path
 import numpy as np
 from . import BaseParser, Run
 from ..confidence import Confidence
+from ..contact_probs import (
+    AF2_DISTOGRAM_CONTACT_CUTOFF,
+    contact_probs_from_distogram,
+)
+
+logger = logging.getLogger(__name__)
 
 class AF2Parser(BaseParser):
     name = "af2"
+    _warned_missing_distogram = False
 
     def detect(self, d: Path) -> bool:
         return (d / "ranking_debug.json").exists()
@@ -22,6 +33,7 @@ def load_model(model: str):
             pae_payload = self._read_json(d / f"pae_{model}.json")
             pae = np.array(pae_payload[0]["predicted_aligned_error"], dtype=float)
             max_pae = float(np.nanmax(pae) if pae.size else np.nan)
+            contact_probs = self._load_contact_probs_from_result_pkl(d, model, pae.shape)
 
             # AF2 rankings
             is_multimer = ("iptm+ptm" in rj) and ("iptm" in rj)
@@ -48,5 +60,99 @@ def load_model(model: str):
                 pae_matrix=pae, max_pae=max_pae,
                 iptm=iptm, ptm=ptm, iptm_ptm=iptm_ptm, confidence_score=conf,
                 plddt_residue=plddt,
+                contact_prob_matrix=contact_probs,
+                contact_prob_source=(
+                    "af2_distogram_lb_lt_12A" if contact_probs is not None else None
+                ),
             )
         return Run(order=order, source="af2", load_model=load_model)
+
+    @classmethod
+    def _load_contact_probs_from_result_pkl(
+        cls, d: Path, model: str, expected_shape: tuple[int, int]
+    ) -> np.ndarray | None:
+        result_pkl = cls._find_result_pkl(d, model)
+        if result_pkl is None:
+            return None
+
+        try:
+            payload = cls._read_pickle(result_pkl)
+        except Exception as e:
+            logger.warning(f"could not read AF2 result pickle {result_pkl}: {e}")
+            return None
+
+        if not isinstance(payload, dict):
+            return None
+        distogram = payload.get("distogram")
+        if not isinstance(distogram, dict):
+            if not cls._warned_missing_distogram:
+                logger.warning(
+                    "AF2 result pickle %s has no distogram; AF2 contact-probability "
+                    "columns will be empty/NaN. Full AlphaPulldown result pickles are "
+                    "required; disable --remove_keys_from_pickles to retain distograms.",
+                    result_pkl,
+                )
+                cls._warned_missing_distogram = True
+            else:
+                logger.debug(
+                    "AF2 result pickle %s has no distogram; contact scores unavailable.",
+                    result_pkl,
+                )
+            return None
+        logits = distogram.get("logits")
+        bin_edges = distogram.get("bin_edges")
+        if logits is None or bin_edges is None:
+            return None
+
+        try:
+            contact_probs = contact_probs_from_distogram(
+                np.asarray(logits),
+                np.asarray(bin_edges),
+                AF2_DISTOGRAM_CONTACT_CUTOFF,
+            )
+        except Exception as e:
+            logger.warning(f"could not derive AF2 contact probabilities from {result_pkl}: {e}")
+            return None
+
+        if contact_probs.shape != expected_shape:
+            logger.warning(
+                f"AF2 contact probability shape {contact_probs.shape} != expected "
+                f"{expected_shape}; skipping contact probabilities."
+            )
+            return None
+
+        return contact_probs
+
+    @staticmethod
+    def _find_result_pkl(d: Path, model: str) -> Path | None:
+        suffixes = ("", ".gz", ".xz")
+        stems = [
+            d / f"result_{model}.pkl",
+            d / model / "result.pkl",
+            d / model / f"result_{model}.pkl",
+        ]
+        candidates: list[Path] = []
+        for stem in stems:
+            candidates.extend(stem.with_name(stem.name + suffix) for suffix in suffixes)
+        candidates.extend(sorted(d.glob(f"result*{model}*.pkl*")))
+        candidates.extend(sorted((d / model).glob("result*.pkl*")) if (d / model).is_dir() else [])
+
+        seen: set[Path] = set()
+        for candidate in candidates:
+            if candidate in seen:
+                continue
+            seen.add(candidate)
+            if candidate.exists():
+                return candidate
+        return None
+
+    @staticmethod
+    def _read_pickle(path: Path):
+        if path.suffix == ".gz":
+            with gzip.open(path, "rb") as f:
+                return pickle.load(f)
+        if path.suffix == ".xz":
+            with lzma.open(path, "rb") as f:
+                return pickle.load(f)
+        with path.open("rb") as f:
+            return pickle.load(f)