shellde 0.2.0__py3-none-any.whl

shellde/campaign.py

@@ -0,0 +1,257 @@
+"""funclib -> active-learning campaign orchestrator (the validated D18->D19 handoff).
+
+This wires the two halves the tool already ships into one flow:
+
+  R0 (cold start): ``funclib.design_library`` proposes the active-site measurement library
+      -- the seed plate -- using only zero-shot / structure signals (no fitness yet). This is
+      where prediction is irreplaceable: with zero measurements the AL surrogate is blind.
+
+  R1.. (engine): measurement-driven active learning takes over. Each round re-fits a fresh
+      surrogate on ALL accumulated measurements and proposes the next plate over the FULL
+      saturation universe of the design positions -- NOT hard-restricted to the funclib
+      library (D19: the non-natural winner is frequently outside the seed; restricting the
+      candidate space caps you below it). funclib only SEEDS; AL explores past it.
+
+``simulate_campaign`` runs this against a lookup oracle (a measured landscape CSV) so the
+handoff can be validated/reproduced (D19/D25). The live R0 plate generation lives in the CLI
+(``shellde campaign`` without ``--simulate``).
+
+Honest scope: this does not predict winners. It increases the probability/efficiency of
+*finding* a better variant per experimental round, conditional on the winner staying in the
+saturation universe and the active-site framing being right (see D19/D25/D27).
+"""
+from __future__ import annotations
+
+import statistics
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from shellde.acquisition import Greedy, UCB
+from shellde.design_space import DesignSpace
+from shellde.features.defaults import default_blocks
+from shellde.features.matrix import FeatureMatrix
+from shellde.funclib import FuncLibLibrary, design_library
+from shellde.loop import run_campaign
+from shellde.oracle import DMSLookupOracle
+from shellde.protocols import Surrogate
+from shellde.surrogate import (
+    EnsembleSurrogate,
+    GlobalEpistasisSurrogate,
+    RankingSurrogate,
+    RFSurrogate,
+    RidgeSurrogate,
+)
+
+_SURROGATE_CLASSES = {
+    "ridge": RidgeSurrogate, "ensemble": EnsembleSurrogate, "rf": RFSurrogate,
+    "ranking": RankingSurrogate, "global_epistasis": GlobalEpistasisSurrogate,
+}
+
+
+@dataclass(frozen=True)
+class SurrogateSpec:
+    """Picklable surrogate factory (a class+kwargs, unlike a lambda) so simulate_campaign can run
+    across processes. Callable like the lambdas the serial path also accepts."""
+
+    name: str
+    random_state: int = 0
+
+    def __call__(self) -> Surrogate:
+        return _SURROGATE_CLASSES[self.name](random_state=self.random_state)
+
+
+def funclib_seed_variants(
+    space: DesignSpace,
+    *,
+    plate: int,
+    ddg: Mapping[int, Mapping[str, float]] | None = None,
+    tolerance: Mapping[str, float] | None = None,
+    **design_kwargs: object,
+) -> tuple[list[str], FuncLibLibrary]:
+    """Build the funclib R0 seed plate: top-``plate`` combinable library variants.
+
+    Returns (combo strings sorted by funclib combo score, the FuncLibLibrary for provenance).
+    On abstain (no signal / everything filtered) the variant list is empty.
+    """
+    lib = design_library(space, ddg=ddg, tolerance=tolerance, **design_kwargs)  # type: ignore[arg-type]
+    variants = [c.variant for c in lib.candidates[:plate]]
+    return variants, lib
+
+
+def single_mutant_variants(space: DesignSpace) -> list[str]:
+    """All single-substitution combo strings over the design positions (EVOLVEpro-style seed)."""
+    wt = space.wt()
+    out: list[str] = []
+    for i in range(space.n_positions):
+        for aa in space.alphabet:
+            if aa == wt[i]:
+                continue
+            out.append(wt[:i] + aa + wt[i + 1 :])
+    return out
+
+
+def stable_variants(
+    variants: Sequence[str],
+    space: DesignSpace,
+    ddg: Mapping[int, Mapping[str, float]],
+    *,
+    ddg_cutoff: float = 2.5,
+    additive_budget: float | None = None,
+    allow_unscored: bool = False,
+) -> list[str]:
+    """Keep combo variants whose every mutation passes the folding-ddG stability gate.
+
+    A mutation at position p->aa passes iff ``ddg[p][aa] <= ddg_cutoff``; WT residues contribute 0.
+    If ``additive_budget`` is set, the summed single-mutation ddG must also be <= the budget
+    (the validated additive-destabilization proxy, D16). This makes predicted stability a PERSISTENT
+    constraint across the whole campaign -- the honest multi-objective decomposition for the
+    asymmetric-data regime (D14/D15: stability is predictable -> constraint; activity needs
+    measurement -> AL objective), not a two-measured-objective Pareto.
+    """
+    wt = space.wt()
+    out: list[str] = []
+    for v in variants:
+        total = 0.0
+        ok = True
+        for i, pos in enumerate(space.positions):
+            aa = v[i]
+            if aa == wt[i]:
+                continue
+            d = ddg.get(pos, {}).get(aa)
+            if d is None:
+                if allow_unscored:
+                    continue
+                ok = False
+                break
+            if d > ddg_cutoff:
+                ok = False
+                break
+            total += d
+        if ok and (additive_budget is None or total <= additive_budget):
+            out.append(v)
+    return out
+
+
+@dataclass(frozen=True)
+class SimulationResult:
+    """Aggregate of a multi-seed simulated campaign against a lookup oracle."""
+
+    strategy: str
+    n_seeds: int
+    winner_reach_rate: float
+    rounds_to_winner_median: float | None
+    best_found_mean: float
+    best_found_percentile_mean: float
+    per_seed: list[dict]
+
+
+def _percentile_of(value: float, sorted_fitness: np.ndarray) -> float:
+    """Fraction of landscape variants with fitness <= ``value`` (empirical percentile)."""
+    return float(np.searchsorted(sorted_fitness, value, side="right")) / float(len(sorted_fitness))
+
+
+_WORKER_STATE: dict = {}
+
+
+def _seed_result(
+    s: int, oracle: DMSLookupOracle, space: DesignSpace, seed_variants, make_surrogate,
+    blocks, beta: float, plate: int, rounds: int,
+) -> tuple[list[float], str, float]:
+    """One independent seeded campaign -> (trajectory, best_variant, best_fitness). Deterministic in s."""
+    acquisition = Greedy() if beta <= 0.0 else UCB(beta)
+    res = run_campaign(
+        oracle, space,
+        make_matrix=lambda: FeatureMatrix(space, blocks(space)),
+        make_surrogate=make_surrogate, acquisition=acquisition,
+        n_init=plate, batch_size=plate, n_rounds=rounds, seed=s,
+        init_variants=list(seed_variants) if seed_variants is not None else None,
+    )
+    return [float(x) for x in res.trajectory], res.best_variant, float(res.best_fitness)
+
+
+def _init_worker(table, space, seed_variants, make_surrogate, blocks, beta, plate, rounds) -> None:
+    _WORKER_STATE.clear()
+    _WORKER_STATE.update(
+        oracle=DMSLookupOracle(table), space=space, seed_variants=seed_variants,
+        make_surrogate=make_surrogate, blocks=blocks, beta=beta, plate=plate, rounds=rounds,
+    )
+
+
+def _seed_worker(s: int) -> tuple[list[float], str, float]:
+    w = _WORKER_STATE
+    return _seed_result(s, w["oracle"], w["space"], w["seed_variants"], w["make_surrogate"],
+                        w["blocks"], w["beta"], w["plate"], w["rounds"])
+
+
+def simulate_campaign(
+    table: Mapping[str, float],
+    space: DesignSpace,
+    *,
+    seed_variants: Sequence[str] | None,
+    strategy: str,
+    plate: int = 95,
+    rounds: int = 3,
+    n_seeds: int = 20,
+    beta: float = 0.0,
+    make_surrogate: Callable[[], Surrogate],
+    make_blocks: Callable[[DesignSpace], list] | None = None,
+    winner_tol: float = 1e-9,
+    n_jobs: int = 1,
+) -> SimulationResult:
+    """Run the funclib->AL handoff against a measured landscape (oracle = lookup table).
+
+    ``seed_variants`` is the R0 seed (funclib library or singles); ``None`` = random R0.
+    The AL pool is the full landscape universe each round, so the loop explores past the seed.
+    ``beta``=0 is greedy exploitation (D25: greedy was not beaten by UCB on a good seed); >0 is UCB.
+    """
+    blocks = make_blocks or default_blocks
+    # Each seed is an independent, deterministic campaign (run_campaign(seed=s)). Serial by default;
+    # n_jobs>1 runs them across processes (NOT threads: the loop body is Python/GIL-bound). The
+    # aggregate is collected in seed order and percentiles/winner are computed centrally, so a
+    # parallel run is byte-identical to serial. Parallel needs picklable make_surrogate/make_blocks
+    # (use SurrogateSpec / the default blocks), not lambdas.
+    if n_jobs and n_jobs > 1 and n_seeds > 1:
+        import multiprocessing as _mp
+        from concurrent.futures import ProcessPoolExecutor
+        # spawn (not fork): the parent may hold BLAS threads, and fork()-with-threads can deadlock.
+        with ProcessPoolExecutor(
+            max_workers=min(n_jobs, n_seeds),
+            mp_context=_mp.get_context("spawn"),
+            initializer=_init_worker,
+            initargs=(dict(table), space, seed_variants, make_surrogate, blocks, beta, plate, rounds),
+        ) as ex:
+            raw = list(ex.map(_seed_worker, range(n_seeds)))
+    else:
+        oracle = DMSLookupOracle(table)
+        raw = [_seed_result(s, oracle, space, seed_variants, make_surrogate, blocks, beta, plate, rounds)
+               for s in range(n_seeds)]
+
+    fitness = np.asarray(list(table.values()), dtype=float)
+    sorted_fitness = np.sort(fitness)
+    global_max = float(fitness.max())
+    per_seed: list[dict] = []
+    for s, (traj, best_variant, best_fitness) in enumerate(raw):
+        # trajectory[0] is best-after-seed (R0); trajectory[k] is best after AL round k.
+        reached = next((i for i, b in enumerate(traj) if b >= global_max - winner_tol), None)
+        per_seed.append({
+            "rng_seed": s,
+            "best_found": best_fitness,
+            "best_variant": best_variant,
+            "best_percentile": _percentile_of(best_fitness, sorted_fitness),
+            "winner_reached": reached is not None,
+            "rounds_to_winner": reached,
+            "trajectory": traj,
+        })
+
+    reaches = [p["rounds_to_winner"] for p in per_seed if p["winner_reached"]]
+    return SimulationResult(
+        strategy=strategy,
+        n_seeds=n_seeds,
+        winner_reach_rate=sum(1 for p in per_seed if p["winner_reached"]) / n_seeds,
+        rounds_to_winner_median=(statistics.median(reaches) if reaches else None),
+        best_found_mean=float(np.mean([p["best_found"] for p in per_seed])),
+        best_found_percentile_mean=float(np.mean([p["best_percentile"] for p in per_seed])),
+        per_seed=per_seed,
+    )

shellde/candidates.py

@@ -0,0 +1,341 @@
+"""Candidate generation: pick positions, then lazily enumerate up-to-N combos.
+
+Two public entry points:
+  - ``select_positions``: narrow a large design space to the top-m positions, by
+    additive main-effect spread (cheap) or an interaction-aware Sobol total-effect
+    estimate (for epistatic targets; additive importance misses interaction-only
+    positions).
+  - ``generate_candidates``: lazily enumerate high-value up-to-``max_mut`` mutation
+    combinations under hard caps (never iterates the Cartesian product), add a random
+    exploration quota so non-additive optima are reachable, then prefilter to a
+    scorable pool with the calibrated surrogate. Returns a ``CandidatePool`` carrying
+    the surrogate Prediction so downstream ranking does not re-predict.
+"""
+from __future__ import annotations
+
+import heapq
+import itertools
+import time
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.features.matrix import FeatureMatrix
+from shellde.protocols import Surrogate
+from shellde.types import Prediction
+
+GENERATED_CAP = 200_000
+SCORED_CAP = 20_000
+DEFAULT_M = 15
+DEFAULT_MAX_MUT = 5
+DEFAULT_EXPLORE_FRAC = 0.2
+DEFAULT_BETA = 1.0
+DEFAULT_TIME_BUDGET_S = 60.0
+
+
+@dataclass
+class CandidatePool:
+    """Prefiltered candidate variants + the surrogate prediction on them."""
+
+    variants: list[str]
+    pred: Prediction
+    mut_count: list[int]
+    pool_positions: list[int]
+    stats: dict = field(default_factory=dict)
+
+
+# --------------------------------------------------------------------------- #
+# Position selection
+# --------------------------------------------------------------------------- #
+def _additive_scores(surrogate: Surrogate, matrix: FeatureMatrix) -> np.ndarray:
+    """Per-(position, AA) additive main-effect, shape (n_positions, q); zeros if no onehot."""
+    space = matrix.space
+    npos, q = space.n_positions, space.q
+    coef = surrogate.main_effects()
+    oh = next((se for name, *se in matrix.column_layout() if name == "onehot"), None)
+    if coef is None or oh is None:
+        return np.zeros((npos, q), dtype=float)
+    start, end = oh
+    block = np.asarray(coef, dtype=float)[start:end]
+    if block.size != npos * q:
+        return np.zeros((npos, q), dtype=float)
+    return block.reshape(npos, q)
+
+
+def _interaction_importance(
+    surrogate: Surrogate, matrix: FeatureMatrix, *, n_backgrounds: int = 64, seed: int = 0
+) -> np.ndarray:
+    """Sobol-style total-effect importance per position via the full surrogate.
+
+    Vary each position over the alphabet on RANDOM backgrounds and average the
+    prediction variance across residues; random backgrounds make the estimate
+    include interaction effects (a position important only through epistasis still
+    scores high), unlike the additive main-effect spread.
+    """
+    space = matrix.space
+    npos, q = space.n_positions, space.q
+    alphabet = space.alphabet
+    rng = np.random.default_rng(seed)
+    nb = min(n_backgrounds, max(1, q**npos))
+    bg = rng.integers(0, q, size=(nb, npos))
+    imp = np.zeros(npos, dtype=float)
+    for j in range(npos):
+        preds = np.zeros((nb, q), dtype=float)
+        for ai in range(q):
+            variants = []
+            for b in range(nb):
+                chars = [alphabet[int(bg[b, k])] for k in range(npos)]
+                chars[j] = alphabet[ai]
+                variants.append("".join(chars))
+            preds[:, ai] = np.asarray(surrogate.predict(matrix.encode(variants)).mean, dtype=float)
+        imp[j] = float(np.mean(preds.var(axis=1)))
+    return imp
+
+
+def select_positions(
+    surrogate: Surrogate,
+    matrix: FeatureMatrix,
+    *,
+    m: int,
+    interaction_aware: bool = False,
+    seed: int = 0,
+) -> list[int]:
+    """Top-m designed positions by importance.
+
+    Default importance is the additive main-effect spread (cheap, but MISSES a
+    position that matters only through epistasis). ``interaction_aware=True`` ranks by
+    a Sobol total-effect estimate over the full surrogate. ``m >= n_positions`` keeps
+    all positions (the safe choice for small combinatorial design spaces).
+    """
+    space = matrix.space
+    if interaction_aware:
+        importance = _interaction_importance(surrogate, matrix, seed=seed)
+    else:
+        importance = _additive_scores(surrogate, matrix).var(axis=1)
+    order = np.argsort(importance)[::-1][: max(1, min(m, space.n_positions))]
+    return sorted(space.positions[int(j)] for j in order)
+
+
+# --------------------------------------------------------------------------- #
+# Combinatorial generation
+# --------------------------------------------------------------------------- #
+def _assignment_to_combo(mut: tuple[tuple[int, str], ...], space: DesignSpace) -> str:
+    chosen = dict(space.reference)
+    for pos, aa in mut:
+        chosen[pos] = aa
+    return "".join(chosen[p] for p in space.positions)
+
+
+def _mut_count(variant: str, space: DesignSpace) -> int:
+    ref = space.wt()
+    return sum(1 for a, b in zip(variant, ref, strict=True) if a != b)
+
+
+def _generate_additive(
+    add: np.ndarray, space: DesignSpace, pool: list[int], *, max_mut: int, gen_cap: int, deadline: float
+) -> set[str]:
+    """Lazy best-first enumeration of high-gain <= max_mut combos (no Cartesian sweep)."""
+    aa_to_int = {a: i for i, a in enumerate(space.alphabet)}
+    subs: dict[int, list[tuple[str, float]]] = {}
+    for p in pool:
+        j = space.position_index(p)
+        wt_score = add[j, aa_to_int[space.reference[p]]]
+        subs[p] = sorted(
+            ((a, float(add[j, aa_to_int[a]] - wt_score)) for a in space.alphabet if a != space.reference[p]),
+            key=lambda kv: kv[1],
+            reverse=True,
+        )
+    gain_of = {p: dict(subs[p]) for p in pool}
+    out: set[str] = set()
+    visited: set[tuple[tuple[int, str], ...]] = {()}
+    visit_cap = max(4 * gen_cap, 1000)
+    heap: list[tuple[float, tuple[tuple[int, str], ...]]] = [(0.0, ())]
+    while heap and len(out) < gen_cap:
+        if time.perf_counter() > deadline:
+            break
+        neg, mut = heapq.heappop(heap)
+        out.add(_assignment_to_combo(mut, space))
+        used = {p for p, _ in mut}
+        if len(mut) < max_mut and len(visited) < visit_cap:
+            for p in pool:
+                if p in used or not subs[p]:
+                    continue
+                a, g = subs[p][0]
+                nm = tuple(sorted((*mut, (p, a))))
+                if nm not in visited:
+                    visited.add(nm)
+                    heapq.heappush(heap, (neg - g, nm))
+        if len(visited) < visit_cap:
+            for i, (p, a) in enumerate(mut):
+                ranked = subs[p]
+                k = next((idx for idx, (aa, _) in enumerate(ranked) if aa == a), -1)
+                if 0 <= k and k + 1 < len(ranked):
+                    na, gna = ranked[k + 1]
+                    nm = tuple(sorted((*mut[:i], (p, na), *mut[i + 1 :])))
+                    if nm not in visited:
+                        visited.add(nm)
+                        heapq.heappush(heap, (neg + gain_of[p][a] - gna, nm))
+    return out
+
+
+def _explore_sample(
+    space: DesignSpace, pool: list[int], *, max_mut: int, n: int, rng: np.random.Generator, deadline: float
+) -> set[str]:
+    """Random non-additive <= max_mut combos over the pool (exploration quota)."""
+    out: set[str] = set()
+    subs = {p: [a for a in space.alphabet if a != space.reference[p]] for p in pool}
+    attempts = 0
+    while len(out) < n and attempts < 20 * n + 50:
+        if time.perf_counter() > deadline:
+            break
+        attempts += 1
+        k = min(int(rng.integers(1, max_mut + 1)), len(pool))
+        chosen_pos = rng.choice(len(pool), size=k, replace=False)
+        mut = tuple(
+            sorted((pool[int(i)], subs[pool[int(i)]][int(rng.integers(0, len(subs[pool[int(i)]])))]) for i in chosen_pos)
+        )
+        out.add(_assignment_to_combo(mut, space))
+    return out
+
+
+def generate_candidates(
+    surrogate: Surrogate,
+    matrix: FeatureMatrix,
+    *,
+    m: int = DEFAULT_M,
+    max_mut: int = DEFAULT_MAX_MUT,
+    pinned_positions: Sequence[int] | None = None,
+    interaction_aware: bool = False,
+    generated_cap: int = GENERATED_CAP,
+    scored_cap: int = SCORED_CAP,
+    explore_frac: float = DEFAULT_EXPLORE_FRAC,
+    beta: float = DEFAULT_BETA,
+    time_budget_s: float = DEFAULT_TIME_BUDGET_S,
+    seed: int = 0,
+) -> CandidatePool:
+    """Generate + prefilter up-to-``max_mut`` mutation candidates under hard caps."""
+    if max_mut < 1:
+        raise ValueError("max_mut must be >= 1")
+    if not 0.0 <= explore_frac < 1.0:
+        raise ValueError("explore_frac must be in [0, 1)")
+    space = matrix.space
+    rng = np.random.default_rng(seed)
+    t0 = time.perf_counter()
+    deadline = t0 + time_budget_s
+
+    if pinned_positions is not None:
+        pool = sorted(int(p) for p in pinned_positions)
+        bad = [p for p in pool if p not in space.positions]
+        if bad:
+            raise ValueError(f"pinned positions {bad} not in design space")
+    else:
+        pool = select_positions(surrogate, matrix, m=m, interaction_aware=interaction_aware, seed=seed)
+
+    add = _additive_scores(surrogate, matrix)
+    n_explore = int(round(generated_cap * explore_frac))
+    generated = _generate_additive(
+        add, space, pool, max_mut=max_mut, gen_cap=max(0, generated_cap - n_explore), deadline=deadline
+    )
+    generated |= _explore_sample(space, pool, max_mut=max_mut, n=n_explore, rng=rng, deadline=deadline)
+    generated.discard(space.wt())  # never let WT dominate the tier ranking
+    if len(generated) > generated_cap:
+        generated = set(sorted(generated)[:generated_cap])
+
+    variants = sorted(generated)
+    if not variants:
+        empty = Prediction(np.zeros(0), np.zeros(0))
+        return CandidatePool([], empty, [], pool, {"n_generated": 0, "n_scored": 0, "n_pool_positions": len(pool)})
+
+    # full-surrogate prefilter (mean + beta*calibrated std), batched, deadline-bounded
+    n_total = len(variants)
+    mu_parts: list[np.ndarray] = []
+    sd_parts: list[np.ndarray] = []
+    for i in range(0, n_total, 4096):
+        if time.perf_counter() > deadline:
+            break
+        chunk = variants[i : i + 4096]
+        pr = surrogate.predict(matrix.encode(chunk))
+        mu_parts.append(np.asarray(pr.mean, dtype=float))
+        sd_parts.append(np.asarray(pr.std, dtype=float))
+    mu = np.concatenate(mu_parts) if mu_parts else np.zeros(0)
+    sigma = np.concatenate(sd_parts) if sd_parts else np.zeros(0)
+    variants = variants[: mu.shape[0]]
+    acq = mu + beta * sigma
+    keep = np.argsort(acq)[::-1][:scored_cap]
+    order = keep[np.argsort(acq[keep])[::-1]]
+    sel = [variants[int(i)] for i in order]
+    return CandidatePool(
+        variants=sel,
+        pred=Prediction(mu[order], sigma[order]),
+        mut_count=[_mut_count(v, space) for v in sel],
+        pool_positions=pool,
+        stats={
+            "n_generated": n_total,
+            "n_scored": len(sel),
+            "n_pool_positions": len(pool),
+            "wall_clock_s": time.perf_counter() - t0,
+        },
+    )
+
+
+def recombine_beneficials(
+    space: DesignSpace,
+    variants: Sequence[str],
+    fitness: Sequence[float],
+    *,
+    wt_fitness: float | None = None,
+    min_gain: float = 0.0,
+    max_mut: int = 4,
+    cap: int = GENERATED_CAP,
+) -> list[str]:
+    """Recombine the mutations of MEASURED variants that beat WT (Arnold-style recombination).
+
+    The evidence-backed strategy (broad measurement then recombine confirmed beneficials, not
+    predictor ranking): pool every substitution that appears in a measured variant with
+    fitness > ``wt_fitness`` (+ ``min_gain``), then enumerate combinations of those beneficial
+    substitutions (one AA per position) up to ``max_mut``, ordered by observed gain, capped.
+
+    ``wt_fitness`` defaults to the measured WT combo's fitness if present, else the 75th percentile
+    of the measured fitness (a "clearly-above-typical" proxy, since true WT is unmeasured). Returns
+    combo strings over ``space.positions`` (WT elsewhere), excluding WT itself.
+    """
+    wt = space.wt()
+    fit = {str(v): float(f) for v, f in zip(variants, fitness, strict=True)}
+    if wt_fitness is None:
+        wt_fitness = fit.get(wt, float(np.quantile(np.asarray(list(fit.values()), dtype=float), 0.75)))
+    thr = wt_fitness + min_gain
+    # best observed gain per (position-index, aa) among beneficial variants
+    mut_gain: dict[tuple[int, str], float] = {}
+    for v, f in fit.items():
+        if f <= thr or len(v) != len(wt):
+            continue
+        gain = f - wt_fitness
+        for i, (a, w) in enumerate(zip(v, wt, strict=True)):
+            if a != w:
+                key = (i, a)
+                if gain > mut_gain.get(key, float("-inf")):
+                    mut_gain[key] = gain
+    if not mut_gain:
+        return []
+    per_pos: dict[int, list[str]] = {}
+    for (i, a), _g in sorted(mut_gain.items(), key=lambda kv: -kv[1]):
+        per_pos.setdefault(i, []).append(a)
+    pos_order = sorted(per_pos, key=lambda i: -max(mut_gain[(i, a)] for a in per_pos[i]))
+    out: list[str] = []
+    seen: set[str] = set()
+    for k in range(1, min(max_mut, len(pos_order)) + 1):
+        for pos_combo in itertools.combinations(pos_order, k):
+            for aa_combo in itertools.product(*(per_pos[i] for i in pos_combo)):
+                chosen = list(wt)
+                for i, a in zip(pos_combo, aa_combo, strict=True):
+                    chosen[i] = a
+                cs = "".join(chosen)
+                if cs == wt or cs in seen:
+                    continue
+                seen.add(cs)
+                out.append(cs)
+                if len(out) >= cap:
+                    return out
+    return out