shellde 0.2.0__py3-none-any.whl

shellde/colab.py

@@ -0,0 +1,257 @@
+"""Notebook / Colab front end: one function call == one directed-evolution round.
+
+Root design: the notebook stays a tiny, stable bootstrap ::
+
+    %pip install -q -U shellde
+    from shellde.colab import plate
+    df = plate(wt=\"\"\"...WT sequence...\"\"\", uniprot=\"Q50L36\", auto_msa=True)
+
+All behaviour lives HERE, versioned with the package, so improvements arrive via ``pip install -U``
+and the notebook file itself never needs re-uploading. This is a thin wrapper over ``shellde round``
+(``cli.main``) -- the science (funclib seed, active learning, calibrated abstention) is unchanged.
+"""
+from __future__ import annotations
+
+import glob
+import os
+from pathlib import Path
+
+from shellde.cli import main as _main
+from shellde.hotspots import rank_hotspots
+
+
+_AA_LETTERS = set("ACDEFGHIKLMNPQRSTVWYXBZUO")
+
+
+def _clean_sequence(wt: str) -> str:
+    """Raw sequence or FASTA text -> bare uppercase amino acids only.
+
+    Drops '>' headers and keeps ONLY amino-acid letters, so stray spaces/newlines/digits/punctuation
+    (and non-latin placeholder text) are removed -- a too-short result then trips plate()'s guard.
+    """
+    body = "".join(line for line in wt.splitlines() if not line.strip().startswith(">"))
+    return "".join(c for c in body.upper() if c in _AA_LETTERS)
+
+
+def _download_links(paths: list[str]) -> None:
+    """Show CLICKABLE download links for the result files (a zip-all + one per file).
+
+    No forced download: the files are embedded as base64 data URIs, so the user clicks what they want.
+    Works in Colab and Jupyter (no server round-trip). No-op outside a notebook.
+    """
+    import base64
+    import html
+    import io
+    import zipfile
+
+    try:
+        from IPython.display import HTML, display
+    except Exception:  # noqa: BLE001 (not in a notebook -> nothing to render)
+        return
+    existing = [p for p in paths if os.path.exists(p)]
+    if not existing:
+        return
+    items = []
+    for p in existing:
+        data = Path(p).read_bytes()
+        b64 = base64.b64encode(data).decode()
+        name = html.escape(os.path.basename(p))
+        items.append(
+            f'<li><a download="{name}" href="data:text/csv;base64,{b64}">{name}</a>'
+            f' <span style="color:#888">({len(data) / 1024:.1f} KB)</span></li>'
+        )
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
+        for p in existing:
+            z.write(p, os.path.basename(p))
+    zb64 = base64.b64encode(buf.getvalue()).decode()
+    display(HTML(
+        '<div style="font-family:sans-serif;line-height:1.6">'
+        '<b>Download results</b> (click):<br>'
+        f'<a download="shellde_round_results.zip" href="data:application/zip;base64,{zb64}">'
+        '&#128230; all files (.zip)</a>'
+        f'<ul style="margin:4px 0">{"".join(items)}</ul>'
+        '</div>'
+    ))
+
+
+def plate(
+    wt: str,
+    *,
+    uniprot: str | None = None,
+    positions: str | None = None,
+    pdb: str | None = None,
+    query: str | None = None,
+    holo: str | None = None,
+    scan: str | None = None,
+    top_k: int = 6,
+    measured: list[str] | str | None = None,
+    auto_msa: bool = True,
+    msa: str | None = None,
+    ddg: str | None = None,
+    if_logprobs: str | None = None,
+    tolerance: str | None = None,
+    wt_fitness: float | None = None,
+    plm: str = "none",
+    plm_rerank: int = 256,
+    exclude_catalytic: bool = False,
+    find_holo: bool = False,
+    model_class: str | None = None,
+    beta: float | None = None,
+    ligand_resnames: str | None = None,
+    min_coverage: float | None = None,
+    timeout: float | None = None,
+    plate_size: int = 95,
+    max_mut: int = 4,
+    outdir: str = "outputs/round",
+    download: bool = True,
+    wt_path: str = "WT.fasta",
+):
+    """Run one round and return the proposed plate as a DataFrame (or None if it abstained).
+
+    ``wt`` is the wild-type sequence (raw or FASTA text). Pass whatever design-region inputs you have
+    -- ``positions=`` (e.g. ``\"183,184,227,228\"``), ``holo=`` (local ligand-bound PDB), ``pdb=`` (RCSB
+    id / local cocrystal), ``uniprot=`` (accession or ``\"auto\"`` = exact sequence match), ``query=``
+    (UniProt text search) -- and the region is picked automatically, most-specific first (positions >
+    holo > pdb+uniprot combine > pdb > uniprot/query > sequence match). Omit ``measured`` for the R0 seed plate; pass measured
+    ``variant,fitness`` CSV path(s) to advance a round. ``auto_msa=True`` builds the R0 seed MSA from
+    the WT via ColabFold MMseqs2 so R0 runs with no uploaded file (natural proteins; GPU-free).
+    Prints the running version + command and shows clickable download links for the result files
+    (a zip of all, plus one per file); nothing is force-downloaded.
+    """
+    try:
+        import importlib.metadata as _md
+        print("shellde", _md.version("shellde"))
+    except Exception:  # noqa: BLE001
+        pass
+
+    seq = _clean_sequence(wt)
+    if len(seq) < 20:
+        raise ValueError(
+            f"wt cleaned to only {len(seq)} amino acids ({seq!r}) -- paste the FULL wild-type protein "
+            "sequence (you likely left the placeholder text, or pasted a fragment). Raw sequence or "
+            "FASTA both work."
+        )
+    Path(wt_path).write_text(">WT\n" + seq + "\n")
+    print(f"WT: {len(seq)} residues -> {wt_path}")
+
+    argv = ["round", wt_path, "--plate", str(plate_size), "--max-mut", str(max_mut), "--outdir", outdir]
+
+    if scan and not (uniprot or positions or pdb or query or holo):
+        positions = hotspots(scan, top_k=top_k)  # niche C: single-mutant scan -> hotspot positions
+    # Auto region dispatch: use whatever inputs are given, most-specific first (positions > holo >
+    # pdb+uniprot combine > pdb > uniprot/query > sequence match). No forced single choice; the chosen
+    # region and any also-provided-but-unused inputs are printed so nothing is silent.
+    given = {k: v for k, v in
+             {"positions": positions, "holo": holo, "pdb": pdb, "uniprot": uniprot, "query": query}.items() if v}
+    if positions:
+        use = {"positions": positions}
+    elif holo:
+        use = {"holo": holo}
+    elif pdb and (uniprot or query):
+        use = {"pdb": pdb, "uniprot": uniprot} if uniprot else {"pdb": pdb, "query": query}
+    elif pdb:
+        use = {"pdb": pdb}
+    elif uniprot or query:
+        use = {"uniprot": uniprot} if uniprot else {"query": query}
+    else:
+        use = {"uniprot": "auto"}  # match the accession from the sequence
+    if not given:
+        print("auto region: no inputs given -> --uniprot auto (exact sequence match)")
+    else:
+        ignored = sorted(set(given) - set(use))
+        how = "combining" if len(use) > 1 else "using"
+        msg = f"auto region: {how} {sorted(use)}"
+        if ignored:
+            msg += f" (also provided, not needed here: {ignored})"
+        print(msg)
+    for key, val in use.items():
+        argv += [f"--{key}", val]
+    if exclude_catalytic:
+        argv += ["--exclude-catalytic"]
+    if find_holo:
+        argv += ["--find-holo"]
+
+    if msa:
+        argv += ["--msa", msa]
+    elif auto_msa:
+        argv += ["--auto-msa"]
+    if ddg:
+        argv += ["--ddg", ddg]
+    if if_logprobs:
+        argv += ["--if-logprobs", if_logprobs]
+    if tolerance:
+        argv += ["--tolerance", tolerance]
+    if measured:
+        meas = [measured] if isinstance(measured, str) else list(measured)
+        argv += ["--measured", *meas]
+    if wt_fitness is not None:
+        argv += ["--wt-fitness", str(wt_fitness)]
+    if plm != "none":
+        argv += ["--plm", plm, "--plm-rerank", str(plm_rerank)]
+    if model_class:
+        argv += ["--model-class", model_class]
+    if beta is not None:
+        argv += ["--beta", str(beta)]
+    if ligand_resnames:
+        argv += ["--ligand-resnames", ligand_resnames]
+    if min_coverage is not None:
+        argv += ["--min-coverage", str(min_coverage)]
+    if timeout is not None:
+        argv += ["--timeout", str(timeout)]
+
+    print("running: shellde " + " ".join(argv))
+    print("-" * 70)
+    _main(argv)
+    print("-" * 70)
+
+    found = sorted(glob.glob(os.path.join(outdir, "round*_plate.csv")), key=os.path.getmtime)
+    if not found:
+        print("no plate written (abstained -- read the message above). Common fixes: give a design "
+              "region (uniprot=/positions=/pdb=); for the first plate keep auto_msa=True or pass "
+              "msa=/ddg=.")
+        return None
+    latest = found[-1]
+    prefix = os.path.basename(latest)[: -len("_plate.csv")]  # e.g. "round0"
+    saved = [os.path.join(outdir, prefix + s) for s in ("_plate.csv", "_all.csv", "_measured_template.csv")]
+    saved = [p for p in saved if os.path.exists(p)]
+    print(f"saved to {outdir}/:  " + ",  ".join(os.path.basename(p) for p in saved))
+    print("  _plate = variants to MAKE/measure | _all = full scored set | _measured_template = fill "
+          "fitness -> next round")
+    try:
+        import pandas as pd
+        df = pd.read_csv(latest)
+    except Exception:  # noqa: BLE001 (pandas missing / unreadable -> still report the path)
+        print("plate written to", latest)
+        return None
+    print(f"plate: {len(df)} variants (top of {prefix})")
+    if download:
+        _download_links(saved)
+    return df
+
+
+def hotspots(
+    scan: str,
+    *,
+    top_k: int = 6,
+    metric: str = "max",
+    baseline: float = 0.0,
+    min_effect: float | None = None,
+    mutation_col: str | None = None,
+    score_col: str | None = None,
+) -> str:
+    """Rank hotspot positions from a whole-protein single-mutant scan CSV (the niche-C bridge).
+
+    ``scan`` is a CSV path with a mutation column (e.g. ``V100F``) and a score column (measured
+    fitness/activity, or a predicted ``y_pred``/``score``). Returns a comma-separated positions
+    string ready for ``plate(positions=...)`` / ``round --positions``. Prints the ranked positions.
+    """
+    from shellde.cli import _read_singles
+
+    singles, mcol, scol = _read_singles(scan, mutation_col, score_col)
+    hs = rank_hotspots(singles, top_k=top_k, metric=metric, baseline=baseline, min_effect=min_effect)
+    positions = ",".join(str(h.position) for h in sorted(hs, key=lambda x: x.position))
+    print(f"hotspots (mutation='{mcol}', score='{scol}', metric={metric}) -> {positions}")
+    for h in hs:
+        print(f"  {h.wt}{h.position}  best {h.best_sub} ({h.best_effect:.4g})  n_beneficial={h.n_beneficial}")
+    return positions

shellde/conformal.py

@@ -0,0 +1,160 @@
+"""Distribution-free conformal prediction intervals for the surrogate.
+
+The calibrated ``Prediction.std`` is a *parametric* uncertainty: variance scaling
+fits one scalar by k-fold, so it is Gaussian-flavoured and, by construction, OFF at
+very low N (``surrogate._calibrate`` returns 1.0 when n < 4). Split conformal sidesteps
+both problems: it gives finite-sample *marginal* coverage with no Gaussian assumption
+and is valid at small N (the interval just gets wider). It needs only the surrogate's
+point prediction (``.mean``), never its ``.std``, so it is robust precisely where the
+variance-scaling path is weakest.
+
+This is split (a.k.a. inductive) conformal: the surrogate is fitted on a TRAIN split,
+the nonconformity scores ``|y - mean|`` are collected on a disjoint CALIBRATION split,
+and coverage holds on any further exchangeable point. The wrapper never fits the
+surrogate; it consumes an already-fitted one, keeping the Prediction(mean, std) contract
+and the surrogate machinery untouched.
+"""
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+
+import numpy as np
+from sklearn.model_selection import KFold
+
+from shellde.protocols import Surrogate
+
+ALPHA_90 = 0.1  # default miscoverage level: a 90% prediction interval
+_FLOOR = 1e-9  # sigma floor for the normalized score (matches surrogate._FLOOR)
+
+
+def conformal_quantile(scores: np.ndarray, alpha: float = ALPHA_90) -> float:
+    """Distribution-free conformal quantile: the ceil((n+1)(1-alpha))-th smallest |score|.
+
+    The index is clipped to n, so at low N (when (n+1)(1-alpha) > n) it saturates at the
+    largest score (the conformal "infinite quantile"), keeping the half-width finite and
+    marginally valid. Shared by ``SplitConformal._quantile`` and the cross-conformal
+    half-width surfaced in ``advisor.data_readiness``.
+    """
+    if not 0.0 < alpha < 1.0:
+        raise ValueError(f"alpha must be in (0, 1), got {alpha}")
+    s = np.sort(np.abs(np.asarray(scores, dtype=float).ravel()))
+    n = s.shape[0]
+    if n < 1:
+        raise ValueError("need at least 1 nonconformity score")
+    k = math.ceil((n + 1) * (1.0 - alpha))
+    return float(s[min(k, n) - 1])
+
+
+class SplitConformal:
+    """Split-conformal interval wrapper around an already-fitted surrogate.
+
+    Usage (disjoint splits): fit ``surrogate`` on a TRAIN split, then ``fit`` this on a
+    disjoint CALIBRATION split, then ``interval(x)`` on new points. Coverage at level
+    ``1 - alpha`` is finite-sample valid (marginal) for exchangeable data.
+    """
+
+    def __init__(self, surrogate: Surrogate) -> None:
+        self.surrogate = surrogate
+        self._scores: np.ndarray | None = None
+
+    def fit(self, x_cal: np.ndarray, y_cal: np.ndarray) -> "SplitConformal":
+        """Store sorted absolute nonconformity scores ``|y - surrogate.mean|`` on the cal split."""
+        x_cal = np.asarray(x_cal, dtype=float)
+        y_cal = np.asarray(y_cal, dtype=float).ravel()
+        if x_cal.ndim != 2 or x_cal.shape[0] != y_cal.shape[0]:
+            raise ValueError(f"x_cal {x_cal.shape} incompatible with y_cal {y_cal.shape}")
+        if x_cal.shape[0] < 1:
+            raise ValueError("need at least 1 calibration sample")
+        mean = np.asarray(self.surrogate.predict(x_cal).mean, dtype=float).ravel()
+        self._scores = np.sort(np.abs(y_cal - mean))
+        return self
+
+    def _quantile(self, alpha: float) -> float:
+        if self._scores is None:
+            raise RuntimeError("SplitConformal not fitted; call fit() first")
+        return conformal_quantile(self._scores, alpha)
+
+    def interval(self, x: np.ndarray, alpha: float = ALPHA_90) -> tuple[np.ndarray, np.ndarray]:
+        """Return ``(lo, hi)`` = mean +/- conformal quantile, vectorized over ``x``."""
+        x = np.asarray(x, dtype=float)
+        q = self._quantile(alpha)
+        mean = np.asarray(self.surrogate.predict(x).mean, dtype=float).ravel()
+        return mean - q, mean + q
+
+
+def empirical_coverage(
+    lo: np.ndarray, hi: np.ndarray, y: np.ndarray
+) -> float:
+    """Fraction of held-out targets that fall inside ``[lo, hi]`` (the realised coverage)."""
+    lo = np.asarray(lo, dtype=float).ravel()
+    hi = np.asarray(hi, dtype=float).ravel()
+    y = np.asarray(y, dtype=float).ravel()
+    if not (lo.shape == hi.shape == y.shape):
+        raise ValueError(f"shape mismatch: lo {lo.shape}, hi {hi.shape}, y {y.shape}")
+    return float(np.mean((y >= lo) & (y <= hi)))
+
+
+def conformal_interval(
+    surrogate: Surrogate,
+    x_cal: np.ndarray,
+    y_cal: np.ndarray,
+    x_query: np.ndarray,
+    *,
+    alpha: float = ALPHA_90,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Standalone helper: calibrate on (x_cal, y_cal) and return the interval at x_query.
+
+    Additive convenience over an already-fitted ``surrogate``; changes no existing
+    signature. The surrogate must have been fitted on data DISJOINT from ``x_cal``.
+    """
+    return SplitConformal(surrogate).fit(x_cal, y_cal).interval(x_query, alpha=alpha)
+
+
+def cross_conformal_normalized_q(
+    x: np.ndarray,
+    y: np.ndarray,
+    make_surrogate: Callable[[], Surrogate],
+    *,
+    k: int = 5,
+    alpha: float = ALPHA_90,
+    seed: int = 0,
+) -> float:
+    """Normalized cross-conformal quantile (heteroscedastic-preserving).
+
+    Runs seeded k-fold; on each out-of-fold point collects the NORMALIZED nonconformity
+    score ``|y - mean_oof| / max(sigma_oof, floor)`` using a surrogate refit on the fold's
+    train split (its ``Prediction.mean`` and ``Prediction.std``). The pooled scores'
+    ``conformal_quantile`` is returned as ``q_norm``: a per-candidate interval half-width
+    is then ``q_norm * sigma(x)``, which preserves the surrogate's heteroscedastic sigma
+    instead of collapsing it to a constant.
+
+    Coverage is MARGINAL and distribution-free only for points EXCHANGEABLE with the
+    measured set. It is NOT valid for out-of-distribution candidates (e.g. higher-order
+    mutation combinations when the measured set is single mutants); the caller must label
+    the scope accordingly. This is not redundant with ``advisor.data_readiness`` (that one
+    is a mean-only additive one-hot Ridge half-width; this one needs the chosen surrogate's
+    calibrated sigma).
+    """
+    x = np.asarray(x, dtype=float)
+    y = np.asarray(y, dtype=float).ravel()
+    n = y.shape[0]
+    if x.ndim != 2 or x.shape[0] != n:
+        raise ValueError(f"x {x.shape} incompatible with y {y.shape}")
+    n_splits = min(k, n)
+    if n_splits < 2:
+        raise ValueError("need at least 2 samples for cross-conformal")
+    kf = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
+    scores: list[float] = []
+    for tr, te in kf.split(x):
+        if len(tr) < 2 or float(np.std(y[tr])) < _FLOOR:
+            # Degenerate fold: predict the train mean, sigma = residual floor.
+            mean_oof = np.full(len(te), float(np.mean(y[tr])) if len(tr) else 0.0)
+            sigma_oof = np.full(len(te), _FLOOR)
+        else:
+            pred = make_surrogate().fit(x[tr], y[tr]).predict(x[te])
+            mean_oof = np.asarray(pred.mean, dtype=float).ravel()
+            sigma_oof = np.asarray(pred.std, dtype=float).ravel()
+        denom = np.maximum(sigma_oof, _FLOOR)
+        scores.extend((np.abs(y[te] - mean_oof) / denom).tolist())
+    return conformal_quantile(np.asarray(scores, dtype=float), alpha)

shellde/consensus.py

@@ -0,0 +1,52 @@
+"""Merge multiple active-site evidence sources into one per-position confidence map.
+
+The evidence ladder picks ONE source. When several independent lines are available (UniProt curated
+catalytic sites, a PDB cocrystal contact shell, a Foldseek holo-homolog shell, a whole-protein scan),
+overlaying them says WHERE they AGREE - and agreement across independent evidence types is the honest
+confidence signal (it does NOT manufacture activity prediction; it says which positions to trust/measure
+first). Pure, no I/O.
+"""
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+# Evidence types that are experimental/curated (a single one already justifies designing there);
+# everything else is a weaker/heuristic prior that wants corroboration.
+STRONG_EVIDENCE = frozenset({"uniprot_catalytic", "pdb_cocrystal_shell", "holo_homolog_shell"})
+
+
+def build_consensus(sources: dict[str, Iterable[int]], *, strong: frozenset[str] = STRONG_EVIDENCE) -> list[dict]:
+    """Overlay named position sources -> sorted per-position ``{position, evidence, tier}``.
+
+    ``sources`` maps an evidence-type name (e.g. ``"pdb_cocrystal_shell"``) to its positions. Tier:
+    ``high`` = supported by >=2 independent evidence types; ``medium`` = exactly one STRONG source;
+    ``low`` = exactly one weak/heuristic source. Deterministic (evidence names sorted).
+    """
+    by_pos: dict[int, set[str]] = {}
+    for name, positions in sources.items():
+        for p in positions:
+            by_pos.setdefault(int(p), set()).add(name)
+    out: list[dict] = []
+    for pos in sorted(by_pos):
+        ev = sorted(by_pos[pos])
+        if len(ev) >= 2:
+            tier = "high"
+        elif any(e in strong for e in ev):
+            tier = "medium"
+        else:
+            tier = "low"
+        out.append({"position": pos, "evidence": ev, "tier": tier})
+    return out
+
+
+def tier_positions(consensus: list[dict], tier: str) -> list[int]:
+    """Positions at a given confidence tier (``high``/``medium``/``low``)."""
+    return [c["position"] for c in consensus if c["tier"] == tier]
+
+
+def consensus_summary(consensus: list[dict]) -> str:
+    """One-line count summary, e.g. ``high=7 medium=78 low=0``."""
+    counts = {"high": 0, "medium": 0, "low": 0}
+    for c in consensus:
+        counts[c["tier"]] = counts.get(c["tier"], 0) + 1
+    return " ".join(f"{t}={counts[t]}" for t in ("high", "medium", "low"))

shellde/design_space.py

@@ -0,0 +1,141 @@
+"""Frozen design space + variant<->assignment normalisation.
+
+A ``DesignSpace`` is a sorted tuple of ABSOLUTE 1-based residue positions plus the
+reference (WT) residue at each. Freezing it is what guarantees a feature column
+layout that is stable across active-learning rounds selecting DIFFERENT candidate
+pools. Variants are accepted as fixed-site combo strings ("FRMNY") or absolute
+mutation tokens ("A42V:L88M") and normalised to a full per-position assignment.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+AA_ALPHABET = "ACDEFGHIKLMNPQRSTVWY"
+
+Mutation = tuple[int, str]  # (1-based absolute position, target AA)
+
+
+@dataclass(frozen=True)
+class DesignSpace:
+    """Frozen set of designed positions + reference residues + alphabet."""
+
+    positions: tuple[int, ...]
+    reference: Mapping[int, str]
+    alphabet: str = AA_ALPHABET
+
+    def __post_init__(self) -> None:
+        if list(self.positions) != sorted(set(self.positions)):
+            raise ValueError("DesignSpace.positions must be sorted and unique")
+        missing = [p for p in self.positions if p not in self.reference]
+        if missing:
+            raise ValueError(f"reference missing residues for positions {missing}")
+        bad = [p for p in self.positions if self.reference[p] not in self.alphabet]
+        if bad:
+            raise ValueError(f"reference residue outside alphabet at positions {bad}")
+
+    @property
+    def n_positions(self) -> int:
+        return len(self.positions)
+
+    @property
+    def q(self) -> int:
+        return len(self.alphabet)
+
+    def position_index(self, pos: int) -> int:
+        try:
+            return self.positions.index(pos)
+        except ValueError as exc:
+            raise ValueError(f"position {pos} not in design space") from exc
+
+    def wt(self) -> str:
+        """The reference combo string over the designed positions."""
+        return "".join(self.reference[p] for p in self.positions)
+
+
+def combo_to_assignment(combo: str, space: DesignSpace) -> dict[int, str]:
+    """Map a fixed-site combo string (one AA per designed position) to an assignment."""
+    if len(combo) != space.n_positions:
+        raise ValueError(f"combo {combo!r} length {len(combo)} != n_positions {space.n_positions}")
+    invalid = sorted({c for c in combo if c not in space.alphabet})
+    if invalid:
+        raise ValueError(f"combo {combo!r} has residues outside alphabet: {invalid}")
+    return {pos: combo[i] for i, pos in enumerate(space.positions)}
+
+
+def tokens_to_assignment(tokens: str, space: DesignSpace) -> dict[int, str]:
+    """Map absolute mutation tokens ('A42V:L88M' / 'A42V,L88M') to a full assignment.
+
+    Unspecified designed positions are filled from the reference (WT) residue. The
+    cited wild-type residue is validated against the reference when given.
+    """
+    assignment = dict(space.reference)
+    raw = tokens.replace(",", ":").replace(" ", ":")
+    for tok in (t for t in raw.split(":") if t):
+        wt, pos, mt = tok[0], tok[1:-1], tok[-1]
+        if not wt.isalpha():
+            raise ValueError(f"mutation token {tok!r} must be [WT][pos][MT], e.g. A42V")
+        try:
+            ipos = int(pos)
+        except ValueError as exc:
+            raise ValueError(f"bad mutation token {tok!r}") from exc
+        if ipos not in space.reference:
+            raise ValueError(f"token {tok!r} position {ipos} not in design space")
+        if space.reference[ipos] != wt:
+            raise ValueError(f"token {tok!r} WT {wt} != reference {space.reference[ipos]} at {ipos}")
+        if mt not in space.alphabet:
+            raise ValueError(f"token {tok!r} target residue {mt!r} outside alphabet")
+        assignment[ipos] = mt
+    return assignment
+
+
+def to_assignment(variant: str, space: DesignSpace) -> dict[int, str]:
+    """Normalise a variant string to a full per-position assignment over the space.
+
+    A pure-alphabetic string of length n_positions is a fixed-site combo; anything
+    containing a digit is treated as mutation tokens.
+    """
+    if not variant:
+        raise ValueError("empty variant string")
+    if variant.isalpha() and len(variant) == space.n_positions:
+        return combo_to_assignment(variant, space)
+    return tokens_to_assignment(variant, space)
+
+
+def mutations_of(variant: str, space: DesignSpace) -> tuple[str, ...]:
+    """Mutation tokens (e.g. 'V39A') where the variant differs from the reference."""
+    a = to_assignment(variant, space)
+    return tuple(
+        f"{space.reference[p]}{p}{a[p]}" for p in space.positions if a[p] != space.reference[p]
+    )
+
+
+def combo_indices(variants: Sequence[str], space: DesignSpace) -> np.ndarray | None:
+    """Vectorized combo-string -> (N, n_positions) alphabet-index array, or None to fall back.
+
+    Returns the per-position alphabet indices for ``variants`` IFF every variant is a valid
+    fixed-site combo over this space (``len == n_positions`` and every residue in ``alphabet``)
+    -- exactly the case where :func:`to_assignment` routes to :func:`combo_to_assignment`. Returns
+    ``None`` (so the caller uses the per-variant slow path) for any token string, wrong length, or
+    out-of-alphabet residue, preserving identical behaviour AND error semantics. Pure speed path.
+    """
+    n, length = len(variants), space.n_positions
+    if n == 0:
+        return np.zeros((0, length), dtype=np.intp)
+    if any(len(v) != length for v in variants):
+        return None
+    try:
+        buf = np.frombuffer("".join(variants).encode("latin-1"), dtype=np.uint8)
+    except (UnicodeEncodeError, ValueError):
+        return None
+    if buf.size != n * length:
+        return None
+    lut = np.full(256, -1, dtype=np.intp)
+    for i, aa in enumerate(space.alphabet):
+        lut[ord(aa)] = i
+    idx = lut[buf].reshape(n, length)
+    if (idx < 0).any():  # a digit (token) or out-of-alphabet residue -> slow path handles/raises
+        return None
+    return idx

shellde/embeddings/__init__.py

@@ -0,0 +1,6 @@
+"""Embedding providers (PLM) + on-disk cache."""
+from __future__ import annotations
+
+from shellde.embeddings.provider import EmbeddingCache, FunctionProvider
+
+__all__ = ["EmbeddingCache", "FunctionProvider"]