shellde 0.2.0__py3-none-any.whl

shellde/naturalness.py

@@ -0,0 +1,142 @@
+"""Zero-shot naturalness providers (PLM likelihood) + a test/compose stub.
+
+Naturalness is FolDE's actual lever (per the 6-paper tool-development synthesis): a
+PLM zero-shot score, distilled into the supervised readout as a WARM START, is what
+prevents the round-2 collapse of pure activity-only few-shot AL. Here a provider maps
+full-length sequences to a per-sequence scalar; it is injected as a feature
+(``features.NaturalnessBlock``) and/or used as a discovery scorer (``zero_shot``).
+
+``FunctionNaturalness`` wraps any ``str -> float`` callable so the plumbing is
+testable without a heavy PLM SDK. The ESM providers are import-guarded.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+
+from shellde.embeddings.provider import EmbeddingCache, _key
+
+
+class FunctionNaturalness:
+    """A NaturalnessProvider backed by a plain ``str -> float`` function (cached)."""
+
+    def __init__(
+        self, fn: Callable[[str], float], *, model_id: str = "fn", cache: EmbeddingCache | None = None
+    ) -> None:
+        self._fn = fn
+        self.model_id = model_id
+        self._cache = cache if cache is not None else EmbeddingCache()
+        self.calls = 0
+
+    def score(self, sequences: Sequence[str]) -> np.ndarray:
+        out: list[float] = []
+        for s in sequences:
+            k = _key(self.model_id, s)
+            v = self._cache.get(k)
+            if v is None:
+                self.calls += 1
+                v = np.asarray([float(self._fn(s))], dtype=float)
+                self._cache.put(k, v)
+            out.append(float(v[0]))
+        return np.asarray(out, dtype=float)
+
+
+class ESM2Naturalness:
+    """ESM-2 pseudo-log-likelihood naturalness (mean per-residue log-prob), import-guarded."""
+
+    def __init__(
+        self, model: str = "facebook/esm2_t30_150M_UR50D", *, device: str | None = None,
+        cache: EmbeddingCache | None = None,
+    ) -> None:
+        self.model = model
+        self.device = device
+        self._cache = cache if cache is not None else EmbeddingCache()
+        self._tok: Any = None
+        self._mdl: Any = None
+        self._torch: Any = None
+
+    def _load(self) -> None:  # pragma: no cover - exercised only with the SDK
+        try:
+            import torch
+            from transformers import AutoModelForMaskedLM, AutoTokenizer  # type: ignore[import-not-found, import-untyped]
+        except ImportError as exc:
+            raise ImportError("ESM2Naturalness requires 'transformers' and 'torch'.") from exc
+        self._torch = torch
+        self._tok = AutoTokenizer.from_pretrained(self.model)
+        self._mdl = AutoModelForMaskedLM.from_pretrained(self.model).eval()
+
+    def score(self, sequences: Sequence[str]) -> np.ndarray:
+        keys = [_key(self.model, s) for s in sequences]
+        misses = [(i, s) for i, s in enumerate(sequences) if self._cache.get(keys[i]) is None]
+        if misses:  # pragma: no cover - exercised only with the SDK
+            if self._mdl is None:
+                self._load()
+            assert self._torch is not None and self._tok is not None and self._mdl is not None
+            with self._torch.no_grad():
+                for i, s in misses:
+                    enc = self._tok(s, return_tensors="pt")
+                    logits = self._mdl(**enc).logits[0]
+                    logp = self._torch.log_softmax(logits, dim=-1)
+                    ids = enc["input_ids"][0]
+                    tok_lp = logp[range(len(ids)), ids][1:-1]  # drop BOS/EOS
+                    self._cache.put(keys[i], np.asarray([float(tok_lp.mean())], dtype=float))
+        out = []
+        for k in keys:
+            v = self._cache.get(k)
+            assert v is not None
+            out.append(float(v[0]))
+        return np.asarray(out, dtype=float)
+
+
+class ESMCNaturalness:
+    """ESM-C naturalness via the EvolutionaryScale SDK (mean per-residue log-prob), import-guarded."""
+
+    def __init__(
+        self, model_name: str = "esmc_300m", *, device: str | None = None,
+        cache: EmbeddingCache | None = None,
+    ) -> None:
+        self.model_name = model_name
+        self.device = device
+        self._cache = cache if cache is not None else EmbeddingCache()
+        self._model: Any = None
+        self._torch: Any = None
+
+    def _load(self) -> None:  # pragma: no cover - exercised only with the SDK
+        try:
+            import torch
+            from esm.models.esmc import ESMC  # type: ignore[import-not-found, import-untyped]
+        except ImportError as exc:
+            raise ImportError("ESMCNaturalness requires the EvolutionaryScale 'esm' SDK.") from exc
+        self._torch = torch
+        dev = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self._model = ESMC.from_pretrained(self.model_name).to(dev).eval()
+
+    def score(self, sequences: Sequence[str]) -> np.ndarray:
+        keys = [_key(self.model_name, s) for s in sequences]
+        misses = [(i, s) for i, s in enumerate(sequences) if self._cache.get(keys[i]) is None]
+        if misses:  # pragma: no cover - exercised only with the SDK
+            if self._model is None:
+                self._load()
+            from esm.sdk.api import ESMProtein, LogitsConfig  # type: ignore[import-not-found, import-untyped]
+
+            assert self._torch is not None and self._model is not None
+            with self._torch.no_grad():
+                for i, s in misses:
+                    tensor = self._model.encode(ESMProtein(sequence=s))
+                    res = self._model.logits(tensor, LogitsConfig(sequence=True))
+                    seq_logits = res.logits.sequence
+                    if seq_logits.dim() == 3:
+                        seq_logits = seq_logits[0]
+                    logp = self._torch.log_softmax(seq_logits.float(), dim=-1)  # (L+2, V)
+                    ids = tensor.sequence.long()  # token ids incl. BOS/EOS
+                    pos = self._torch.arange(1, ids.shape[0] - 1)  # residues only
+                    tok_lp = logp[pos, ids[pos]]
+                    self._cache.put(keys[i], np.asarray([float(tok_lp.mean())], dtype=float))
+        out = []
+        for k in keys:
+            v = self._cache.get(k)
+            assert v is not None
+            out.append(float(v[0]))
+        return np.asarray(out, dtype=float)

shellde/oracle.py

@@ -0,0 +1,94 @@
+"""Oracles: the function-aligned information source (one Protocol, two impls).
+
+  - ``DMSLookupOracle``: ground-truth lookup over a measured (variant -> fitness)
+    table; ``universe`` is the measurable pool (real-DMS benchmarks).
+  - ``SyntheticOracle``: additive backbone + ``alpha`` * pairwise epistasis over a
+    design space, with the global optimum knowable from the enumerated ``universe``
+    (tunable-epistasis benchmarks). alpha=0 is purely additive.
+
+Both satisfy ``protocols.Oracle`` and expose an enumerable ``universe`` so the
+active-learning loop is uniform across them.
+"""
+from __future__ import annotations
+
+import itertools
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace, to_assignment
+
+_UNIVERSE_CAP = 200_000
+
+
+class DMSLookupOracle:
+    """Ground-truth fitness by lookup over a measured table."""
+
+    def __init__(self, table: Mapping[str, float]) -> None:
+        self._table = {str(k): float(v) for k, v in table.items()}
+        self._universe = list(self._table)
+
+    def evaluate(self, variants: Sequence[str]) -> np.ndarray:
+        try:
+            return np.array([self._table[str(v)] for v in variants], dtype=float)
+        except KeyError as exc:
+            raise KeyError(f"variant {exc} not in the DMS lookup table") from exc
+
+    @property
+    def universe(self) -> list[str] | None:
+        return list(self._universe)
+
+
+class SyntheticOracle:
+    """Additive + alpha*pairwise-epistatic landscape over a design space.
+
+    Fields h and couplings J are drawn once from ``seed``; ``alpha`` then scales the
+    pairwise term, so sweeping alpha walks from the additive null into progressively
+    stronger epistasis with the same underlying landscape.
+    """
+
+    def __init__(self, space: DesignSpace, *, alpha: float = 0.0, seed: int = 0) -> None:
+        self.space = space
+        self.alpha = alpha
+        n, q = space.n_positions, space.q
+        rng = np.random.default_rng(seed)
+        self._h = rng.standard_normal((n, q))
+        self._j = rng.standard_normal((n, n, q, q))
+        self._aa = {a: i for i, a in enumerate(space.alphabet)}
+        self._universe: list[str] | None = None
+
+    def evaluate(self, variants: Sequence[str]) -> np.ndarray:
+        space = self.space
+        n = space.n_positions
+        out = np.zeros(len(variants), dtype=float)
+        for k, v in enumerate(variants):
+            a = to_assignment(v, space)
+            idx = [self._aa[a[p]] for p in space.positions]
+            f = sum(self._h[i, idx[i]] for i in range(n))
+            if self.alpha != 0.0:
+                for i in range(n):
+                    for j in range(i + 1, n):
+                        f += self.alpha * self._j[i, j, idx[i], idx[j]]
+            out[k] = f
+        return out
+
+    @property
+    def universe(self) -> list[str] | None:
+        if self._universe is None:
+            space = self.space
+            total = space.q**space.n_positions
+            if total > _UNIVERSE_CAP:
+                raise ValueError(
+                    f"SyntheticOracle universe {total} exceeds cap {_UNIVERSE_CAP}; "
+                    "use a smaller design space for enumerated benchmarks"
+                )
+            self._universe = ["".join(c) for c in itertools.product(space.alphabet, repeat=space.n_positions)]
+        return list(self._universe)
+
+    def optimum(self) -> tuple[str, float]:
+        """The true global optimum over the enumerated universe (for regret)."""
+        uni = self.universe
+        assert uni is not None
+        f = self.evaluate(uni)
+        i = int(np.argmax(f))
+        return uni[i], float(f[i])

shellde/plm.py

@@ -0,0 +1,145 @@
+"""First-class PLM (ESM-C) front end for the real-data recommend path.
+
+At the N=95 constraint, explicit pairwise epistasis is unestimable (one pairwise
+term over q=20 is 19^2=361 params > 95 samples), so the lever that actually wins is
+a protein language model: ESM-C embeddings carry IMPLICIT epistasis from pretraining
+(attention makes them non-additive) and N=95 only fits a cheap readout on top.
+
+ESM-C CPU inference is slow and unbatched, so scoring a 200k-candidate pool with it
+is intractable AND wasteful (embeddings dilute on dense combinatorial libraries).
+The honest, correct integration is a TWO-STAGE re-rank: a cheap one-hot surrogate
+shortlists high-value candidates, then a PLM surrogate fit on the measured
+embeddings re-scores ONLY the shortlist. The PLM block earns its place on the front
+end, not on the bulk scoring. Embeds exactly ``len(measured) + top_k`` sequences.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from pathlib import Path
+
+import numpy as np
+
+from shellde.candidates import CandidatePool
+from shellde.design_space import DesignSpace
+from shellde.embeddings.provider import EmbeddingCache
+from shellde.features import (
+    FeatureBlock,
+    FeatureMatrix,
+    NaturalnessBlock,
+    OneHotBlock,
+    PLMEmbeddingBlock,
+)
+from shellde.protocols import EmbeddingProvider, NaturalnessProvider, Surrogate
+from shellde.surrogate import RidgeSurrogate
+
+
+def read_fasta(path: str | Path) -> str:
+    """Return the first record's sequence from a FASTA file (whitespace stripped)."""
+    seq: list[str] = []
+    started = False
+    for line in Path(path).read_text().splitlines():
+        if line.startswith(">"):
+            if started:
+                break
+            started = True
+            continue
+        seq.append(line.strip())
+    out = "".join(seq).replace(" ", "")
+    if not out:
+        raise ValueError(f"{path}: no sequence found in FASTA")
+    return out
+
+
+def build_provider(
+    kind: str, *, model: str = "", cache: EmbeddingCache | None = None
+) -> EmbeddingProvider:
+    """Construct an import-guarded PLM provider ('esmc' | 'esm2').
+
+    The heavy SDK import is lazy (inside the provider), so an unavailable SDK only
+    raises when ``embed`` is first called, with a clear install message.
+    """
+    if kind == "esmc":
+        from shellde.embeddings.esmc import ESMCProvider
+
+        return ESMCProvider(model or "esmc_300m", cache=cache)
+    if kind == "esm2":
+        from shellde.embeddings.esm2 import ESM2Provider
+
+        return ESM2Provider(model or "facebook/esm2_t30_150M_UR50D", cache=cache)
+    raise ValueError(f"unknown PLM kind {kind!r} (expected 'esmc' or 'esm2')")
+
+
+def build_naturalness(
+    kind: str, *, model: str = "", cache: EmbeddingCache | None = None
+) -> NaturalnessProvider:
+    """Construct an import-guarded zero-shot naturalness provider ('esmc' | 'esm2')."""
+    if kind == "esmc":
+        from shellde.naturalness import ESMCNaturalness
+
+        return ESMCNaturalness(model or "esmc_300m", cache=cache)
+    if kind == "esm2":
+        from shellde.naturalness import ESM2Naturalness
+
+        return ESM2Naturalness(model or "facebook/esm2_t30_150M_UR50D", cache=cache)
+    raise ValueError(f"unknown naturalness kind {kind!r} (expected 'esmc' or 'esm2')")
+
+
+def plm_rerank(
+    pool: CandidatePool,
+    space: DesignSpace,
+    measured_variants: Sequence[str],
+    fitness: np.ndarray,
+    *,
+    provider: EmbeddingProvider | None = None,
+    naturalness: "NaturalnessProvider | None" = None,
+    wt_sequence: str,
+    top_k: int,
+    surrogate_factory: Callable[[], Surrogate] = lambda: RidgeSurrogate(random_state=0),
+    use_onehot: bool = True,
+    plm_pooling: str = "mean",
+) -> CandidatePool:
+    """Stage-2: refit a PLM/naturalness surrogate on measured data, re-score the shortlist.
+
+    Adds the (cost-bearing) PLM embedding and/or naturalness blocks only on the bounded
+    top-``top_k`` shortlist, not the 200k-candidate base scoring. At least one of
+    ``provider`` (embedding) or ``naturalness`` must be given. ``stats['wt_mean']`` is
+    the refit surrogate's WT prediction (for risk flagging). Embeds/scores only
+    ``len(measured) + top_k`` sequences, so CPU PLM inference stays tractable.
+    """
+    if provider is None and naturalness is None:
+        raise ValueError("plm_rerank needs at least one of provider (embedding) or naturalness")
+    if not pool.variants:
+        return pool
+    k = max(1, min(top_k, len(pool.variants)))
+    shortlist = list(pool.variants[:k])
+    mut_count = list(pool.mut_count[:k])
+
+    blocks: list[FeatureBlock] = []
+    if use_onehot:
+        blocks.append(OneHotBlock(space))
+    dim = 0
+    if provider is not None:
+        blocks.append(PLMEmbeddingBlock(space, provider, wt_sequence, pooling=plm_pooling))
+        dim += int(provider.dim)
+    if naturalness is not None:
+        blocks.append(NaturalnessBlock(space, naturalness, wt_sequence))
+        dim += 1
+    matrix = FeatureMatrix(space, blocks)
+
+    measured = [str(v) for v in measured_variants]
+    sur = surrogate_factory().fit(matrix.encode(measured), np.asarray(fitness, dtype=float))
+    pred = sur.predict(matrix.encode(shortlist))
+    wt_mean = float(sur.predict(matrix.encode([space.wt()])).mean[0])
+
+    stats = dict(pool.stats)
+    stats.update({
+        "plm_reranked": k, "plm_dim": dim, "wt_mean": wt_mean,
+        "plm_embedding": provider is not None, "plm_naturalness": naturalness is not None,
+    })
+    return CandidatePool(
+        variants=shortlist,
+        pred=pred,
+        mut_count=mut_count,
+        pool_positions=pool.pool_positions,
+        stats=stats,
+    )

shellde/prereg.py

@@ -0,0 +1,41 @@
+"""Pre-registration: a locked protocol + deterministic hash stamped into outputs.
+
+Locking the protocol once (metric, budget, seeds, statistics, decision rule) before
+running is the falsifiability guard: results are interpreted against a fixed bar, not
+a metric/seed/universe chosen after seeing the numbers.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+PROTOCOL: dict[str, Any] = {
+    "version": 1,
+    "primary_metric": "normalised_simple_regret",  # (true_max - best_found)/(true_max - true_min)
+    "secondary_metrics": ["best_found", "heldout_spearman"],
+    "n_seeds": 20,
+    "budget": {"n_init": 95, "batch_size": 95, "n_rounds": 3},
+    "alpha_spectrum": [0.0, 0.5, 1.0, 1.5, 2.0],
+    "statistics": {"alpha": 0.05, "correction": "benjamini_hochberg", "paired": True},
+    "methods": ["ours", "additive_greedy", "rf_greedy", "ucb_bo"],
+    "decision_rule": (
+        "ours must never be FDR-significantly worse than a baseline on any landscape; "
+        "superiority is claimed only on an FDR-significant paired regret reduction with "
+        "CI lower bound > 0."
+    ),
+}
+
+
+def prereg_hash(protocol: dict[str, Any] | None = None) -> str:
+    """Deterministic SHA-256 over the canonical protocol JSON (sorted keys)."""
+    blob = json.dumps(protocol if protocol is not None else PROTOCOL, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(blob.encode()).hexdigest()
+
+
+def stamp(payload: dict[str, Any]) -> dict[str, Any]:
+    """Return a copy of ``payload`` stamped with the protocol + its hash."""
+    out = dict(payload)
+    out["protocol"] = PROTOCOL
+    out["prereg_hash"] = prereg_hash()
+    return out

shellde/protocols.py

@@ -0,0 +1,87 @@
+"""Structural typing contracts (typing.Protocol) for the pluggable pieces.
+
+These pin the seams that make ShellDE one coherent pipeline instead of two
+lineages: a single Surrogate contract (calibrated std), a single Acquisition
+contract, a single Oracle contract, and an EmbeddingProvider for PLM features.
+Concrete implementations live in their own modules and need no registration.
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.types import Prediction
+
+
+@runtime_checkable
+class Surrogate(Protocol):
+    """A fitted fitness model over a feature matrix; returns calibrated uncertainty."""
+
+    def fit(
+        self, x: np.ndarray, y: np.ndarray, sample_weight: np.ndarray | None = None
+    ) -> "Surrogate":
+        ...
+
+    def predict(self, x: np.ndarray) -> Prediction:
+        """Return mean + CALIBRATED predictive std (1 sigma)."""
+        ...
+
+    def main_effects(self) -> np.ndarray | None:
+        """Additive per-column coefficients for the position-selection prior (or None)."""
+        ...
+
+
+@runtime_checkable
+class Acquisition(Protocol):
+    """Selects which candidates to measure next, given surrogate predictions."""
+
+    def select(
+        self,
+        pred: Prediction,
+        k: int,
+        *,
+        candidates: Sequence[str],
+        space: DesignSpace,
+    ) -> list[int]:
+        """Return indices into ``candidates`` for the chosen batch of size <= k."""
+        ...
+
+
+@runtime_checkable
+class Oracle(Protocol):
+    """Maps variants to (measured) fitness; the only function-aligned information source."""
+
+    def evaluate(self, variants: Sequence[str]) -> np.ndarray:
+        ...
+
+    @property
+    def universe(self) -> list[str] | None:
+        """The measurable variant pool for lookup oracles, else None (enumerable space)."""
+        ...
+
+
+@runtime_checkable
+class EmbeddingProvider(Protocol):
+    """Produces per-sequence embeddings (PLM); implementations cache on disk."""
+
+    dim: int
+
+    def embed(self, sequences: Sequence[str]) -> np.ndarray:
+        """Return an (len(sequences), dim) float array of mean-pooled embeddings."""
+        ...
+
+
+@runtime_checkable
+class NaturalnessProvider(Protocol):
+    """Produces a per-sequence zero-shot naturalness score (PLM likelihood-style).
+
+    Higher means more 'natural' under the model's pretraining prior. Used as a
+    warm-start signal (a feature and/or discovery scorer); implementations cache.
+    """
+
+    def score(self, sequences: Sequence[str]) -> np.ndarray:
+        """Return a (len(sequences),) float array of per-sequence naturalness."""
+        ...

shellde/rank.py

@@ -0,0 +1,103 @@
+"""Rank a candidate pool into per-mutation-count tiers + a global ranking.
+
+Each ranked entry carries evidence: predicted mean/std, acquisition value, the
+mutation tokens, the baseline gain over WT, and a CALIBRATED risk flag. The risk
+flag is the opportunity-cost framing made concrete: a recommendation is "confident"
+when its predicted gain over WT exceeds one calibrated standard deviation, and
+"speculative" (the predicted gain is within the calibrated sigma, so the ranking of
+that candidate is not resolved by the current data) otherwise. It is a gain-vs-sigma
+comparison only; it does not detect or classify epistasis.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from shellde.acquisition import _ScoreAcquisition
+from shellde.candidates import CandidatePool
+from shellde.design_space import DesignSpace, mutations_of
+from shellde.types import Candidate
+
+
+@dataclass
+class RankingResult:
+    """Per-mutation-count tiers + a global (count-mixing) exploratory ranking."""
+
+    per_count: dict[int, list[Candidate]] = field(default_factory=dict)
+    global_ranking: list[Candidate] = field(default_factory=list)
+    global_is_exploratory: bool = True
+
+
+def _risk_flag(gain: float, sigma: float) -> str:
+    return "confident" if gain >= max(sigma, 0.0) else "speculative"
+
+
+def rank_candidates(
+    pool: CandidatePool,
+    space: DesignSpace,
+    *,
+    acquisition: _ScoreAcquisition,
+    per_count_k: int = 10,
+    global_k: int = 20,
+    wt_mean: float = 0.0,
+    conformal_q: float | None = None,
+    max_measured_nmut: int | None = None,
+) -> RankingResult:
+    """Build per-count tiers + a global ranking from a scored candidate pool.
+
+    When ``conformal_q`` is given (the normalized cross-conformal quantile from the
+    MEASURED set), each candidate gets an ADDITIVE conformal abstain gate on its
+    WT-relative gain: ``conformal_lo = gain - conformal_q * sigma`` and ``abstain`` is
+    True when that lower end is <= 0 (cannot conclude the candidate beats WT at this
+    conformal level). ``conformal_scope`` is ``"conformal"`` when the candidate is
+    in-distribution (``n_mut <= max_measured_nmut``) and ``"conformal_oob"`` when it is a
+    higher-order combination out of the calibration distribution. The sigma-based
+    ``risk`` flag is kept exactly as-is. When ``conformal_q`` is None the evidence dict is
+    byte-identical to the no-conformal behaviour (no new keys).
+    """
+    n = len(pool.variants)
+    if n == 0:
+        return RankingResult({}, [], True)
+    scores = np.asarray(acquisition.score(pool.pred), dtype=float)
+    mean = np.asarray(pool.pred.mean, dtype=float)
+    std = np.asarray(pool.pred.std, dtype=float)
+
+    entries: list[Candidate] = []
+    for i, v in enumerate(pool.variants):
+        gain = float(mean[i] - wt_mean)
+        sigma = float(std[i])
+        n_mut = int(pool.mut_count[i])
+        evidence: dict = {"baseline_gain": gain, "risk": _risk_flag(gain, sigma)}
+        if conformal_q is not None:
+            conformal_lo = gain - conformal_q * sigma
+            evidence["conformal_lo"] = conformal_lo
+            evidence["abstain"] = bool(conformal_lo <= 0.0)
+            evidence["conformal_scope"] = (
+                "conformal_oob"
+                if (max_measured_nmut is not None and n_mut > max_measured_nmut)
+                else "conformal"
+            )
+        entries.append(
+            Candidate(
+                variant=v,
+                n_mut=n_mut,
+                mu=float(mean[i]),
+                sigma=sigma,
+                acq=float(scores[i]),
+                mutations=mutations_of(v, space),
+                evidence=evidence,
+            )
+        )
+
+    by_acq = sorted(entries, key=lambda c: c.acq, reverse=True)
+    per_count: dict[int, list[Candidate]] = {}
+    for c in by_acq:
+        per_count.setdefault(c.n_mut, [])
+        if len(per_count[c.n_mut]) < per_count_k:
+            per_count[c.n_mut].append(c)
+    return RankingResult(
+        per_count=dict(sorted(per_count.items())),
+        global_ranking=by_acq[:global_k],
+        global_is_exploratory=True,
+    )