From f8b1f54305d152f9b6565916b44b2e3904003f18 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Thu, 9 Jul 2026 13:39:55 +0200 Subject: [PATCH 1/6] Add interface contact probability scores --- README.md | 9 +- .../evaluate_contact_probability_scores.py | 384 ++++++++++++++++++ src/alphajudge/confidence.py | 2 + src/alphajudge/contact_probs.py | 95 +++++ src/alphajudge/interface.py | 19 + src/alphajudge/parsers/af2.py | 97 +++++ src/alphajudge/parsers/af3.py | 95 +++++ src/alphajudge/runner.py | 4 + test/test_parsers_and_runner.py | 124 +++++- 9 files changed, 824 insertions(+), 5 deletions(-) create mode 100755 scripts/evaluate_contact_probability_scores.py create mode 100644 src/alphajudge/contact_probs.py diff --git a/README.md b/README.md index bc95642e..bee33d6c 100644 --- a/README.md +++ b/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_expected_contacts`, `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, interface_expected_contacts**: AF3 native `contact_probs`, or AF2 distogram-derived `P(distance < 8 A)` when full result pickles are available - **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. diff --git a/scripts/evaluate_contact_probability_scores.py b/scripts/evaluate_contact_probability_scores.py new file mode 100755 index 00000000..81335dd9 --- /dev/null +++ b/scripts/evaluate_contact_probability_scores.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +"""Evaluate AlphaJudge contact-probability scores on a labelled prediction tree. + +The benchmark tree is expected to look like: + + ////predictions/ + +Rows are written without modifying the prediction directories. By default the +script scores all chain pairs, including pairs without coordinate-detected +contacts, because contact probability is a model-confidence quantity and can be +defined even when the sampled structure has no interface. +""" + +from __future__ import annotations + +import argparse +import csv +import logging +import math +from pathlib import Path +from typing import Iterable + +import numpy as np + +from alphajudge.complex import Complex +from alphajudge.interface import Interface +from alphajudge.meta_score import interface_meta_score +from alphajudge.parsers import pick_parser + +logger = logging.getLogger("contact_probability_benchmark") + +DEFAULT_ROOT = Path( + "/g/transform/kosinski/dima/IntAct_BioGRID_STRING/benchmark_26/predictions" +) +DEFAULT_SCORES = ( + "interface_contact_prob_max", + "interface_contact_prob_top10_mean", + "interface_expected_contacts", + "iptm", + "iptm_ptm", + "confidence_score", + "pDockQ/mpDockQ", + "interface_pDockQ2", + "interface_ipSAE", + "interface_LIS", + "interface_cLIS", + "interface_iLIS", + "interface_contact_pairs", + "average_interface_pae", +) +LOWER_IS_BETTER = {"average_interface_pae"} + + +def _safe_float(value) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + return float("nan") + return parsed if math.isfinite(parsed) else float("nan") + + +def _infer_metadata(run_dir: Path, root: Path) -> dict[str, str | int]: + parts = run_dir.resolve().relative_to(root.resolve()).parts + organism = parts[0] if len(parts) > 0 else "" + backend = parts[1] if len(parts) > 1 else "" + pair_set = parts[2] if len(parts) > 2 else "" + if pair_set.startswith("pos"): + label = 1 + elif pair_set.startswith("neg"): + label = 0 + else: + label = -1 + return { + "organism": organism, + "backend": backend, + "pair_set": pair_set, + "label": label, + "pair_id": run_dir.name, + "source_dir": str(run_dir.resolve()), + } + + +def _discover_run_dirs( + root: Path, + limit_per_group: int | None = None, + *, + organism_filter: str | None = None, + backend_filter: str | None = None, + pair_set_filter: str | None = None, +) -> list[Path]: + candidates: list[Path] = [] + for predictions_dir in sorted(root.glob("*/*/*/predictions")): + if not predictions_dir.is_dir(): + continue + try: + organism, backend, pair_set, _ = predictions_dir.relative_to(root).parts + except ValueError: + organism, backend, pair_set = "", "", "" + if organism_filter and organism != organism_filter: + continue + if backend_filter and backend != backend_filter: + continue + if pair_set_filter and pair_set != pair_set_filter: + continue + group_hits: list[Path] = [] + for child in sorted(p for p in predictions_dir.iterdir() if p.is_dir()): + try: + pick_parser(child) + except Exception: + continue + group_hits.append(child) + if limit_per_group is not None and len(group_hits) >= limit_per_group: + break + candidates.extend(group_hits) + + if candidates: + return candidates + + for child in sorted(p for p in root.rglob("*") if p.is_dir()): + try: + pick_parser(child) + except Exception: + continue + metadata = _infer_metadata(child, root) + if organism_filter and metadata["organism"] != organism_filter: + continue + if backend_filter and metadata["backend"] != backend_filter: + continue + if pair_set_filter and metadata["pair_set"] != pair_set_filter: + continue + candidates.append(child) + return candidates[:limit_per_group] if limit_per_group is not None else candidates + + +def _interface_label(iface: Interface) -> str: + return f"{iface.chain1[0].get_parent().id}_{iface.chain2[0].get_parent().id}" + + +def _all_chain_pair_interfaces(comp: Complex) -> Iterable[Interface]: + for i in range(len(comp._chains)): # noqa: SLF001 - benchmark helper + for j in range(i + 1, len(comp._chains)): # noqa: SLF001 + yield Interface(comp._chains[i], comp._chains[j], comp) # noqa: SLF001 + + +def _row_for_interface(job: str, model: str, confidence, global_score: float, iface: Interface) -> dict: + pd2, _ = iface.pDockQ2() + iptm_val = iface.iptm_chainpair if iface.iptm_chainpair is not None else confidence.iptm + row = { + "jobs": job, + "model_used": model, + "interface": _interface_label(iface), + "iptm_ptm": float(confidence.iptm_ptm) if confidence.iptm_ptm is not None else float("nan"), + "iptm": float(iptm_val) if iptm_val is not None else float("nan"), + "ptm": float(confidence.ptm) if confidence.ptm is not None else float("nan"), + "confidence_score": float(confidence.confidence_score) + if confidence.confidence_score is not None + else float("nan"), + "pDockQ/mpDockQ": global_score, + "average_interface_pae": iface.average_interface_pae, + "interface_average_plddt": iface.average_interface_plddt, + "interface_num_intf_residues": iface.num_intf_residues, + "interface_contact_pairs": iface.contact_pairs, + "interface_contact_prob_source": confidence.contact_prob_source or "", + "interface_contact_prob_max": iface.contact_prob_max, + "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, + "interface_expected_contacts": iface.expected_contacts, + "interface_score": iface.score_complex, + "interface_pDockQ2": pd2, + "interface_ipSAE": iface.ipsae(), + "interface_LIS": iface.lis(), + "interface_cLIS": iface.clis(), + "interface_iLIS": iface.ilis(), + } + row["interface_meta_score"] = interface_meta_score(row) + return row + + +def score_run( + run_dir: Path, + root: Path, + *, + contact_thresh: float, + pae_filter: float, + ipsae_pae_cutoff: float, + models_to_analyse: str, + chain_pairs: str, +) -> tuple[list[dict], str | None]: + metadata = _infer_metadata(run_dir, root) + try: + parser = pick_parser(run_dir) + run = parser.parse_run(run_dir) + models = [run.order[0]] if models_to_analyse == "best" else list(run.order) + except Exception as e: + return [], f"{run_dir}: parser setup failed: {e}" + + rows: list[dict] = [] + for model in models: + try: + structure, confidence = run.load_model(model) + comp = Complex(structure, confidence, contact_thresh, pae_filter, ipsae_pae_cutoff) + global_score = ( + comp.mpDockQ + if comp.num_chains > 2 + else (comp.interfaces[0].pDockQ if comp.interfaces else float("nan")) + ) + interfaces = ( + list(_all_chain_pair_interfaces(comp)) + if chain_pairs == "all" + else list(comp.interfaces) + ) + for iface in interfaces: + if chain_pairs == "detected" and iface.num_intf_residues == 0: + continue + if ( + iface.num_intf_residues > 0 + and math.isfinite(iface.average_interface_pae) + and iface.average_interface_pae > pae_filter + ): + continue + row = {**metadata, "parser": parser.name} + row.update(_row_for_interface(run_dir.name, model, confidence, global_score, iface)) + rows.append(row) + except Exception as e: + return rows, f"{run_dir} {model}: scoring failed: {e}" + return rows, None + + +def _write_rows(path: Path, rows: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if not rows: + path.write_text("") + return + + fieldnames: list[str] = [] + seen: set[str] = set() + for row in rows: + for key in row: + if key not in seen: + seen.add(key) + fieldnames.append(key) + + with path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + for row in rows: + writer.writerow({key: row.get(key, "") for key in fieldnames}) + + +def _average_ranks(values: np.ndarray) -> np.ndarray: + order = np.argsort(values) + ranks = np.empty(values.size, dtype=float) + i = 0 + while i < values.size: + j = i + 1 + while j < values.size and values[order[j]] == values[order[i]]: + j += 1 + ranks[order[i:j]] = (i + 1 + j) / 2.0 + i = j + return ranks + + +def _auroc(labels: np.ndarray, scores: np.ndarray) -> float: + n_pos = int(np.count_nonzero(labels == 1)) + n_neg = int(np.count_nonzero(labels == 0)) + if n_pos == 0 or n_neg == 0: + return float("nan") + ranks = _average_ranks(scores) + sum_pos = float(np.sum(ranks[labels == 1])) + return (sum_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg) + + +def _average_precision(labels: np.ndarray, scores: np.ndarray) -> float: + n_pos = int(np.count_nonzero(labels == 1)) + if n_pos == 0: + return float("nan") + order = np.argsort(-scores) + y = labels[order] + precision = np.cumsum(y == 1) / (np.arange(y.size) + 1) + return float(np.sum(precision[y == 1]) / n_pos) + + +def benchmark_metrics(rows: list[dict], score_names: Iterable[str]) -> list[dict]: + groups: list[tuple[str, str, list[dict]]] = [("all", "all", rows)] + for backend in sorted({str(r.get("backend", "")) for r in rows}): + groups.append((backend, "all", [r for r in rows if r.get("backend") == backend])) + for organism in sorted({str(r.get("organism", "")) for r in rows}): + groups.append(("all", organism, [r for r in rows if r.get("organism") == organism])) + for backend in sorted({str(r.get("backend", "")) for r in rows}): + for organism in sorted({str(r.get("organism", "")) for r in rows}): + subset = [ + r + for r in rows + if r.get("backend") == backend and r.get("organism") == organism + ] + groups.append((backend, organism, subset)) + + metrics: list[dict] = [] + for backend, organism, subset in groups: + labels_all = np.asarray([int(r.get("label", -1)) for r in subset], dtype=int) + valid_labels = (labels_all == 0) | (labels_all == 1) + for score_name in score_names: + raw_scores = np.asarray([_safe_float(r.get(score_name)) for r in subset], dtype=float) + valid = valid_labels & np.isfinite(raw_scores) + labels = labels_all[valid] + scores = raw_scores[valid] + if score_name in LOWER_IS_BETTER: + scores = -scores + metrics.append( + { + "backend": backend, + "organism": organism, + "score": score_name, + "n": int(valid.sum()), + "n_pos": int(np.count_nonzero(labels == 1)), + "n_neg": int(np.count_nonzero(labels == 0)), + "auroc": _auroc(labels, scores), + "average_precision": _average_precision(labels, scores), + } + ) + return metrics + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--predictions-root", default=str(DEFAULT_ROOT)) + parser.add_argument("--out", required=True, help="Output per-interface/per-chain-pair CSV.") + parser.add_argument("--metrics-out", required=True, help="Output AUROC/AP summary CSV.") + parser.add_argument("--models-to-analyse", default="best", choices=("best", "all")) + parser.add_argument("--chain-pairs", default="all", choices=("all", "detected")) + parser.add_argument("--contact-thresh", type=float, default=8.0) + parser.add_argument("--pae-filter", type=float, default=100.0) + parser.add_argument("--ipsae-pae-cutoff", type=float, default=10.0) + parser.add_argument("--limit-per-group", type=int, default=None) + parser.add_argument("--organism", default=None) + parser.add_argument("--backend", default=None, choices=("af2", "af3")) + parser.add_argument("--pair-set", default=None, choices=("pos_pairs", "neg_pairs")) + parser.add_argument("--scores", nargs="*", default=list(DEFAULT_SCORES)) + parser.add_argument("--log-level", default="INFO") + args = parser.parse_args() + + logging.basicConfig(level=getattr(logging, args.log_level.upper()), format="%(levelname)s: %(message)s") + + root = Path(args.predictions_root) + run_dirs = _discover_run_dirs( + root, + args.limit_per_group, + organism_filter=args.organism, + backend_filter=args.backend, + pair_set_filter=args.pair_set, + ) + logger.info("found %d run directories under %s", len(run_dirs), root) + + rows: list[dict] = [] + errors: list[str] = [] + for i, run_dir in enumerate(run_dirs, start=1): + logger.info("scoring %d/%d %s", i, len(run_dirs), run_dir) + scored, error = score_run( + run_dir, + root, + contact_thresh=args.contact_thresh, + pae_filter=args.pae_filter, + ipsae_pae_cutoff=args.ipsae_pae_cutoff, + models_to_analyse=args.models_to_analyse, + chain_pairs=args.chain_pairs, + ) + rows.extend(scored) + if error: + logger.warning(error) + errors.append(error) + + _write_rows(Path(args.out), rows) + _write_rows(Path(args.metrics_out), benchmark_metrics(rows, args.scores)) + + if errors: + errors_path = Path(args.metrics_out).with_name(Path(args.metrics_out).stem + "_errors.txt") + errors_path.write_text("\n".join(errors) + "\n") + logger.warning("wrote %d errors to %s", len(errors), errors_path) + logger.info("wrote %d score rows to %s", len(rows), args.out) + logger.info("wrote metric summary to %s", args.metrics_out) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/alphajudge/confidence.py b/src/alphajudge/confidence.py index 770b6a82..58db9d04 100644 --- a/src/alphajudge/confidence.py +++ b/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 diff --git a/src/alphajudge/contact_probs.py b/src/alphajudge/contact_probs.py new file mode 100644 index 00000000..704e5310 --- /dev/null +++ b/src/alphajudge/contact_probs.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import numpy as np + +AF2_DISTOGRAM_CONTACT_CUTOFF = 8.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). + + This follows the AlphaPulldown diagnostics fallback convention: append an + infinite final upper bound, clip the requested cutoff to 3-20 A, softmax + over bins, then sum bins whose upper bound is strictly below the cutoff. + """ + 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]}" + ) + + dtype = logits_arr.dtype if np.issubdtype(logits_arr.dtype, np.floating) else np.float32 + work = logits_arr.astype(dtype, copy=False) + shifted = work - np.max(work, axis=-1, keepdims=True) + np.exp(shifted, out=shifted) + + upper_bounds = np.concatenate([edges, [np.inf]]) + clipped_cutoff = float(np.clip(float(contact_cutoff), 3.0, 20.0)) + contact_bins = upper_bounds < clipped_cutoff + if not np.any(contact_bins): + contact_bins[0] = True + + numerator = shifted[..., contact_bins].sum(axis=-1) + denominator = shifted.sum(axis=-1) + return numerator / denominator + + +def symmetrize_contact_probs(matrix: np.ndarray) -> tuple[np.ndarray, float]: + """Return a nan-aware symmetric copy and the maximum pre-symmetry delta.""" + m = np.asarray(matrix, dtype=float) + if m.ndim != 2 or m.shape[0] != m.shape[1]: + raise ValueError(f"contact probability matrix must be square, got {m.shape}") + + mt = m.T + sym = (m + mt) / 2.0 + + left_nan = np.isnan(m) + right_nan = np.isnan(mt) + sym[left_nan & ~right_nan] = mt[left_nan & ~right_nan] + sym[right_nan & ~left_nan] = m[right_nan & ~left_nan] + sym[left_nan & right_nan] = np.nan + + delta = np.abs(m - mt) + finite = np.isfinite(delta) + max_delta = float(np.nanmax(delta[finite])) if np.any(finite) else 0.0 + return sym, max_delta + + +def summarize_contact_prob_block( + matrix: np.ndarray | None, + idx1: np.ndarray, + idx2: np.ndarray, + top_n: int = 10, +) -> tuple[float, float, float]: + """Summarize inter-chain contact probabilities as max, top-N mean, and sum.""" + if matrix is None or idx1.size == 0 or idx2.size == 0: + nan = float("nan") + return nan, nan, nan + + m = np.asarray(matrix, dtype=float) + if m.ndim != 2: + nan = float("nan") + return nan, nan, nan + if idx1.max(initial=-1) >= m.shape[0] or idx2.max(initial=-1) >= m.shape[1]: + nan = float("nan") + return nan, nan, nan + + block = m[np.ix_(idx1, idx2)].ravel() + finite = block[np.isfinite(block)] + if finite.size == 0: + nan = float("nan") + return nan, 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)), float(np.sum(finite)) diff --git a/src/alphajudge/interface.py b/src/alphajudge/interface.py index 485bcfda..d847885f 100644 --- a/src/alphajudge/interface.py +++ b/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,22 @@ def ilis(self) -> float: return 0.0 return float(math.sqrt(lis * clis)) + @cached_property + def contact_probability_scores(self) -> tuple[float, 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] + + @cached_property + def expected_contacts(self) -> float: + return self.contact_probability_scores[2] + @property def polar(self) -> float: return self._frac(POLAR_RES) diff --git a/src/alphajudge/parsers/af2.py b/src/alphajudge/parsers/af2.py index dd1d1cec..17c08914 100644 --- a/src/alphajudge/parsers/af2.py +++ b/src/alphajudge/parsers/af2.py @@ -1,8 +1,19 @@ 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, + symmetrize_contact_probs, +) + +logger = logging.getLogger(__name__) class AF2Parser(BaseParser): name = "af2" @@ -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,90 @@ 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_lt_8A" 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): + 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 + + contact_probs, max_delta = symmetrize_contact_probs(contact_probs) + if max_delta > 1e-6: + logger.warning( + f"AF2 contact probabilities in {result_pkl} were asymmetric " + f"(max abs delta {max_delta:.3g}); symmetrized." + ) + 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) diff --git a/src/alphajudge/parsers/af3.py b/src/alphajudge/parsers/af3.py index 389d9dc7..0250a018 100644 --- a/src/alphajudge/parsers/af3.py +++ b/src/alphajudge/parsers/af3.py @@ -6,6 +6,7 @@ import numpy as np from . import BaseParser, Run from ..confidence import Confidence +from ..contact_probs import symmetrize_contact_probs logger = logging.getLogger(__name__) @@ -53,12 +54,15 @@ def load_model(model: str): chain_pair_iptm = [list(r) if isinstance(r, (list, tuple)) else [] for r in chain_pair_iptm_raw] pae, max_pae = self._normalize_pae_af3(matrix, chains, cid) + contact_probs = self._normalize_contact_probs_af3(matrix, chains, cid) plddt = self._plddt(chains, rim) return struct, Confidence( pae_matrix=pae, max_pae=max_pae, iptm=iptm, ptm=ptm, iptm_ptm=iptm_ptm, confidence_score=ranking_score, plddt_residue=plddt, chain_pair_iptm=chain_pair_iptm, + contact_prob_matrix=contact_probs, + contact_prob_source="af3_contact_probs" if contact_probs is not None else None, ) return Run(order=order, source="af3", load_model=load_model) @@ -252,3 +256,94 @@ def _normalize_pae_af3(matrix: dict, chains, cid) -> tuple[np.ndarray, float]: ) return pae, float(max_pae) + + @classmethod + def _normalize_contact_probs_af3(cls, matrix: dict, chains, cid) -> np.ndarray | None: + raw = matrix.get("contact_probs") + if raw is None: + return None + + expected_shape = (sum(len(cid[c.id]) for c in chains),) * 2 + contact_probs = np.asarray(raw, dtype=float) + if contact_probs.ndim != 2: + logger.warning( + f"contact_probs must be a 2D matrix, got shape {contact_probs.shape}; " + "skipping contact probabilities." + ) + return None + + if contact_probs.shape == expected_shape: + aligned = contact_probs + else: + aligned = cls._align_token_pair_matrix_to_residues( + contact_probs, + matrix.get("token_chain_ids"), + chains, + cid, + expected_shape, + ) + if aligned is None: + logger.warning( + f"contact_probs shape {contact_probs.shape} != expected {expected_shape}; " + "skipping contact probabilities." + ) + return None + + aligned, max_delta = symmetrize_contact_probs(aligned) + if max_delta > 1e-6: + logger.warning( + f"AF3 contact_probs were asymmetric (max abs delta {max_delta:.3g}); " + "symmetrized." + ) + return aligned + + @staticmethod + def _align_token_pair_matrix_to_residues( + token_matrix: np.ndarray, + token_chain_ids, + chains, + cid, + expected_shape: tuple[int, int], + ) -> np.ndarray | None: + if not isinstance(token_chain_ids, (list, tuple)): + return None + if len(token_chain_ids) != token_matrix.shape[0] or token_matrix.shape[0] != token_matrix.shape[1]: + return None + + ids = [str(x) for x in token_chain_ids] + seen: list[str] = [] + for chain_id in ids: + if chain_id not in seen: + seen.append(chain_id) + + token_indices_by_chain: dict[str, np.ndarray] = {} + for i, chain in enumerate(chains): + residue_indices = np.asarray(cid.get(chain.id, []), dtype=int) + if residue_indices.size == 0: + token_indices_by_chain[chain.id] = np.array([], dtype=int) + continue + + exact = np.asarray([k for k, chain_id in enumerate(ids) if chain_id == str(chain.id)], dtype=int) + ordered = ( + np.asarray([k for k, chain_id in enumerate(ids) if chain_id == seen[i]], dtype=int) + if i < len(seen) + else np.array([], dtype=int) + ) + token_indices = exact if exact.size == residue_indices.size else ordered + if token_indices.size != residue_indices.size: + return None + token_indices_by_chain[chain.id] = token_indices + + residue_matrix = np.full(expected_shape, np.nan, dtype=float) + for chi in chains: + ri = np.asarray(cid.get(chi.id, []), dtype=int) + ti = token_indices_by_chain.get(chi.id, np.array([], dtype=int)) + if ri.size == 0: + continue + for chj in chains: + rj = np.asarray(cid.get(chj.id, []), dtype=int) + tj = token_indices_by_chain.get(chj.id, np.array([], dtype=int)) + if rj.size == 0: + continue + residue_matrix[np.ix_(ri, rj)] = token_matrix[np.ix_(ti, tj)] + return residue_matrix diff --git a/src/alphajudge/runner.py b/src/alphajudge/runner.py index 7b845871..2cd838e9 100644 --- a/src/alphajudge/runner.py +++ b/src/alphajudge/runner.py @@ -120,6 +120,10 @@ def process( "interface_hydrophobic": iface.hydrophobic, "interface_charged": iface.charged, "interface_contact_pairs": iface.contact_pairs, + "interface_contact_prob_source": confidence.contact_prob_source or "", + "interface_contact_prob_max": iface.contact_prob_max, + "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, + "interface_expected_contacts": iface.expected_contacts, "interface_score": iface.score_complex, "interface_pDockQ2": pd2, "interface_ipSAE": iface.ipsae(), diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index 882ea4b8..ac8051fe 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -4,6 +4,7 @@ import json import logging import math +import pickle import shutil import subprocess from pathlib import Path @@ -13,7 +14,9 @@ import pytest from alphajudge.parsers import pick_parser +from alphajudge.parsers.af2 import AF2Parser from alphajudge.parsers.af3 import AF3Parser +from alphajudge.contact_probs import contact_probs_from_distogram from alphajudge.runner import process, process_many @@ -37,6 +40,10 @@ "interface_hydrophobic", "interface_charged", "interface_contact_pairs", + "interface_contact_prob_source", + "interface_contact_prob_max", + "interface_contact_prob_top10_mean", + "interface_expected_contacts", "interface_score", "interface_pDockQ2", "interface_ipSAE", @@ -52,7 +59,12 @@ } # columns that are expected to be numeric (but may be NaN) -EXPECTED_NUMERIC_COLUMNS = EXPECTED_OUTPUT_COLUMNS - {"jobs", "model_used", "interface"} +EXPECTED_NUMERIC_COLUMNS = EXPECTED_OUTPUT_COLUMNS - { + "jobs", + "model_used", + "interface", + "interface_contact_prob_source", +} # ------------------------- @@ -390,6 +402,71 @@ def test_clis_ilis_math_is_deterministic(): assert iface.ilis() == 0.0 +def test_contact_probability_scores_math_is_deterministic(): + """ + Contact-probability summaries are computed over all residue pairs between + the two chains: max, mean of the ten largest values, and expected contacts + as the sum of probabilities. + """ + from alphajudge.interface import Interface + + iface = object.__new__(Interface) + iface._idx1 = np.array([0, 1]) + iface._idx2 = np.array([2, 3]) + iface._contact_prob = np.array( + [ + [0.0, 0.0, 0.8, 0.4], + [0.0, 0.0, 0.2, 0.1], + [0.8, 0.2, 0.0, 0.0], + [0.4, 0.1, 0.0, 0.0], + ] + ) + + assert nearly_equal(iface.contact_prob_max, 0.8) + assert nearly_equal(iface.contact_prob_top10_mean, 0.375) + assert nearly_equal(iface.expected_contacts, 1.5) + + missing = object.__new__(Interface) + missing._idx1 = np.array([0, 1]) + missing._idx2 = np.array([2, 3]) + missing._contact_prob = None + assert math.isnan(missing.contact_prob_max) + assert math.isnan(missing.contact_prob_top10_mean) + assert math.isnan(missing.expected_contacts) + + +def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: Path): + """ + AF2 contact probabilities are distogram softmax mass for bins with upper + bound strictly below 8 A, matching AlphaPulldown's diagnostics fallback. + """ + probs = np.array( + [ + [[0.70, 0.20, 0.05, 0.05], [0.10, 0.20, 0.30, 0.40]], + [[0.20, 0.30, 0.10, 0.40], [0.05, 0.05, 0.20, 0.70]], + ], + dtype=float, + ) + logits = np.log(probs) + bin_edges = np.array([4.0, 8.0, 12.0]) + + direct = contact_probs_from_distogram(logits, bin_edges) + expected_asym = np.array([[0.70, 0.10], [0.20, 0.05]]) + assert np.allclose(direct, expected_asym) + + run_dir = tmp_path / "af2_result" + run_dir.mkdir() + with (run_dir / "result_model_1.pkl").open("wb") as f: + pickle.dump({"distogram": {"logits": logits, "bin_edges": bin_edges}}, f) + + parsed = AF2Parser._load_contact_probs_from_result_pkl( + run_dir, "model_1", expected_shape=(2, 2) + ) + assert parsed is not None + expected_sym = np.array([[0.70, 0.15], [0.15, 0.05]]) + assert np.allclose(parsed, expected_sym) + + # ------------------------- # AF3 runner: best/all + score checks # ------------------------- @@ -443,9 +520,54 @@ def test_af3_parser_accepts_alphapulldown_layout(af3_dir_src: Path): _, conf = run.load_model(run.order[0]) assert conf.pae_matrix.shape[0] == len(conf.plddt_residue) + assert conf.contact_prob_matrix is not None + assert conf.contact_prob_matrix.shape == conf.pae_matrix.shape + assert conf.contact_prob_source == "af3_contact_probs" assert np.isfinite(conf.confidence_score) +def test_af3_contact_probability_scores_match_raw_contact_probs( + tmp_path: Path, af3_dir_src: Path +): + from alphajudge.complex import Complex + + parser = pick_parser(af3_dir_src) + run = parser.parse_run(af3_dir_src) + best_model = run.order[0] + structure, conf = run.load_model(best_model) + + raw = _load_json(af3_dir_src / best_model / "confidences.json") + raw_contact_probs = np.asarray(raw["contact_probs"], dtype=float) + raw_sym = 0.5 * (raw_contact_probs + raw_contact_probs.T) + + assert conf.contact_prob_matrix is not None + assert np.allclose(conf.contact_prob_matrix, raw_sym) + + comp = Complex(structure, conf, 8.0, 100.0, 10.0) + assert comp.interfaces + iface = comp.interfaces[0] + label = f"{iface.chain1[0].get_parent().id}_{iface.chain2[0].get_parent().id}" + + af3_dir = copy_run_dir(af3_dir_src, tmp_path) + process( + str(af3_dir), + 8.0, + 100.0, + "best", + 10.0, + per_run_csv_name="interfaces_contact_probs.csv", + skip_pae_png=True, + skip_biophysical_scores=True, + ) + + rows = read_csv_rows(af3_dir / "interfaces_contact_probs.csv") + row = next(r for r in rows if r["interface"] == label) + assert row["interface_contact_prob_source"] == "af3_contact_probs" + assert nearly_equal(row["interface_contact_prob_max"], iface.contact_prob_max) + assert nearly_equal(row["interface_contact_prob_top10_mean"], iface.contact_prob_top10_mean) + assert nearly_equal(row["interface_expected_contacts"], iface.expected_contacts) + + @pytest.mark.parametrize("models_to_analyse", ["best", "all"]) def test_af3_runner_outputs_have_expected_scores(tmp_path: Path, af3_dir_src: Path, models_to_analyse: str): af3_dir = copy_run_dir(af3_dir_src, tmp_path) From 85b129dad234911226820ea7ee473aef8a352607 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Thu, 9 Jul 2026 19:13:50 +0200 Subject: [PATCH 2/6] Address contact probability review feedback --- README.md | 4 +- .../evaluate_contact_probability_scores.py | 91 ++++++++++++++++--- src/alphajudge/contact_probs.py | 86 +++++++++--------- src/alphajudge/interface.py | 6 +- src/alphajudge/parsers/af2.py | 12 +-- src/alphajudge/parsers/af3.py | 72 +++++++++++++-- src/alphajudge/runner.py | 1 + test/test_parsers_and_runner.py | 46 +++++++++- 8 files changed, 239 insertions(+), 79 deletions(-) diff --git a/README.md b/README.md index bee33d6c..83865f8e 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,7 @@ process_many( ) ``` -Key outputs per interface include: `average_interface_pae`, `interface_average_plddt`, `interface_contact_pairs`, `interface_contact_prob_max`, `interface_contact_prob_top10_mean`, `interface_expected_contacts`, `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`. +Key outputs per interface include: `average_interface_pae`, `interface_average_plddt`, `interface_contact_pairs`, `interface_contact_prob_max`, `interface_contact_prob_top10_mean`, `interface_expected_contacts`, `interface_confident_contacts`, `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`. --- @@ -208,7 +208,7 @@ AlphaJudge writes `interfaces.csv` with one row per interface (and includes the - **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_cLIS, interface_iLIS** -- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean, interface_expected_contacts**: AF3 native `contact_probs`, or AF2 distogram-derived `P(distance < 8 A)` when full result pickles are available +- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean, interface_expected_contacts, interface_confident_contacts**: AF3 native `contact_probs`, or AF2 distogram-derived contact probability when full result pickles retain the `distogram` key. `interface_confident_contacts` counts inter-chain residue pairs with contact probability >= 0.5. - **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. diff --git a/scripts/evaluate_contact_probability_scores.py b/scripts/evaluate_contact_probability_scores.py index 81335dd9..053ee2f2 100755 --- a/scripts/evaluate_contact_probability_scores.py +++ b/scripts/evaluate_contact_probability_scores.py @@ -17,11 +17,17 @@ import csv import logging import math +import sys from pathlib import Path from typing import Iterable import numpy as np +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +if SRC_ROOT.exists(): + sys.path.insert(0, str(SRC_ROOT)) + from alphajudge.complex import Complex from alphajudge.interface import Interface from alphajudge.meta_score import interface_meta_score @@ -36,6 +42,7 @@ "interface_contact_prob_max", "interface_contact_prob_top10_mean", "interface_expected_contacts", + "interface_confident_contacts", "iptm", "iptm_ptm", "confidence_score", @@ -164,6 +171,7 @@ def _row_for_interface(job: str, model: str, confidence, global_score: float, if "interface_contact_prob_max": iface.contact_prob_max, "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, "interface_expected_contacts": iface.expected_contacts, + "interface_confident_contacts": iface.confident_contacts, "interface_score": iface.score_complex, "interface_pDockQ2": pd2, "interface_ipSAE": iface.ipsae(), @@ -184,16 +192,17 @@ def score_run( ipsae_pae_cutoff: float, models_to_analyse: str, chain_pairs: str, -) -> tuple[list[dict], str | None]: +) -> tuple[list[dict], list[str]]: metadata = _infer_metadata(run_dir, root) try: parser = pick_parser(run_dir) run = parser.parse_run(run_dir) models = [run.order[0]] if models_to_analyse == "best" else list(run.order) except Exception as e: - return [], f"{run_dir}: parser setup failed: {e}" + return [], [f"{run_dir}: parser setup failed: {e}"] rows: list[dict] = [] + errors: list[str] = [] for model in models: try: structure, confidence = run.load_model(model) @@ -221,8 +230,9 @@ def score_run( row.update(_row_for_interface(run_dir.name, model, confidence, global_score, iface)) rows.append(row) except Exception as e: - return rows, f"{run_dir} {model}: scoring failed: {e}" - return rows, None + errors.append(f"{run_dir} {model}: scoring failed: {e}") + continue + return rows, errors def _write_rows(path: Path, rows: list[dict]) -> None: @@ -275,11 +285,32 @@ def _average_precision(labels: np.ndarray, scores: np.ndarray) -> float: return float("nan") order = np.argsort(-scores) y = labels[order] - precision = np.cumsum(y == 1) / (np.arange(y.size) + 1) - return float(np.sum(precision[y == 1]) / n_pos) + sorted_scores = scores[order] + ap = 0.0 + tp = 0 + seen = 0 + i = 0 + while i < y.size: + j = i + 1 + while j < y.size and sorted_scores[j] == sorted_scores[i]: + j += 1 + group_pos = int(np.count_nonzero(y[i:j] == 1)) + tp += group_pos + seen = j + if group_pos: + ap += (group_pos / n_pos) * (tp / seen) + i = j + return float(ap) -def benchmark_metrics(rows: list[dict], score_names: Iterable[str]) -> list[dict]: + +def benchmark_metrics( + rows: list[dict], + score_names: Iterable[str], + *, + row_mode: str = "complete-case", +) -> list[dict]: + score_names = list(score_names) groups: list[tuple[str, str, list[dict]]] = [("all", "all", rows)] for backend in sorted({str(r.get("backend", "")) for r in rows}): groups.append((backend, "all", [r for r in rows if r.get("backend") == backend])) @@ -298,9 +329,21 @@ def benchmark_metrics(rows: list[dict], score_names: Iterable[str]) -> list[dict for backend, organism, subset in groups: labels_all = np.asarray([int(r.get("label", -1)) for r in subset], dtype=int) valid_labels = (labels_all == 0) | (labels_all == 1) + raw_by_score = { + score_name: np.asarray([_safe_float(r.get(score_name)) for r in subset], dtype=float) + for score_name in score_names + } + complete_valid = valid_labels.copy() + if row_mode == "complete-case": + for raw_scores in raw_by_score.values(): + complete_valid &= np.isfinite(raw_scores) for score_name in score_names: - raw_scores = np.asarray([_safe_float(r.get(score_name)) for r in subset], dtype=float) - valid = valid_labels & np.isfinite(raw_scores) + raw_scores = raw_by_score[score_name] + valid = ( + complete_valid + if row_mode == "complete-case" + else (valid_labels & np.isfinite(raw_scores)) + ) labels = labels_all[valid] scores = raw_scores[valid] if score_name in LOWER_IS_BETTER: @@ -313,6 +356,7 @@ def benchmark_metrics(rows: list[dict], score_names: Iterable[str]) -> list[dict "n": int(valid.sum()), "n_pos": int(np.count_nonzero(labels == 1)), "n_neg": int(np.count_nonzero(labels == 0)), + "row_mode": row_mode, "auroc": _auroc(labels, scores), "average_precision": _average_precision(labels, scores), } @@ -335,6 +379,21 @@ def main() -> int: parser.add_argument("--backend", default=None, choices=("af2", "af3")) parser.add_argument("--pair-set", default=None, choices=("pos_pairs", "neg_pairs")) parser.add_argument("--scores", nargs="*", default=list(DEFAULT_SCORES)) + parser.add_argument( + "--metric-row-mode", + default="complete-case", + choices=("complete-case", "per-score"), + help=( + "complete-case compares all requested scores on the same finite rows; " + "per-score maximizes row availability independently for each score." + ), + ) + parser.add_argument( + "--checkpoint-every", + type=int, + default=100, + help="Rewrite the partial score CSV every N runs so long jobs leave usable progress.", + ) parser.add_argument("--log-level", default="INFO") args = parser.parse_args() @@ -352,9 +411,10 @@ def main() -> int: rows: list[dict] = [] errors: list[str] = [] + out_path = Path(args.out) for i, run_dir in enumerate(run_dirs, start=1): logger.info("scoring %d/%d %s", i, len(run_dirs), run_dir) - scored, error = score_run( + scored, run_errors = score_run( run_dir, root, contact_thresh=args.contact_thresh, @@ -364,12 +424,17 @@ def main() -> int: chain_pairs=args.chain_pairs, ) rows.extend(scored) - if error: + for error in run_errors: logger.warning(error) errors.append(error) + if args.checkpoint_every > 0 and i % args.checkpoint_every == 0: + _write_rows(out_path, rows) - _write_rows(Path(args.out), rows) - _write_rows(Path(args.metrics_out), benchmark_metrics(rows, args.scores)) + _write_rows(out_path, rows) + _write_rows( + Path(args.metrics_out), + benchmark_metrics(rows, args.scores, row_mode=args.metric_row_mode), + ) if errors: errors_path = Path(args.metrics_out).with_name(Path(args.metrics_out).stem + "_errors.txt") diff --git a/src/alphajudge/contact_probs.py b/src/alphajudge/contact_probs.py index 704e5310..5899d87e 100644 --- a/src/alphajudge/contact_probs.py +++ b/src/alphajudge/contact_probs.py @@ -1,6 +1,7 @@ from __future__ import annotations import numpy as np +from scipy.special import logsumexp AF2_DISTOGRAM_CONTACT_CUTOFF = 8.0 @@ -13,9 +14,12 @@ def contact_probs_from_distogram( """ Convert AF2 distogram logits to P(distance < contact_cutoff). - This follows the AlphaPulldown diagnostics fallback convention: append an - infinite final upper bound, clip the requested cutoff to 3-20 A, softmax - over bins, then sum bins whose upper bound is strictly below the 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 includes the bin that straddles 8 A in the standard AF2 + distogram, instead of dropping it because its upper edge is slightly above + 8 A. """ logits_arr = np.asarray(logits) if logits_arr.ndim != 3: @@ -28,68 +32,64 @@ def contact_probs_from_distogram( f"{logits_arr.shape[-1]}" ) - dtype = logits_arr.dtype if np.issubdtype(logits_arr.dtype, np.floating) else np.float32 - work = logits_arr.astype(dtype, copy=False) - shifted = work - np.max(work, axis=-1, keepdims=True) - np.exp(shifted, out=shifted) + work = ( + logits_arr + if np.issubdtype(logits_arr.dtype, np.floating) + else logits_arr.astype(np.float32) + ) - upper_bounds = np.concatenate([edges, [np.inf]]) + lower_bounds = np.concatenate([[0.0], edges]) clipped_cutoff = float(np.clip(float(contact_cutoff), 3.0, 20.0)) - contact_bins = upper_bounds < clipped_cutoff - if not np.any(contact_bins): - contact_bins[0] = True - - numerator = shifted[..., contact_bins].sum(axis=-1) - denominator = shifted.sum(axis=-1) - return numerator / denominator - - -def symmetrize_contact_probs(matrix: np.ndarray) -> tuple[np.ndarray, float]: - """Return a nan-aware symmetric copy and the maximum pre-symmetry delta.""" - m = np.asarray(matrix, dtype=float) - if m.ndim != 2 or m.shape[0] != m.shape[1]: - raise ValueError(f"contact probability matrix must be square, got {m.shape}") - - mt = m.T - sym = (m + mt) / 2.0 - - left_nan = np.isnan(m) - right_nan = np.isnan(mt) - sym[left_nan & ~right_nan] = mt[left_nan & ~right_nan] - sym[right_nan & ~left_nan] = m[right_nan & ~left_nan] - sym[left_nan & right_nan] = np.nan - - delta = np.abs(m - mt) - finite = np.isfinite(delta) - max_delta = float(np.nanmax(delta[finite])) if np.any(finite) else 0.0 - return sym, max_delta + n_contact_bins = int(np.count_nonzero(lower_bounds < clipped_cutoff)) + n_contact_bins = max(1, min(n_contact_bins, work.shape[-1])) + numerator = logsumexp(work[..., :n_contact_bins], axis=-1) + denominator = logsumexp(work, axis=-1) + return np.exp(numerator - denominator) def summarize_contact_prob_block( matrix: np.ndarray | None, idx1: np.ndarray, idx2: np.ndarray, top_n: int = 10, -) -> tuple[float, float, float]: - """Summarize inter-chain contact probabilities as max, top-N mean, and sum.""" +) -> tuple[float, float, float, float]: + """Summarize one inter-chain block as max, top-N mean, sum, and count >=0.5.""" if matrix is None or idx1.size == 0 or idx2.size == 0: nan = float("nan") - return nan, nan, nan + return nan, nan, nan, nan m = np.asarray(matrix, dtype=float) if m.ndim != 2: nan = float("nan") - return nan, nan, nan + return nan, nan, nan, nan if idx1.max(initial=-1) >= m.shape[0] or idx2.max(initial=-1) >= m.shape[1]: nan = float("nan") - return nan, nan, nan + return nan, nan, 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, nan + return nan, nan, 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)), float(np.sum(finite)) + return ( + float(np.max(finite)), + float(np.mean(top)), + float(np.sum(finite)), + float(np.count_nonzero(finite >= 0.5)), + ) diff --git a/src/alphajudge/interface.py b/src/alphajudge/interface.py index d847885f..429dacb2 100644 --- a/src/alphajudge/interface.py +++ b/src/alphajudge/interface.py @@ -217,7 +217,7 @@ def ilis(self) -> float: return float(math.sqrt(lis * clis)) @cached_property - def contact_probability_scores(self) -> tuple[float, float, float]: + def contact_probability_scores(self) -> tuple[float, float, float, float]: return summarize_contact_prob_block(self._contact_prob, self._idx1, self._idx2) @cached_property @@ -232,6 +232,10 @@ def contact_prob_top10_mean(self) -> float: def expected_contacts(self) -> float: return self.contact_probability_scores[2] + @cached_property + def confident_contacts(self) -> float: + return self.contact_probability_scores[3] + @property def polar(self) -> float: return self._frac(POLAR_RES) diff --git a/src/alphajudge/parsers/af2.py b/src/alphajudge/parsers/af2.py index 17c08914..60965264 100644 --- a/src/alphajudge/parsers/af2.py +++ b/src/alphajudge/parsers/af2.py @@ -10,7 +10,6 @@ from ..contact_probs import ( AF2_DISTOGRAM_CONTACT_CUTOFF, contact_probs_from_distogram, - symmetrize_contact_probs, ) logger = logging.getLogger(__name__) @@ -61,7 +60,9 @@ def load_model(model: str): iptm=iptm, ptm=ptm, iptm_ptm=iptm_ptm, confidence_score=conf, plddt_residue=plddt, contact_prob_matrix=contact_probs, - contact_prob_source="af2_distogram_lt_8A" if contact_probs is not None else None, + contact_prob_source=( + "af2_distogram_lb_lt_8A" if contact_probs is not None else None + ), ) return Run(order=order, source="af2", load_model=load_model) @@ -83,6 +84,7 @@ def _load_contact_probs_from_result_pkl( return None distogram = payload.get("distogram") if not isinstance(distogram, dict): + logger.debug(f"AF2 result pickle {result_pkl} has no distogram; contact scores unavailable.") return None logits = distogram.get("logits") bin_edges = distogram.get("bin_edges") @@ -106,12 +108,6 @@ def _load_contact_probs_from_result_pkl( ) return None - contact_probs, max_delta = symmetrize_contact_probs(contact_probs) - if max_delta > 1e-6: - logger.warning( - f"AF2 contact probabilities in {result_pkl} were asymmetric " - f"(max abs delta {max_delta:.3g}); symmetrized." - ) return contact_probs @staticmethod diff --git a/src/alphajudge/parsers/af3.py b/src/alphajudge/parsers/af3.py index 0250a018..9196bb6c 100644 --- a/src/alphajudge/parsers/af3.py +++ b/src/alphajudge/parsers/af3.py @@ -6,7 +6,7 @@ import numpy as np from . import BaseParser, Run from ..confidence import Confidence -from ..contact_probs import symmetrize_contact_probs +from ..geometry import is_pae_token_residue logger = logging.getLogger(__name__) @@ -278,6 +278,7 @@ def _normalize_contact_probs_af3(cls, matrix: dict, chains, cid) -> np.ndarray | aligned = cls._align_token_pair_matrix_to_residues( contact_probs, matrix.get("token_chain_ids"), + matrix.get("token_res_ids") or matrix.get("token_residue_ids"), chains, cid, expected_shape, @@ -289,18 +290,13 @@ def _normalize_contact_probs_af3(cls, matrix: dict, chains, cid) -> np.ndarray | ) return None - aligned, max_delta = symmetrize_contact_probs(aligned) - if max_delta > 1e-6: - logger.warning( - f"AF3 contact_probs were asymmetric (max abs delta {max_delta:.3g}); " - "symmetrized." - ) return aligned @staticmethod def _align_token_pair_matrix_to_residues( token_matrix: np.ndarray, token_chain_ids, + token_res_ids, chains, cid, expected_shape: tuple[int, int], @@ -311,6 +307,18 @@ def _align_token_pair_matrix_to_residues( return None ids = [str(x) for x in token_chain_ids] + if isinstance(token_res_ids, (list, tuple)) and len(token_res_ids) == len(ids): + residue_matrix = AF3Parser._align_token_pair_matrix_by_residue_ids( + token_matrix, + ids, + token_res_ids, + chains, + cid, + expected_shape, + ) + if residue_matrix is not None: + return residue_matrix + seen: list[str] = [] for chain_id in ids: if chain_id not in seen: @@ -347,3 +355,53 @@ def _align_token_pair_matrix_to_residues( continue residue_matrix[np.ix_(ri, rj)] = token_matrix[np.ix_(ti, tj)] return residue_matrix + + @staticmethod + def _align_token_pair_matrix_by_residue_ids( + token_matrix: np.ndarray, + token_chain_ids: list[str], + token_res_ids, + chains, + cid, + expected_shape: tuple[int, int], + ) -> np.ndarray | None: + token_lookup: dict[tuple[str, int], int] = {} + for token_idx, (chain_id, raw_res_id) in enumerate(zip(token_chain_ids, token_res_ids)): + try: + res_id = int(raw_res_id) + except (TypeError, ValueError): + continue + token_lookup.setdefault((chain_id, res_id), token_idx) + + token_indices_by_chain: dict[str, np.ndarray] = {} + for chain in chains: + residue_indices = np.asarray(cid.get(chain.id, []), dtype=int) + if residue_indices.size == 0: + token_indices_by_chain[chain.id] = np.array([], dtype=int) + continue + + kept = [res for res in chain if is_pae_token_residue(res)] + if len(kept) != residue_indices.size: + return None + + token_indices: list[int] = [] + for residue in kept: + token_idx = token_lookup.get((str(chain.id), int(residue.id[1]))) + if token_idx is None: + return None + token_indices.append(token_idx) + token_indices_by_chain[chain.id] = np.asarray(token_indices, dtype=int) + + residue_matrix = np.full(expected_shape, np.nan, dtype=float) + for chi in chains: + ri = np.asarray(cid.get(chi.id, []), dtype=int) + ti = token_indices_by_chain.get(chi.id, np.array([], dtype=int)) + if ri.size == 0: + continue + for chj in chains: + rj = np.asarray(cid.get(chj.id, []), dtype=int) + tj = token_indices_by_chain.get(chj.id, np.array([], dtype=int)) + if rj.size == 0: + continue + residue_matrix[np.ix_(ri, rj)] = token_matrix[np.ix_(ti, tj)] + return residue_matrix diff --git a/src/alphajudge/runner.py b/src/alphajudge/runner.py index 2cd838e9..99a48028 100644 --- a/src/alphajudge/runner.py +++ b/src/alphajudge/runner.py @@ -124,6 +124,7 @@ def process( "interface_contact_prob_max": iface.contact_prob_max, "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, "interface_expected_contacts": iface.expected_contacts, + "interface_confident_contacts": iface.confident_contacts, "interface_score": iface.score_complex, "interface_pDockQ2": pd2, "interface_ipSAE": iface.ipsae(), diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index ac8051fe..4e40b2a4 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -1,6 +1,7 @@ from __future__ import annotations import csv +import gzip import json import logging import math @@ -44,6 +45,7 @@ "interface_contact_prob_max", "interface_contact_prob_top10_mean", "interface_expected_contacts", + "interface_confident_contacts", "interface_score", "interface_pDockQ2", "interface_ipSAE", @@ -425,6 +427,7 @@ def test_contact_probability_scores_math_is_deterministic(): assert nearly_equal(iface.contact_prob_max, 0.8) assert nearly_equal(iface.contact_prob_top10_mean, 0.375) assert nearly_equal(iface.expected_contacts, 1.5) + assert nearly_equal(iface.confident_contacts, 1.0) missing = object.__new__(Interface) missing._idx1 = np.array([0, 1]) @@ -433,12 +436,14 @@ def test_contact_probability_scores_math_is_deterministic(): assert math.isnan(missing.contact_prob_max) assert math.isnan(missing.contact_prob_top10_mean) assert math.isnan(missing.expected_contacts) + assert math.isnan(missing.confident_contacts) def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: Path): """ AF2 contact probabilities are distogram softmax mass for bins with upper - bound strictly below 8 A, matching AlphaPulldown's diagnostics fallback. + lower bound is below 8 A, matching the published AF2 contact-probability + convention and including the bin that straddles 8 A. """ probs = np.array( [ @@ -451,20 +456,19 @@ def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: P bin_edges = np.array([4.0, 8.0, 12.0]) direct = contact_probs_from_distogram(logits, bin_edges) - expected_asym = np.array([[0.70, 0.10], [0.20, 0.05]]) + expected_asym = np.array([[0.90, 0.30], [0.50, 0.10]]) assert np.allclose(direct, expected_asym) run_dir = tmp_path / "af2_result" run_dir.mkdir() - with (run_dir / "result_model_1.pkl").open("wb") as f: + with gzip.open(run_dir / "result_model_1.pkl.gz", "wb") as f: pickle.dump({"distogram": {"logits": logits, "bin_edges": bin_edges}}, f) parsed = AF2Parser._load_contact_probs_from_result_pkl( run_dir, "model_1", expected_shape=(2, 2) ) assert parsed is not None - expected_sym = np.array([[0.70, 0.15], [0.15, 0.05]]) - assert np.allclose(parsed, expected_sym) + assert np.allclose(parsed, expected_asym) # ------------------------- @@ -566,6 +570,38 @@ def test_af3_contact_probability_scores_match_raw_contact_probs( assert nearly_equal(row["interface_contact_prob_max"], iface.contact_prob_max) assert nearly_equal(row["interface_contact_prob_top10_mean"], iface.contact_prob_top10_mean) assert nearly_equal(row["interface_expected_contacts"], iface.expected_contacts) + assert nearly_equal(row["interface_confident_contacts"], iface.confident_contacts) + + +def test_af3_contact_probs_alignment_uses_token_res_ids_with_extra_tokens(): + from Bio.PDB.Atom import Atom + from Bio.PDB.Chain import Chain + from Bio.PDB.Residue import Residue + + def residue(chain: Chain, resseq: int, serial: int) -> None: + res = Residue((" ", resseq, " "), "ALA", "") + res.add(Atom("CA", np.zeros(3), 1.0, 1.0, " ", "CA", serial, element="C")) + chain.add(res) + + chain_a = Chain("A") + chain_b = Chain("B") + residue(chain_a, 1, 1) + residue(chain_a, 2, 2) + residue(chain_b, 1, 3) + + token_matrix = np.arange(16, dtype=float).reshape(4, 4) + aligned = AF3Parser._align_token_pair_matrix_to_residues( + token_matrix, + ["A", "A", "A", "B"], + [1, 99, 2, 1], + [chain_a, chain_b], + {"A": [0, 1], "B": [2]}, + (3, 3), + ) + + assert aligned is not None + expected = token_matrix[np.ix_([0, 2, 3], [0, 2, 3])] + assert np.allclose(aligned, expected) @pytest.mark.parametrize("models_to_analyse", ["best", "all"]) From c6b50a0349fb22f5eb029404364356f511ac6afa Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Fri, 10 Jul 2026 11:24:00 +0200 Subject: [PATCH 3/6] Trim contact probability score outputs --- README.md | 4 +- .../evaluate_contact_probability_scores.py | 4 -- src/alphajudge/contact_probs.py | 35 +++++++++------- src/alphajudge/interface.py | 10 +---- src/alphajudge/parsers/af2.py | 15 ++++++- src/alphajudge/runner.py | 2 - test/test_parsers_and_runner.py | 42 ++++++++++++++----- 7 files changed, 68 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 83865f8e..9fbd1878 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,7 @@ process_many( ) ``` -Key outputs per interface include: `average_interface_pae`, `interface_average_plddt`, `interface_contact_pairs`, `interface_contact_prob_max`, `interface_contact_prob_top10_mean`, `interface_expected_contacts`, `interface_confident_contacts`, `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`. +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`. --- @@ -208,7 +208,7 @@ AlphaJudge writes `interfaces.csv` with one row per interface (and includes the - **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_cLIS, interface_iLIS** -- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean, interface_expected_contacts, interface_confident_contacts**: AF3 native `contact_probs`, or AF2 distogram-derived contact probability when full result pickles retain the `distogram` key. `interface_confident_contacts` counts inter-chain residue pairs with contact probability >= 0.5. +- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean**: AF3 native `contact_probs`, or AF2 distogram-derived contact probability 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. diff --git a/scripts/evaluate_contact_probability_scores.py b/scripts/evaluate_contact_probability_scores.py index 053ee2f2..e892bbcc 100755 --- a/scripts/evaluate_contact_probability_scores.py +++ b/scripts/evaluate_contact_probability_scores.py @@ -41,8 +41,6 @@ DEFAULT_SCORES = ( "interface_contact_prob_max", "interface_contact_prob_top10_mean", - "interface_expected_contacts", - "interface_confident_contacts", "iptm", "iptm_ptm", "confidence_score", @@ -170,8 +168,6 @@ def _row_for_interface(job: str, model: str, confidence, global_score: float, if "interface_contact_prob_source": confidence.contact_prob_source or "", "interface_contact_prob_max": iface.contact_prob_max, "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, - "interface_expected_contacts": iface.expected_contacts, - "interface_confident_contacts": iface.confident_contacts, "interface_score": iface.score_complex, "interface_pDockQ2": pd2, "interface_ipSAE": iface.ipsae(), diff --git a/src/alphajudge/contact_probs.py b/src/alphajudge/contact_probs.py index 5899d87e..19f0dee4 100644 --- a/src/alphajudge/contact_probs.py +++ b/src/alphajudge/contact_probs.py @@ -1,7 +1,6 @@ from __future__ import annotations import numpy as np -from scipy.special import logsumexp AF2_DISTOGRAM_CONTACT_CUTOFF = 8.0 @@ -43,28 +42,37 @@ def contact_probs_from_distogram( n_contact_bins = int(np.count_nonzero(lower_bounds < clipped_cutoff)) n_contact_bins = max(1, min(n_contact_bins, work.shape[-1])) - numerator = logsumexp(work[..., :n_contact_bins], axis=-1) - denominator = logsumexp(work, axis=-1) - return np.exp(numerator - denominator) + 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, float, float]: - """Summarize one inter-chain block as max, top-N mean, sum, and count >=0.5.""" +) -> 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, nan, nan + return nan, nan m = np.asarray(matrix, dtype=float) if m.ndim != 2: nan = float("nan") - return nan, nan, nan, 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, nan, nan + return nan, nan block = m[np.ix_(idx1, idx2)].ravel() if ( @@ -83,13 +91,8 @@ def summarize_contact_prob_block( finite = block[np.isfinite(block)] if finite.size == 0: nan = float("nan") - return nan, nan, nan, 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)), - float(np.sum(finite)), - float(np.count_nonzero(finite >= 0.5)), - ) + return float(np.max(finite)), float(np.mean(top)) diff --git a/src/alphajudge/interface.py b/src/alphajudge/interface.py index 429dacb2..6b2a1525 100644 --- a/src/alphajudge/interface.py +++ b/src/alphajudge/interface.py @@ -217,7 +217,7 @@ def ilis(self) -> float: return float(math.sqrt(lis * clis)) @cached_property - def contact_probability_scores(self) -> tuple[float, float, float, float]: + def contact_probability_scores(self) -> tuple[float, float]: return summarize_contact_prob_block(self._contact_prob, self._idx1, self._idx2) @cached_property @@ -228,14 +228,6 @@ def contact_prob_max(self) -> float: def contact_prob_top10_mean(self) -> float: return self.contact_probability_scores[1] - @cached_property - def expected_contacts(self) -> float: - return self.contact_probability_scores[2] - - @cached_property - def confident_contacts(self) -> float: - return self.contact_probability_scores[3] - @property def polar(self) -> float: return self._frac(POLAR_RES) diff --git a/src/alphajudge/parsers/af2.py b/src/alphajudge/parsers/af2.py index 60965264..f3cbfd9f 100644 --- a/src/alphajudge/parsers/af2.py +++ b/src/alphajudge/parsers/af2.py @@ -16,6 +16,7 @@ class AF2Parser(BaseParser): name = "af2" + _warned_missing_distogram = False def detect(self, d: Path) -> bool: return (d / "ranking_debug.json").exists() @@ -84,7 +85,19 @@ def _load_contact_probs_from_result_pkl( return None distogram = payload.get("distogram") if not isinstance(distogram, dict): - logger.debug(f"AF2 result pickle {result_pkl} has no distogram; contact scores unavailable.") + 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") diff --git a/src/alphajudge/runner.py b/src/alphajudge/runner.py index 99a48028..9bd0be47 100644 --- a/src/alphajudge/runner.py +++ b/src/alphajudge/runner.py @@ -123,8 +123,6 @@ def process( "interface_contact_prob_source": confidence.contact_prob_source or "", "interface_contact_prob_max": iface.contact_prob_max, "interface_contact_prob_top10_mean": iface.contact_prob_top10_mean, - "interface_expected_contacts": iface.expected_contacts, - "interface_confident_contacts": iface.confident_contacts, "interface_score": iface.score_complex, "interface_pDockQ2": pd2, "interface_ipSAE": iface.ipsae(), diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index 4e40b2a4..f29f09d9 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -44,8 +44,6 @@ "interface_contact_prob_source", "interface_contact_prob_max", "interface_contact_prob_top10_mean", - "interface_expected_contacts", - "interface_confident_contacts", "interface_score", "interface_pDockQ2", "interface_ipSAE", @@ -407,8 +405,7 @@ def test_clis_ilis_math_is_deterministic(): def test_contact_probability_scores_math_is_deterministic(): """ Contact-probability summaries are computed over all residue pairs between - the two chains: max, mean of the ten largest values, and expected contacts - as the sum of probabilities. + the two chains: max and mean of the ten largest values. """ from alphajudge.interface import Interface @@ -426,8 +423,6 @@ def test_contact_probability_scores_math_is_deterministic(): assert nearly_equal(iface.contact_prob_max, 0.8) assert nearly_equal(iface.contact_prob_top10_mean, 0.375) - assert nearly_equal(iface.expected_contacts, 1.5) - assert nearly_equal(iface.confident_contacts, 1.0) missing = object.__new__(Interface) missing._idx1 = np.array([0, 1]) @@ -435,8 +430,6 @@ def test_contact_probability_scores_math_is_deterministic(): missing._contact_prob = None assert math.isnan(missing.contact_prob_max) assert math.isnan(missing.contact_prob_top10_mean) - assert math.isnan(missing.expected_contacts) - assert math.isnan(missing.confident_contacts) def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: Path): @@ -471,6 +464,37 @@ def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: P assert np.allclose(parsed, expected_asym) +def test_af2_missing_distogram_warns_once_and_returns_none( + tmp_path: Path, caplog: pytest.LogCaptureFixture +): + run_dir = tmp_path / "af2_result" + run_dir.mkdir() + with (run_dir / "result_model_1.pkl").open("wb") as f: + pickle.dump({"plddt": [90.0]}, f) + with (run_dir / "result_model_2.pkl").open("wb") as f: + pickle.dump({"plddt": [80.0]}, f) + + old_warned = AF2Parser._warned_missing_distogram + AF2Parser._warned_missing_distogram = False + try: + caplog.set_level(logging.WARNING, logger="alphajudge.parsers.af2") + parsed = AF2Parser._load_contact_probs_from_result_pkl( + run_dir, "model_1", expected_shape=(1, 1) + ) + assert parsed is None + assert "has no distogram" in caplog.text + assert "--remove_keys_from_pickles" in caplog.text + + caplog.clear() + parsed = AF2Parser._load_contact_probs_from_result_pkl( + run_dir, "model_2", expected_shape=(1, 1) + ) + assert parsed is None + assert "has no distogram" not in caplog.text + finally: + AF2Parser._warned_missing_distogram = old_warned + + # ------------------------- # AF3 runner: best/all + score checks # ------------------------- @@ -569,8 +593,6 @@ def test_af3_contact_probability_scores_match_raw_contact_probs( assert row["interface_contact_prob_source"] == "af3_contact_probs" assert nearly_equal(row["interface_contact_prob_max"], iface.contact_prob_max) assert nearly_equal(row["interface_contact_prob_top10_mean"], iface.contact_prob_top10_mean) - assert nearly_equal(row["interface_expected_contacts"], iface.expected_contacts) - assert nearly_equal(row["interface_confident_contacts"], iface.confident_contacts) def test_af3_contact_probs_alignment_uses_token_res_ids_with_extra_tokens(): From 6fce50a27d65593d3913ca9c79f5814f91e26203 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Fri, 10 Jul 2026 12:10:56 +0200 Subject: [PATCH 4/6] Use Humphreys AF2 contact cutoff --- README.md | 2 +- src/alphajudge/contact_probs.py | 7 ++++--- src/alphajudge/parsers/af2.py | 2 +- test/test_parsers_and_runner.py | 15 +++++++++++---- 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 9fbd1878..ec804cec 100644 --- a/README.md +++ b/README.md @@ -208,7 +208,7 @@ AlphaJudge writes `interfaces.csv` with one row per interface (and includes the - **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_cLIS, interface_iLIS** -- **interface_contact_prob_source, interface_contact_prob_max, interface_contact_prob_top10_mean**: AF3 native `contact_probs`, or AF2 distogram-derived contact probability 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_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. diff --git a/src/alphajudge/contact_probs.py b/src/alphajudge/contact_probs.py index 19f0dee4..d2a70e46 100644 --- a/src/alphajudge/contact_probs.py +++ b/src/alphajudge/contact_probs.py @@ -2,7 +2,7 @@ import numpy as np -AF2_DISTOGRAM_CONTACT_CUTOFF = 8.0 +AF2_DISTOGRAM_CONTACT_CUTOFF = 12.0 def contact_probs_from_distogram( @@ -16,9 +16,10 @@ def contact_probs_from_distogram( 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 includes the bin that straddles 8 A in the standard AF2 + 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 - 8 A. + 12 A. """ logits_arr = np.asarray(logits) if logits_arr.ndim != 3: diff --git a/src/alphajudge/parsers/af2.py b/src/alphajudge/parsers/af2.py index f3cbfd9f..857efd3f 100644 --- a/src/alphajudge/parsers/af2.py +++ b/src/alphajudge/parsers/af2.py @@ -62,7 +62,7 @@ def load_model(model: str): plddt_residue=plddt, contact_prob_matrix=contact_probs, contact_prob_source=( - "af2_distogram_lb_lt_8A" if contact_probs is not None else None + "af2_distogram_lb_lt_12A" if contact_probs is not None else None ), ) return Run(order=order, source="af2", load_model=load_model) diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index f29f09d9..c70b87ad 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -434,9 +434,9 @@ def test_contact_probability_scores_math_is_deterministic(): def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: Path): """ - AF2 contact probabilities are distogram softmax mass for bins with upper - lower bound is below 8 A, matching the published AF2 contact-probability - convention and including the bin that straddles 8 A. + AF2 contact probabilities are distogram softmax mass for bins whose + lower bound is below 12 A, matching the Humphreys et al. + contact-probability convention and including the bin that straddles 12 A. """ probs = np.array( [ @@ -449,9 +449,16 @@ def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: P bin_edges = np.array([4.0, 8.0, 12.0]) direct = contact_probs_from_distogram(logits, bin_edges) - expected_asym = np.array([[0.90, 0.30], [0.50, 0.10]]) + expected_asym = np.array([[0.95, 0.60], [0.60, 0.30]]) assert np.allclose(direct, expected_asym) + standard_edges = np.linspace(2.3125, 21.6875, 63) + standard_logits = np.zeros((1, 1, 64), dtype=float) + assert np.allclose( + contact_probs_from_distogram(standard_logits, standard_edges), + np.array([[32 / 64]]), + ) + run_dir = tmp_path / "af2_result" run_dir.mkdir() with gzip.open(run_dir / "result_model_1.pkl.gz", "wb") as f: From 67e4841aa6d7a57c2e0015c6236b2a33b2e4c8c5 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Fri, 10 Jul 2026 12:42:23 +0200 Subject: [PATCH 5/6] Move benchmark helpers to manual tests --- scripts/extract_pisa_molref.py | 2 -- src/alphajudge/meta_score.py | 15 +++++++------- src/alphajudge/report.py | 4 ++-- test/manual/README.md | 7 +++++++ .../evaluate_contact_probability_scores.py | 12 +++++------ .../manual}/freeze_metascore_quantiles.py | 20 +++++++++++-------- 6 files changed, 34 insertions(+), 26 deletions(-) create mode 100644 test/manual/README.md rename {scripts => test/manual}/evaluate_contact_probability_scores.py (97%) rename {scripts => test/manual}/freeze_metascore_quantiles.py (87%) diff --git a/scripts/extract_pisa_molref.py b/scripts/extract_pisa_molref.py index 53998931..bb303c58 100644 --- a/scripts/extract_pisa_molref.py +++ b/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() diff --git a/src/alphajudge/meta_score.py b/src/alphajudge/meta_score.py index 658b7c26..a26e88d6 100644 --- a/src/alphajudge/meta_score.py +++ b/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 = { diff --git a/src/alphajudge/report.py b/src/alphajudge/report.py index 2bc615c5..9faaf646 100644 --- a/src/alphajudge/report.py +++ b/src/alphajudge/report.py @@ -81,7 +81,7 @@ _AJ_DARK = "#111111" _REPORT_TITLE = "AlphaJudge Interface validation Report" -_BENCHMARK_TAG = "benchmark_26 positives (final_sync_20260523, n=3,878 interacting AF2/AF3 pairs)" +_BENCHMARK_TAG = "AlphaJudge interacting reference set (n=3,878 AF2/AF3 pairs)" _GRADIENT = np.tile(np.linspace(0.0, 1.0, 1024), (2, 1)) @@ -1451,7 +1451,7 @@ def _aggregate_cover_page( info = [ "This report scores AlphaFold-predicted complexes against the", - "AlphaJudge benchmark_26 interacting (positive) reference set.", + "AlphaJudge interacting (positive) reference set.", "All percentiles are archive percentiles; higher is better.", ] _draw_info_box(fig, x=0.13, y=0.54, w=0.74, h=0.11, lines=info) diff --git a/test/manual/README.md b/test/manual/README.md new file mode 100644 index 00000000..4fb785de --- /dev/null +++ b/test/manual/README.md @@ -0,0 +1,7 @@ +# Manual Benchmark Helpers + +These scripts are developer-only helpers for reproducing benchmark comparisons +or regenerating frozen calibration constants. They are not run by pytest and are +not installed as AlphaJudge command-line tools. + +Run them manually with explicit local input/output paths. diff --git a/scripts/evaluate_contact_probability_scores.py b/test/manual/evaluate_contact_probability_scores.py similarity index 97% rename from scripts/evaluate_contact_probability_scores.py rename to test/manual/evaluate_contact_probability_scores.py index e892bbcc..733cfaf1 100755 --- a/scripts/evaluate_contact_probability_scores.py +++ b/test/manual/evaluate_contact_probability_scores.py @@ -1,5 +1,8 @@ #!/usr/bin/env python3 -"""Evaluate AlphaJudge contact-probability scores on a labelled prediction tree. +"""Manual AUROC/AP comparison for AlphaJudge scores on a labelled prediction tree. + +This is a developer benchmark helper, not part of the release CLI and not run +by pytest. Invoke it manually with explicit local benchmark paths. The benchmark tree is expected to look like: @@ -23,7 +26,7 @@ import numpy as np -REPO_ROOT = Path(__file__).resolve().parents[1] +REPO_ROOT = Path(__file__).resolve().parents[2] SRC_ROOT = REPO_ROOT / "src" if SRC_ROOT.exists(): sys.path.insert(0, str(SRC_ROOT)) @@ -35,9 +38,6 @@ logger = logging.getLogger("contact_probability_benchmark") -DEFAULT_ROOT = Path( - "/g/transform/kosinski/dima/IntAct_BioGRID_STRING/benchmark_26/predictions" -) DEFAULT_SCORES = ( "interface_contact_prob_max", "interface_contact_prob_top10_mean", @@ -362,7 +362,7 @@ def benchmark_metrics( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--predictions-root", default=str(DEFAULT_ROOT)) + parser.add_argument("--predictions-root", required=True) parser.add_argument("--out", required=True, help="Output per-interface/per-chain-pair CSV.") parser.add_argument("--metrics-out", required=True, help="Output AUROC/AP summary CSV.") parser.add_argument("--models-to-analyse", default="best", choices=("best", "all")) diff --git a/scripts/freeze_metascore_quantiles.py b/test/manual/freeze_metascore_quantiles.py similarity index 87% rename from scripts/freeze_metascore_quantiles.py rename to test/manual/freeze_metascore_quantiles.py index e973704f..ba601e38 100644 --- a/scripts/freeze_metascore_quantiles.py +++ b/test/manual/freeze_metascore_quantiles.py @@ -1,5 +1,8 @@ #!/usr/bin/env python3 -"""Freeze the AlphaJudge metascore calibration deciles from the benchmark. +"""Manually freeze AlphaJudge metascore calibration deciles from a benchmark CSV. + +This is a developer calibration helper, not part of the release CLI and not run +by pytest. Invoke it manually with an explicit benchmark CSV path. The percentile sliders in AlphaJudge reports map each raw interface descriptor onto a frozen percentile scale (``BENCHMARK_QUANTILES`` in @@ -20,7 +23,7 @@ runs in a stock AlphaJudge install. Usage: - python scripts/freeze_metascore_quantiles.py \ + python test/manual/freeze_metascore_quantiles.py \ --input-csv .../benchmark_best....csv \ [--label-filter positive] # use "all" to reproduce the legacy scale """ @@ -30,21 +33,22 @@ import argparse import csv import math +import sys from pathlib import Path import numpy as np +REPO_ROOT = Path(__file__).resolve().parents[2] +SRC_ROOT = REPO_ROOT / "src" +if SRC_ROOT.exists(): + sys.path.insert(0, str(SRC_ROOT)) + from alphajudge.meta_score import ( BENCHMARK_QUANTILES, CALIBRATION_LEVELS, FEATURE_DIRECTIONS, ) -DEFAULT_BENCHMARK_CSV = Path( - "/scratch/dima/benchmark_26/final_sync_20260523_225722/staged_benchmark/" - "benchmark_best.final_sync_20260523_225722_force_recompute_nointerfacefix.csv" -) - def _safe_float(value) -> float: try: @@ -62,7 +66,7 @@ def feature_deciles(rows: list[dict[str, str]], feature: str, direction: float) def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--input-csv", default=str(DEFAULT_BENCHMARK_CSV)) + parser.add_argument("--input-csv", required=True) parser.add_argument( "--label-filter", default="positive", From f50189a0b3f83c549e56ed11b838f97d53b2372c Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Fri, 10 Jul 2026 12:46:14 +0200 Subject: [PATCH 6/6] Guard AF2 contact probability parsing --- test/test_parsers_and_runner.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/test/test_parsers_and_runner.py b/test/test_parsers_and_runner.py index c70b87ad..8f41dd83 100644 --- a/test/test_parsers_and_runner.py +++ b/test/test_parsers_and_runner.py @@ -458,6 +458,12 @@ def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: P contact_probs_from_distogram(standard_logits, standard_edges), np.array([[32 / 64]]), ) + boundary_mass_logits = np.full((1, 1, 64), -1000.0, dtype=float) + boundary_mass_logits[..., 31] = 0.0 + assert np.allclose( + contact_probs_from_distogram(boundary_mass_logits, standard_edges), + np.array([[1.0]]), + ) run_dir = tmp_path / "af2_result" run_dir.mkdir() @@ -471,6 +477,32 @@ def test_af2_distogram_contact_probs_softmax_cutoff_is_deterministic(tmp_path: P assert np.allclose(parsed, expected_asym) +def test_af2_distogram_loader_uses_requested_model_not_last_glob(tmp_path: Path): + run_dir = tmp_path / "af2_result" + run_dir.mkdir() + bin_edges = np.array([4.0, 8.0, 12.0]) + + def write_distogram(model: str, contact_prob: float) -> None: + logits = np.full((1, 1, 4), -1000.0, dtype=float) + logits[..., 0] = np.log(contact_prob) + logits[..., 3] = np.log(1.0 - contact_prob) + with (run_dir / f"result_{model}.pkl").open("wb") as f: + pickle.dump( + {"distogram": {"logits": logits, "bin_edges": bin_edges}}, + f, + ) + + write_distogram("model_1", 0.25) + write_distogram("model_2", 0.75) + + parsed = AF2Parser._load_contact_probs_from_result_pkl( + run_dir, "model_1", expected_shape=(1, 1) + ) + + assert parsed is not None + assert np.allclose(parsed, np.array([[0.25]])) + + def test_af2_missing_distogram_warns_once_and_returns_none( tmp_path: Path, caplog: pytest.LogCaptureFixture ):