shellde 0.2.0__py3-none-any.whl

shellde/embeddings/esm2.py

@@ -0,0 +1,76 @@
+"""ESM-2 (Meta) embedding provider via transformers (import-guarded, cached).
+
+Provided as an alternative/baseline PLM front end. Mean-pools the last hidden state
+over residues (excluding BOS/EOS). transformers + torch are optional; the import is
+lazy and raises a clear error when absent.
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+from shellde.embeddings.provider import EmbeddingCache, _key
+
+_WIDTH = {
+    "facebook/esm2_t6_8M_UR50D": 320,
+    "facebook/esm2_t12_35M_UR50D": 480,
+    "facebook/esm2_t30_150M_UR50D": 640,
+    "facebook/esm2_t33_650M_UR50D": 1280,
+}
+
+
+class ESM2Provider:
+    """Mean-pooled ESM-2 embeddings; satisfies protocols.EmbeddingProvider."""
+
+    def __init__(
+        self, model: str = "facebook/esm2_t30_150M_UR50D", *, dim: int | None = None,
+        device: str | None = None, cache: EmbeddingCache | None = None,
+    ) -> None:
+        self.model = model
+        self.dim = dim if dim is not None else _WIDTH.get(model, 640)
+        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:
+        try:
+            import torch
+            from transformers import AutoModel, AutoTokenizer  # type: ignore[import-not-found, import-untyped]
+        except ImportError as exc:  # pragma: no cover - exercised only without the SDK
+            raise ImportError(
+                "ESM2Provider requires 'transformers' and 'torch'. Install them to use ESM-2."
+            ) from exc
+        self._torch = torch
+        self._tok = AutoTokenizer.from_pretrained(self.model)
+        self._mdl = AutoModel.from_pretrained(self.model).eval()
+
+    def embed(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:
+            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 j in range(0, len(misses), 16):
+                    batch = misses[j : j + 16]
+                    enc = self._tok([s for _, s in batch], return_tensors="pt", padding=True)
+                    rep = self._mdl(**enc).last_hidden_state
+                    mask = enc["attention_mask"].unsqueeze(-1).float()
+                    mask[:, 0, :] = 0.0
+                    for b in range(len(batch)):
+                        n = int(enc["attention_mask"][b].sum())
+                        mask[b, n - 1, :] = 0.0
+                    pooled = (rep * mask).sum(1) / mask.sum(1).clamp(min=1.0)
+                    for b, (i, _) in enumerate(batch):
+                        self._cache.put(keys[i], pooled[b].cpu().numpy())
+        out: list[np.ndarray] = []
+        for k in keys:
+            v = self._cache.get(k)
+            assert v is not None
+            out.append(v)
+        return np.stack(out, axis=0)

shellde/embeddings/esmc.py

@@ -0,0 +1,94 @@
+"""ESM-C (EvolutionaryScale) embedding provider (import-guarded, cached).
+
+ESM-C was measured substantially more activity-predictive than ESM-2 / one-hot in
+the single-mutant regime, so it is the default PLM front end. The EvolutionaryScale
+``esm`` SDK is optional and distinct from fair-esm; the import is lazy and raises a
+clear error when absent, so the rest of the package works without it.
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+from shellde.embeddings.provider import EmbeddingCache, _key
+
+_WIDTH = {"esmc_300m": 960, "esmc_600m": 1152}
+
+
+class ESMCProvider:
+    """Mean-pooled per-residue ESM-C embeddings; satisfies protocols.EmbeddingProvider."""
+
+    def __init__(
+        self, model_name: str = "esmc_300m", *, device: str | None = None,
+        cache: EmbeddingCache | None = None,
+    ) -> None:
+        self.model_name = model_name
+        self.dim = _WIDTH.get(model_name, 960)
+        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:
+        try:
+            import torch
+            from esm.models.esmc import ESMC  # type: ignore[import-not-found, import-untyped]
+        except ImportError as exc:  # pragma: no cover - exercised only without the SDK
+            raise ImportError(
+                "ESMCProvider requires the EvolutionaryScale 'esm' SDK "
+                "(esm.models.esmc.ESMC), which is separate from fair-esm. Install 'esm'."
+            ) 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 embed(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:
+            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, return_embeddings=True))
+                    emb = res.embeddings[0].float().cpu().numpy()  # (L+2, D): BOS..residues..EOS
+                    self._cache.put(keys[i], emb[1:-1].mean(axis=0))
+        out: list[np.ndarray] = []
+        for k in keys:
+            v = self._cache.get(k)
+            assert v is not None
+            out.append(v)
+        return np.stack(out, axis=0)
+
+    def embed_residues(self, sequences: Sequence[str]) -> list[np.ndarray]:
+        """Per-residue embeddings (one (L, D) float array per sequence; BOS/EOS stripped).
+
+        Cached under a ``:res`` namespace key so they never collide with the mean-pool
+        cache. Returns a list (not a stacked array) because residue arrays differ in L.
+        """
+        keys = [_key(self.model_name + ":res", s) for s in sequences]
+        misses = [(i, s) for i, s in enumerate(sequences) if self._cache.get(keys[i]) is None]
+        if misses:
+            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, return_embeddings=True))
+                    emb = res.embeddings[0].float().cpu().numpy()  # (L+2, D): BOS..residues..EOS
+                    self._cache.put(keys[i], emb[1:-1])
+        out: list[np.ndarray] = []
+        for k in keys:
+            v = self._cache.get(k)
+            assert v is not None
+            out.append(v)
+        return out

shellde/embeddings/provider.py

@@ -0,0 +1,105 @@
+"""Embedding cache + a test/compose-friendly FunctionProvider.
+
+Providers turn sequences into mean-pooled embeddings on demand and cache them so
+repeated variants (across active-learning rounds) are embedded once. The cache key
+hashes (model_id, sequence) so keys stay short and persistable. FunctionProvider
+wraps any ``str -> vector`` callable, which makes the embedding plumbing testable
+without a heavy PLM SDK and lets callers compose custom features.
+"""
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Callable, Sequence
+from pathlib import Path
+
+import numpy as np
+
+
+def _key(model_id: str, seq: str) -> str:
+    return f"{model_id}:{hashlib.sha1(seq.encode()).hexdigest()}"  # noqa: S324 (cache key, not security)
+
+
+class EmbeddingCache:
+    """In-memory embedding cache with optional ``.npz`` persistence."""
+
+    def __init__(self, path: str | Path | None = None) -> None:
+        self.path = Path(path) if path else None
+        self._store: dict[str, np.ndarray] = {}
+        if self.path is not None and self.path.exists():
+            self.load()
+
+    def get(self, key: str) -> np.ndarray | None:
+        return self._store.get(key)
+
+    def put(self, key: str, vec: np.ndarray) -> None:
+        self._store[key] = np.asarray(vec, dtype=float)
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def save(self) -> None:
+        if self.path is not None:
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+            np.savez(self.path, **self._store)  # type: ignore[arg-type]
+
+    def load(self) -> None:
+        if self.path is not None and self.path.exists():
+            data = np.load(self.path)
+            self._store = {k: data[k] for k in data.files}
+
+
+class FunctionProvider:
+    """An EmbeddingProvider backed by a plain ``str -> 1D array`` function (cached).
+
+    Satisfies ``protocols.EmbeddingProvider``. ``calls`` counts cache misses so tests
+    can assert that repeated sequences are embedded only once.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        fn: Callable[[str], np.ndarray],
+        *,
+        model_id: str = "fn",
+        cache: EmbeddingCache | None = None,
+        residue_fn: Callable[[str], np.ndarray] | None = None,
+    ) -> None:
+        self.dim = dim
+        self._fn = fn
+        self.model_id = model_id
+        self._cache = cache if cache is not None else EmbeddingCache()
+        self.calls = 0
+        self._residue_fn = residue_fn
+
+    def embed(self, sequences: Sequence[str]) -> np.ndarray:
+        out: list[np.ndarray] = []
+        for s in sequences:
+            key = _key(self.model_id, s)
+            vec = self._cache.get(key)
+            if vec is None:
+                self.calls += 1
+                vec = np.asarray(self._fn(s), dtype=float)
+                if vec.shape != (self.dim,):
+                    raise ValueError(f"fn returned {vec.shape}, expected ({self.dim},)")
+                self._cache.put(key, vec)
+            out.append(vec)
+        return np.stack(out, axis=0)
+
+    def embed_residues(self, sequences: Sequence[str]) -> list[np.ndarray]:
+        """Per-residue embeddings via ``residue_fn`` (one (L, D) array per sequence).
+
+        Raises if no ``residue_fn`` was supplied (no silent fallback to mean-pool).
+        """
+        if self._residue_fn is None:
+            raise ValueError(
+                "FunctionProvider has no residue_fn; embed_residues (site pooling) unsupported"
+            )
+        out: list[np.ndarray] = []
+        for s in sequences:
+            arr = np.asarray(self._residue_fn(s), dtype=float)
+            if arr.ndim != 2 or arr.shape != (len(s), self.dim):
+                raise ValueError(
+                    f"residue_fn returned {arr.shape}, expected ({len(s)}, {self.dim})"
+                )
+            out.append(arr)
+        return out

shellde/features/__init__.py

@@ -0,0 +1,22 @@
+"""Feature blocks + matrix composition."""
+from __future__ import annotations
+
+from shellde.features.base import FeatureBlock
+from shellde.features.defaults import default_blocks
+from shellde.features.embedding import PLMEmbeddingBlock
+from shellde.features.matrix import FeatureMatrix
+from shellde.features.inverse_folding import InverseFoldingBlock
+from shellde.features.naturalness import NaturalnessBlock
+from shellde.features.onehot import OneHotBlock
+from shellde.features.pairwise import GatedPairwiseBlock
+
+__all__ = [
+    "FeatureBlock",
+    "FeatureMatrix",
+    "GatedPairwiseBlock",
+    "InverseFoldingBlock",
+    "NaturalnessBlock",
+    "OneHotBlock",
+    "PLMEmbeddingBlock",
+    "default_blocks",
+]

shellde/features/base.py

@@ -0,0 +1,49 @@
+"""FeatureBlock: one fixed-width, column-stable feature group.
+
+A block carries NO gate of its own. Whether an optional block is included is decided by the
+caller, per flag and per command (see `shellde.gating` and the CLI command functions).
+
+A block declares a fixed ``dim`` and whether it is ``present``. When absent (no
+structure, no embeddings, ...) it contributes a zero-weight column group of its
+declared width, so the matrix width is identical regardless of data availability
+and the surrogate auto-weights the unused columns down. ``encode`` is shape-checked
+so a buggy block fails loud instead of silently corrupting the matrix.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+
+
+class FeatureBlock(ABC):
+    """One feature block contributing a fixed-width column group."""
+
+    name: str
+
+    @property
+    @abstractmethod
+    def dim(self) -> int: ...
+
+    @property
+    def present(self) -> bool:
+        return True
+
+    @abstractmethod
+    def _encode_present(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray: ...
+
+    def encode(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray:
+        n = len(assignments)
+        if not self.present:
+            return np.zeros((n, self.dim), dtype=float)
+        out = self._encode_present(assignments, space)
+        if out.shape != (n, self.dim):
+            raise ValueError(f"block {self.name} produced {out.shape}, expected {(n, self.dim)}")
+        return out

shellde/features/defaults.py

@@ -0,0 +1,11 @@
+"""THE single source of truth for the shipped default feature blocks."""
+from __future__ import annotations
+
+from shellde.design_space import DesignSpace
+from shellde.features.base import FeatureBlock
+from shellde.features.onehot import OneHotBlock
+
+
+def default_blocks(space: DesignSpace) -> list[FeatureBlock]:
+    """THE shipped additive default block set, consumed by both `cli.recommend` and the bench `ours` method so the head-to-head always measures what ships. `GatedPairwiseBlock` and `InverseFoldingBlock` are OPT-INs the CALLER appends on top, never part of the default. Do not read "opt-in" as "data-gated": pairwise is CV-gated only under `--auto-gate-pairwise` (`--contacts-pdb` alone opens it unconditionally), and IF is CV-gated on `recommend` only (`cmd_campaign`, and `cmd_round` through it, appends it with no gate). `--plm` / `--naturalness` features do not pass through here at all; `plm_rerank` adds them, ungated, on any command that offers the flags."""
+    return [OneHotBlock(space)]

shellde/features/embedding.py

@@ -0,0 +1,80 @@
+"""PLMEmbeddingBlock: on-demand PLM embedding of design variants (single-mutant front end).
+
+Splices each variant's design-position residues into a full WT background sequence,
+then asks an EmbeddingProvider for the mean-pooled embedding. This is the PLM-native
+path for the regime where PLM implicit epistasis wins (single-mutant / long context);
+on dense combinatorial libraries embeddings dilute, so the surrogate auto-weights it
+down (it earns its place on the front end, not 200k-candidate scoring).
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.features.base import FeatureBlock
+from shellde.protocols import EmbeddingProvider
+
+
+class PLMEmbeddingBlock(FeatureBlock):
+    """Per-variant PLM embedding spliced into a full WT background (optional block)."""
+
+    def __init__(
+        self,
+        space: DesignSpace,
+        provider: EmbeddingProvider | None,
+        wt_sequence: str | None,
+        *,
+        name: str = "plm",
+        pooling: str = "mean",
+    ) -> None:
+        if pooling not in {"mean", "site"}:
+            raise ValueError(f"pooling must be 'mean' or 'site', got {pooling!r}")
+        self.name = name
+        self._provider = provider
+        self._wt = wt_sequence
+        self._pooling = pooling
+        if provider is not None and wt_sequence is not None:
+            for p in space.positions:
+                if p < 1 or p > len(wt_sequence):
+                    raise ValueError(f"position {p} outside wt_sequence length {len(wt_sequence)}")
+                if wt_sequence[p - 1] != space.reference[p]:
+                    raise ValueError(
+                        f"wt_sequence[{p - 1}]={wt_sequence[p - 1]!r} != reference "
+                        f"{space.reference[p]!r} at position {p}"
+                    )
+        self._dim = provider.dim if provider is not None else 0
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    @property
+    def present(self) -> bool:
+        return self._provider is not None and self._wt is not None
+
+    def _encode_present(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray:
+        assert self._provider is not None and self._wt is not None
+        seqs: list[str] = []
+        for a in assignments:
+            chars = list(self._wt)
+            for p in space.positions:
+                chars[p - 1] = a[p]
+            seqs.append("".join(chars))
+        if self._pooling == "mean":
+            return np.asarray(self._provider.embed(seqs), dtype=float)
+        embed_residues = getattr(self._provider, "embed_residues", None)
+        if embed_residues is None:
+            raise ValueError(
+                f"pooling='site' requires the provider to implement embed_residues, "
+                f"but {type(self._provider).__name__} does not"
+            )
+        per_residue = embed_residues(seqs)  # list of (L, D) arrays, one per variant
+        rows: list[np.ndarray] = []
+        for res in per_residue:
+            res = np.asarray(res, dtype=float)
+            rows.append(np.mean([res[p - 1] for p in space.positions], axis=0))
+        return np.stack(rows, axis=0)

shellde/features/inverse_folding.py

@@ -0,0 +1,66 @@
+"""InverseFoldingBlock: a structure inverse-folding score as one optional feature.
+
+An inverse-folding model (ProteinMPNN) gives per-position amino-acid log-probabilities
+GIVEN the backbone in a single pass, so a variant's score is a cheap lookup (sum over
+mutated positions of log P(mut) - log P(wt)), not a per-variant model call. This IF
+feature is kept HONEST: the evidence (V2) showed IF is
+protein dependent (helped TrpB, inert on GB1, slight negative on ParPgb), so it is an
+optional block. The block itself carries NO gate; the CALLER decides. `cmd_recommend`
+(`src/shellde/cli.py`) runs it through the additive-null CV gate (`gating.should_include_block`)
+by default, so on `recommend` it is opened only when this protein's held-out data supports it.
+`cmd_campaign` -- and `cmd_round`, which calls it -- appends this block UNCONDITIONALLY whenever
+`--if-logprobs` is supplied, with no gate.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.features.base import FeatureBlock
+
+
+class InverseFoldingBlock(FeatureBlock):
+    """Single-column inverse-folding log-odds-vs-WT score (optional, structure-derived)."""
+
+    name = "inverse_folding"
+
+    def __init__(
+        self,
+        space: DesignSpace,
+        logprobs: Mapping[int, Mapping[str, float]] | None,
+        *,
+        relative_to_wt: bool = True,
+        name: str = "inverse_folding",
+    ) -> None:
+        self.name = name
+        self._relative = relative_to_wt
+        self._lp = logprobs
+        if logprobs is not None:
+            missing = [p for p in space.positions if p not in logprobs]
+            if missing:
+                raise ValueError(f"inverse-folding logprobs missing positions {missing}")
+
+    @property
+    def dim(self) -> int:
+        return 1
+
+    @property
+    def present(self) -> bool:
+        return self._lp is not None
+
+    def _encode_present(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray:
+        assert self._lp is not None
+        out = np.zeros((len(assignments), 1), dtype=float)
+        for i, a in enumerate(assignments):
+            s = 0.0
+            for p in space.positions:
+                table = self._lp[p]
+                s += float(table.get(a[p], 0.0))
+                if self._relative:
+                    s -= float(table.get(space.reference[p], 0.0))
+            out[i, 0] = s
+        return out

shellde/features/matrix.py

@@ -0,0 +1,50 @@
+"""FeatureMatrix: compose feature blocks into one column-stable design matrix.
+
+The column layout depends only on (space, blocks), never on the candidate pool of a
+given round, so train and predict matrices stay aligned across active-learning
+rounds that select different pools. This is the single encoding path (there is no
+separate legacy encoder).
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from shellde.design_space import DesignSpace, combo_indices, to_assignment
+from shellde.features.base import FeatureBlock
+
+
+@dataclass
+class FeatureMatrix:
+    space: DesignSpace
+    blocks: list[FeatureBlock]
+
+    @property
+    def width(self) -> int:
+        return sum(b.dim for b in self.blocks)
+
+    def column_layout(self) -> list[tuple[str, int, int]]:
+        layout: list[tuple[str, int, int]] = []
+        cursor = 0
+        for b in self.blocks:
+            layout.append((b.name, cursor, cursor + b.dim))
+            cursor += b.dim
+        return layout
+
+    def active_blocks(self) -> list[str]:
+        return [b.name for b in self.blocks if b.present]
+
+    def encode(self, variants: Sequence[str]) -> np.ndarray:
+        if not self.blocks:
+            return np.zeros((len(variants), 0), dtype=float)
+        # Fast path: every variant is a valid combo AND every block supports index encoding.
+        # Skips per-variant assignment dicts + per-element Python loops; byte-identical output.
+        if all(hasattr(b, "encode_indices") for b in self.blocks):
+            idx = combo_indices(variants, self.space)
+            if idx is not None:
+                return np.hstack([b.encode_indices(idx, self.space) for b in self.blocks])
+        assignments = [to_assignment(v, self.space) for v in variants]
+        parts = [b.encode(assignments, self.space) for b in self.blocks]
+        return np.hstack(parts)

shellde/features/naturalness.py

@@ -0,0 +1,65 @@
+"""NaturalnessBlock: a per-variant zero-shot naturalness score as one feature column.
+
+This is the protocol-clean realisation of FolDE's naturalness warm-start: rather than
+a separate pretrain stage, the PLM zero-shot naturalness of each variant (spliced into
+the full WT background, like the embedding block) is exposed as a feature the surrogate
+weights. Because the column is available for UNMEASURED variants too, predictions stay
+anchored to the pretraining prior where activity data is sparse, which is exactly the
+anti-round-2-collapse benefit. Pairs naturally with the ranking surrogate.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.features.base import FeatureBlock
+from shellde.protocols import NaturalnessProvider
+
+
+class NaturalnessBlock(FeatureBlock):
+    """Single-column per-variant naturalness spliced into a full WT background (optional)."""
+
+    name = "naturalness"
+
+    def __init__(
+        self,
+        space: DesignSpace,
+        provider: NaturalnessProvider | None,
+        wt_sequence: str | None,
+        *,
+        name: str = "naturalness",
+    ) -> None:
+        self.name = name
+        self._provider = provider
+        self._wt = wt_sequence
+        if provider is not None and wt_sequence is not None:
+            for p in space.positions:
+                if p < 1 or p > len(wt_sequence):
+                    raise ValueError(f"position {p} outside wt_sequence length {len(wt_sequence)}")
+                if wt_sequence[p - 1] != space.reference[p]:
+                    raise ValueError(
+                        f"wt_sequence[{p - 1}]={wt_sequence[p - 1]!r} != reference "
+                        f"{space.reference[p]!r} at position {p}"
+                    )
+
+    @property
+    def dim(self) -> int:
+        return 1
+
+    @property
+    def present(self) -> bool:
+        return self._provider is not None and self._wt is not None
+
+    def _encode_present(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray:
+        assert self._provider is not None and self._wt is not None
+        seqs: list[str] = []
+        for a in assignments:
+            chars = list(self._wt)
+            for p in space.positions:
+                chars[p - 1] = a[p]
+            seqs.append("".join(chars))
+        return np.asarray(self._provider.score(seqs), dtype=float).reshape(-1, 1)

shellde/features/onehot.py

@@ -0,0 +1,55 @@
+"""OneHotBlock: absolute-position-keyed one-hot over the frozen design space.
+
+Always present. This is the cheap, always-available, additive-friendly feature
+that the surrogate falls back to when richer blocks (PLM, structure) are absent or
+uninformative (e.g. non-native function), per the architecture's scope notes.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+from shellde.design_space import DesignSpace
+from shellde.features.base import FeatureBlock
+
+
+class OneHotBlock(FeatureBlock):
+    name = "onehot"
+
+    def __init__(self, space: DesignSpace) -> None:
+        self._dim = space.n_positions * space.q
+        self._aa_to_int = {aa: i for i, aa in enumerate(space.alphabet)}
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    def _encode_present(
+        self, assignments: Sequence[Mapping[int, str]], space: DesignSpace
+    ) -> np.ndarray:
+        q = space.q
+        x = np.zeros((len(assignments), space.n_positions * q), dtype=float)
+        for i, a in enumerate(assignments):
+            for j, pos in enumerate(space.positions):
+                idx = self._aa_to_int.get(a[pos])
+                if idx is not None:
+                    x[i, j * q + idx] = 1.0
+        return x
+
+    def encode_indices(self, idx: np.ndarray, space: DesignSpace) -> np.ndarray:
+        """Vectorized one-hot from a (N, n_positions) alphabet-index array (fast path).
+
+        Produces the IDENTICAL matrix to ``_encode_present`` for combo variants: bit set at
+        column ``j * q + idx[i, j]`` for every variant i, position j. Used by FeatureMatrix.encode
+        when every variant is a valid combo (see design_space.combo_indices); otherwise the
+        per-assignment path runs. Indices are guaranteed in-alphabet by combo_indices.
+        """
+        n, length = idx.shape
+        q = space.q
+        x = np.zeros((n, length * q), dtype=float)
+        if n:
+            cols = (idx + np.arange(length, dtype=np.intp) * q).ravel()
+            rows = np.repeat(np.arange(n, dtype=np.intp), length)
+            x[rows, cols] = 1.0
+        return x