fix: Align sample selection with official AF3

.publish/manifest.json

@@ -1,6 +1,6 @@
 {
   "schema_version": 1,
-  "source_commit": "77cc17f7700eec989a54bb1ac129c0cdc97203f5",
+  "source_commit": "c6089bc41e327ccd8b6cda2d911ebc94f1d7a6b4",
   "paths": [
     "CMakeLists.txt",
     "LICENSE",
@@ -859,7 +859,7 @@
     "src/alphafold3_mlx/jax_rng.py": "ea7154510d7398351e2aba7148a361e62566bfa075d0b186c719ccf2b970161e",
     "src/alphafold3_mlx/model/__init__.py": "1a57850227b93f46be413ca87016078dd2ad90463f67ffc242afe535a366c323",
     "src/alphafold3_mlx/model/inference.py": "11b44e55784bc0146de1248594d5c891444ddbed1a928cbc9fd813fce2839867",
-    "src/alphafold3_mlx/model/model.py": "de9d9ab21d6f8897c0d82e8d78dfd07614c484227f512e84a6a5834cdc5c2365",
+    "src/alphafold3_mlx/model/model.py": "b001879685bc927d365c4e64113035a689ccb82fe6f35a819a2db886ccb54364",
     "src/alphafold3_mlx/model/msa_sampling.py": "625e261d2a23eabdd5dbb8453ece8e889c18ea02887e60e3e011e31aa90f408e",
     "src/alphafold3_mlx/model/recycling.py": "87841f1cfca598cb7a16493f1c8bb9ed3493b895efe5009df8b9911c213032ee",
     "src/alphafold3_mlx/modules/__init__.py": "b91717e9fe7c7a1b8cb7d2e1ef55d37a829e4fc89e27e51a67bf0a30f71bee44",
@@ -891,8 +891,8 @@
     "src/alphafold3_mlx/pipeline/msa_cache.py": "c45893dd228573a7ad38749ac30f662ca751329d98a93af8ffcdf6e0faf097e8",
     "src/alphafold3_mlx/pipeline/output_handler.py": "077cb500cf16b711774fb78b724fa3586e8917c210016cae04c5e3c8b0f081a7",
     "src/alphafold3_mlx/pipeline/progress.py": "e802892627f20b9cd8efc7102e4bee7e0092b849619fcc4fc0e4a64f58deb8d6",
-    "src/alphafold3_mlx/pipeline/ranking.py": "3fa26b3a138cb9f9d9abcfaa781ee920f0ebbd847f8ccc59c1ab94dbd7dda8b5",
-    "src/alphafold3_mlx/pipeline/runner.py": "2bd8e899850bc7a9e8b8b5fbda476e6dd91cdca72039928be1208bcebd2a81af",
+    "src/alphafold3_mlx/pipeline/ranking.py": "6d2ae5a213bb127e870c975a8d76857d0f096503702bb37e7ee4efd472bc6bb4",
+    "src/alphafold3_mlx/pipeline/runner.py": "cb49e7124eaa0943fb05db67d7df45777eed3c768da32968226ad0b4c255e402",
     "src/alphafold3_mlx/pipeline/structure_parser.py": "deb8935de396dd70ef582aa2fbf2f6e5225fcaafed9af7d435d74862218cb3d5",
     "src/alphafold3_mlx/pyproject.toml": "c9fba5ae80ed2d10db609cfa095f7638774c8a8077245735d6897958f7339eab",
     "src/alphafold3_mlx/reference/__init__.py": "1a83c98ce8e3a29c4c610e6873680971b898dbeb6aecc36cdba0c3e0bf204a31",
@@ -966,7 +966,7 @@
     "tests/integration/test_large_proteins.py": "7244b43a8e9baca54149bae1961a8f1c3e945374bbf5cc55e3dbe958ae347679",
     "tests/integration/test_msa_processing.py": "68a8e21bfc055f22e15a8759f584b055a1099fefcab6996d9ae96b7e58604eb3",
     "tests/integration/test_precision_modes.py": "748bacda9515e3cea7b09de8481890b764ce9d86d1549cafc093c50790d39e75",
-    "tests/integration/test_ranking.py": "14bbc238a4eeb1eb848f1c12dc254ef1c0541b2bc1c9a57825c0e04601f23226",
+    "tests/integration/test_ranking.py": "01103a946738bb6d883e116a31bdffb01693ca0e6d4d4097c610a9c0b3fa0f86",
     "tests/integration/test_recycling_stability.py": "0e555c1f9736cec72d9833b9c62e82322ff5d4d835de78ad84bed1f5bb40e002",
     "tests/integration/test_restrained_inference.py": "f3481ce23079b43c78ae6a195b2fb72d6d2feece647a4ef20e02c5aa6dd5c8b3",
     "tests/integration/test_restraint_pipeline_smoke.py": "7fe339fac9637ba9222a53ead9f9bf8093ec4770fafe57be6a1c5fcce3e16a54",

src/alphafold3_mlx/model/model.py

@@ -36,7 +36,7 @@
 from alphafold3_mlx.network.diffusion_head import DiffusionHead
 from alphafold3_mlx.network.confidence_head import ConfidenceHead
 from alphafold3_mlx.network.atom_cross_attention import AtomCrossAttEncoder
-from alphafold3_mlx.jax_rng import haiku_next_rng_keys
+from alphafold3_mlx.jax_rng import official_af3_model_and_diffusion_keys
 from alphafold3_mlx.model.recycling import run_recycling_loop
 
 if TYPE_CHECKING:
@@ -363,7 +363,15 @@ def __call__(
         # Canonical run_alphafold obtains the trunk model key through the
         # first Haiku next_rng_key() call. This key drives per-recycle MSA
         # sampling and must match JAX exactly.
-        model_key = haiku_next_rng_keys(key, 1)[0]
+        model_key, official_diffusion_key = official_af3_model_and_diffusion_keys(
+            key, num_recycles=self.config.num_recycles
+        )
+        if guidance_fn is None:
+            diffusion_key = mx.array(official_diffusion_key)
+        else:
+            # Restraint guidance is an MLX-only extension with a separately
+            # calibrated stochastic stream.
+            _, diffusion_key = mx.random.split(key)
 
         # Extract features
         token_features = batch.token_features
@@ -593,10 +601,7 @@ def _set_contact_positions(token_i, token_j, num_res):
         atom37_mask = atom37_mask * seq_mask[:, :, None]  # [batch, residues, 37]
 
         # Generate coordinates via diffusion (directly in atom37 format).
-        # Keep the calibrated MLX diffusion stream for product stability.
-        # Switching this custom restrained sampler to the JAX random stream
-        # regresses K48 interface recovery despite identical trunk outputs.
-        _, diffusion_key = mx.random.split(key)
+        # Use the post-trunk Haiku key position from the official AF3 schedule.
 
         # Build a minimal diffusion batch from FeatureBatch inputs.
         from alphafold3_mlx.atom_layout import GatherInfo

src/alphafold3_mlx/pipeline/ranking.py

@@ -13,6 +13,75 @@
 from typing import Any, Literal
 
 
+_IPTM_WEIGHT = 0.8
+_FRACTION_DISORDERED_WEIGHT = 0.5
+_CLASH_PENALTY = 100.0
+
+
+def compute_structure_quality_metrics(
+    *,
+    atom_positions: Any,
+    atom_mask: Any,
+    atom_names: Any,
+    element_symbols: Any,
+    comp_ids: Any,
+    chain_ids: Any,
+    residue_indices: Any,
+    chain_types: Any,
+) -> tuple[list[float], list[bool]]:
+    """Compute the structure-dependent terms in the official AF3 ranking."""
+    import numpy as np
+
+    from alphafold3 import structure
+    from alphafold3.model import confidences
+
+    positions = np.asarray(atom_positions)
+    masks = np.asarray(atom_mask, dtype=bool)
+    names = np.asarray(atom_names, dtype=object)
+    elements = np.asarray(element_symbols, dtype=object)
+    residues = np.asarray(comp_ids, dtype=object)
+    chains = np.asarray(chain_ids, dtype=object)
+    residue_ids = np.asarray(residue_indices, dtype=np.int32)
+    types = np.asarray(chain_types, dtype=object)
+
+    if positions.ndim != 4 or positions.shape[-1] != 3:
+        raise ValueError("atom_positions must have shape [samples, residues, atoms, 3]")
+    if masks.shape != positions.shape[:-1]:
+        raise ValueError("atom_mask must match atom_positions sample/residue/atom axes")
+    if names.shape != positions.shape[1:3] or elements.shape != names.shape:
+        raise ValueError("atom metadata must match the residue/atom layout")
+    for values, field_name in (
+        (residues, "comp_ids"),
+        (chains, "chain_ids"),
+        (residue_ids, "residue_indices"),
+        (types, "chain_types"),
+    ):
+        if values.shape != (positions.shape[1],):
+            raise ValueError(f"{field_name} must match the residue axis")
+
+    fraction_disordered_scores = []
+    has_clash_scores = []
+    for sample_positions, sample_mask in zip(positions, masks, strict=True):
+        predicted_structure = structure.from_res_arrays(
+            atom_mask=sample_mask,
+            atom_x=sample_positions[..., 0],
+            atom_y=sample_positions[..., 1],
+            atom_z=sample_positions[..., 2],
+            atom_name=names,
+            atom_element=elements,
+            chain_id=chains,
+            chain_type=types,
+            res_id=residue_ids,
+            res_name=residues,
+        )
+        fraction_disordered_scores.append(
+            float(confidences.fraction_disordered(predicted_structure))
+        )
+        has_clash_scores.append(bool(confidences.has_clash(predicted_structure)))
+
+    return fraction_disordered_scores, has_clash_scores
+
+
 @dataclass
 class RankingScores:
     """Confidence scores for a single sample.
@@ -28,6 +97,9 @@ class RankingScores:
     iptm: float
     mean_plddt: float
     plddt_variance: float = 0.0
+    fraction_disordered: float = 0.0
+    has_clash: bool = False
+    ranking_score: float | None = None
 
     def to_dict(self) -> dict[str, float]:
         """Convert to JSON-serializable dict."""
@@ -36,6 +108,9 @@ def to_dict(self) -> dict[str, float]:
             "iptm": self.iptm,
             "mean_plddt": self.mean_plddt,
             "plddt_variance": self.plddt_variance,
+            "fraction_disordered": self.fraction_disordered,
+            "has_clash": self.has_clash,
+            "ranking_score": self.ranking_score,
         }
 
 
@@ -52,7 +127,7 @@ class SampleRanking:
 
     ranked_indices: list[int]
     scores: dict[int, RankingScores]
-    ranking_metric: Literal["pTM", "ipTM"]
+    ranking_metric: Literal["pTM", "ipTM", "ranking_score"]
     is_complex: bool
 
     @property
@@ -63,6 +138,11 @@ def best_index(self) -> int:
     @property
     def best_score(self) -> float:
         """Score of the best-ranked sample."""
+        if self.ranking_metric == "ranking_score":
+            score = self.scores[self.best_index].ranking_score
+            if score is None:
+                raise ValueError("ranking_score metric selected without scores")
+            return score
         metric_name = "ptm" if self.ranking_metric == "pTM" else "iptm"
         return getattr(self.scores[self.best_index], metric_name)
 
@@ -86,6 +166,9 @@ def to_ranking_debug_dict(self) -> dict[str, Any]:
                 "ptm": score.ptm,
                 "iptm": score.iptm,
                 "mean_plddt": score.mean_plddt,
+                "ranking_score": score.ranking_score,
+                "fraction_disordered": score.fraction_disordered,
+                "has_clash": score.has_clash,
             })
 
         aggregate = compute_aggregate_metrics(self)
@@ -104,6 +187,8 @@ def rank_samples(
     iptm_scores: list[float],
     plddt_scores: list[list[float]],
     is_complex: bool,
+    fraction_disordered_scores: list[float] | None = None,
+    has_clash_scores: list[bool] | None = None,
 ) -> SampleRanking:
     """Rank samples by confidence metric.
 
@@ -121,9 +206,33 @@ def rank_samples(
     """
     num_samples = len(ptm_scores)
 
-    # Choose ranking metric
-    if is_complex:
-        ranking_metric: Literal["pTM", "ipTM"] = "ipTM"
+    use_official_score = (
+        fraction_disordered_scores is not None or has_clash_scores is not None
+    )
+    if fraction_disordered_scores is None:
+        fraction_disordered_scores = [0.0] * num_samples
+    if has_clash_scores is None:
+        has_clash_scores = [False] * num_samples
+    if len(fraction_disordered_scores) != num_samples:
+        raise ValueError("fraction_disordered_scores must match sample count")
+    if len(has_clash_scores) != num_samples:
+        raise ValueError("has_clash_scores must match sample count")
+
+    if use_official_score:
+        ranking_metric: Literal["pTM", "ipTM", "ranking_score"] = "ranking_score"
+        metric_values = [
+            (
+                _IPTM_WEIGHT * iptm_scores[i]
+                + (1.0 - _IPTM_WEIGHT) * ptm_scores[i]
+                if is_complex
+                else ptm_scores[i]
+            )
+            + _FRACTION_DISORDERED_WEIGHT * fraction_disordered_scores[i]
+            - _CLASH_PENALTY * has_clash_scores[i]
+            for i in range(num_samples)
+        ]
+    elif is_complex:
+        ranking_metric = "ipTM"
         metric_values = iptm_scores
     else:
         ranking_metric = "pTM"
@@ -153,6 +262,9 @@ def rank_samples(
             iptm=iptm_scores[i],
             mean_plddt=mean_plddt,
             plddt_variance=variance,
+            fraction_disordered=fraction_disordered_scores[i],
+            has_clash=has_clash_scores[i],
+            ranking_score=metric_values[i] if use_official_score else None,
         )
 
     return SampleRanking(

src/alphafold3_mlx/pipeline/runner.py

@@ -211,7 +211,7 @@ def run(self) -> "OutputBundle":
 
             # Rank samples (not a timed stage per spec)
             is_complex = auto_detect_complex(self.input_json.chain_ids)
-            ranking = self._rank_samples(result, is_complex)
+            ranking = self._rank_samples(result, is_complex, batch)
 
             # Write outputs - timing.json stage: output_writing
             self.progress.on_stage_start("output_writing")
@@ -666,7 +666,9 @@ def _build_guidance_fn(self, batch: Any, model: Any) -> Any:
 
         return guidance_fn
 
-    def _rank_samples(self, result: Any, is_complex: bool) -> "SampleRanking":
+    def _rank_samples(
+        self, result: Any, is_complex: bool, batch: Any
+    ) -> "SampleRanking":
         """Rank samples by confidence.
 
         Args:
@@ -677,7 +679,14 @@ def _rank_samples(self, result: Any, is_complex: bool) -> "SampleRanking":
             SampleRanking with ranked indices.
         """
         import numpy as np
-        from alphafold3_mlx.pipeline.ranking import rank_samples
+        from alphafold3.constants import mmcif_names
+        from alphafold3_mlx.pipeline.output_handler import (
+            build_structure_atom_metadata,
+        )
+        from alphafold3_mlx.pipeline.ranking import (
+            compute_structure_quality_metrics,
+            rank_samples,
+        )
 
         np_result = result.to_numpy()
 
@@ -711,7 +720,59 @@ def _rank_samples(self, result: Any, is_complex: bool) -> "SampleRanking":
             residue_mean_plddt = residue_mean_plddt[residue_mask]
             plddt_scores.append(residue_mean_plddt.tolist())
 
-        return rank_samples(ptm_scores, iptm_scores, plddt_scores, is_complex)
+        token_mask = np.asarray(batch.token_features.mask).reshape(-1) > 0
+        atom_metadata = build_structure_atom_metadata(
+            batch, self.input_json.input.chains
+        )
+        required_metadata = {
+            "atom_names", "element_symbols", "comp_ids", "chain_ids", "is_ligand",
+        }
+        missing_metadata = required_metadata.difference(atom_metadata)
+        if missing_metadata:
+            raise ValueError(
+                "Cannot compute official AF3 ranking without structure metadata: "
+                + ", ".join(sorted(missing_metadata))
+            )
+
+        comp_ids = np.asarray(atom_metadata["comp_ids"])[token_mask]
+        chain_ids = np.asarray(atom_metadata["chain_ids"])[token_mask]
+        is_ligand = np.asarray(atom_metadata["is_ligand"], dtype=bool)[token_mask]
+        chain_types = np.empty(comp_ids.shape, dtype=object)
+        for chain_id in dict.fromkeys(chain_ids.tolist()):
+            chain_mask = chain_ids == chain_id
+            if np.all(is_ligand[chain_mask]):
+                chain_type = mmcif_names.NON_POLYMER_CHAIN
+            else:
+                chain_type = mmcif_names.guess_polymer_type(
+                    comp_ids[chain_mask].tolist()
+                )
+            chain_types[chain_mask] = chain_type
+
+        fraction_disordered_scores, has_clash_scores = (
+            compute_structure_quality_metrics(
+                atom_positions=np_result["atom_positions"][:, token_mask],
+                atom_mask=np_result["atom_mask"][:, token_mask],
+                atom_names=np.asarray(atom_metadata["atom_names"])[token_mask],
+                element_symbols=np.asarray(atom_metadata["element_symbols"])[
+                    token_mask
+                ],
+                comp_ids=comp_ids,
+                chain_ids=chain_ids,
+                residue_indices=np.asarray(batch.token_features.residue_index)[
+                    token_mask
+                ],
+                chain_types=chain_types,
+            )
+        )
+
+        return rank_samples(
+            ptm_scores,
+            iptm_scores,
+            plddt_scores,
+            is_complex,
+            fraction_disordered_scores=fraction_disordered_scores,
+            has_clash_scores=has_clash_scores,
+        )
 
     def _write_outputs(
         self,