dee/core/active_learning.py

"""Active-learning surrogate for the Directed-Evolution loop (Design→Build→Test→Learn).

A user evolves a protein (round 1, zero-shot ESM-2 ΔLL), orders the top library,
measures variants on the bench, and logs the results. This module fits a light
surrogate on THEIR measurements and produces an *adjusted* per-mutation score
that the existing simulated-annealing search (`dee.optimizer.search.evolve`)
consumes unchanged — so round 2 is conditioned on real data, not just the prior.

Design (deliberately simple + honest, runs in milliseconds on CPU; numpy only,
no scikit-learn/scipy in the image):

  * The SA objective is ADDITIVE over single-site effects:
        fitness(variant) = Σ_{m∈variant} score(m)
    In round 1, score(m) = ΔLL_m (the ESM-2 wild-type-marginal prior).

  * We model the measured fitness as a ridge over single-site effects WITH the
    ΔLL prior as a feature:
        y_std ≈ b + w_prior · (Σ ΔLL of the variant) + Σ_{m} β_m · 1[m∈variant]
    Fit [b, w_prior, β] by closed-form ridge (β strongly shrunk toward 0, so
    mutations the user never measured fall back to the pure prior). y is
    standardized so the (arbitrary) assay scale doesn't matter — ranking is
    scale-invariant anyway.

  * The round-2 per-mutation acquisition score (additive ⇒ plugs straight into
    the SA) is:
        score(m) = w_prior · ΔLL_m + β_m + κ · uncertainty_m
    where uncertainty_m rewards under-measured mutations (exploration). Unseen
    mutations get β_m = 0 and high uncertainty → the search explores them while
    still respecting the ΔLL prior.

  * Below a small floor of measurements we DON'T pretend to learn: we return the
    pure prior with an honest note. No overclaiming.

Privacy: operates only on the data passed in (one user's own measurements).
"""
from __future__ import annotations

import re
from dataclasses import dataclass, replace
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np

# Need at least this many measured variants before we trust a learned signal;
# below it, round 2 = re-search on the zero-shot prior (still a fresh library).
MIN_MEASUREMENTS = 4

_LABEL_RE = re.compile(r"^([A-Za-z])(\d+)([A-Za-z*])$")   # e.g. "W58L"


def parse_label(label: str) -> Optional[Tuple[int, str]]:
    """'W58L' → (57, 'L')  (0-indexed position, mutant AA). None if malformed."""
    m = _LABEL_RE.match((label or "").strip())
    if not m:
        return None
    pos = int(m.group(2)) - 1
    if pos < 0:
        return None
    return (pos, m.group(3).upper())


def parse_mutations(labels: str | Sequence[str]) -> List[Tuple[int, str]]:
    """Accept 'W58L,K204R' or ['W58L','K204R'] → [(57,'L'),(203,'R')]."""
    if isinstance(labels, str):
        parts = re.split(r"[,\s;]+", labels.strip())
    else:
        parts = list(labels)
    out = []
    for p in parts:
        pm = parse_label(p)
        if pm is not None:
            out.append(pm)
    return out


@dataclass
class Surrogate:
    """Result of fitting on a user's measurements."""
    adjusted: Dict[Tuple[int, str], float]   # (pos, mut_aa) → round-2 acquisition score
    w_prior: float                           # learned weight on the ΔLL prior
    n_train: int                             # measured variants used
    n_effects: int                           # mutations that got a learned correction
    learned: bool                            # False ⇒ fell back to the prior
    note: str

    def adjust_pool(self, pool: list) -> list:
        """Return a copy of a `search.Mutation` pool with delta_ll replaced by
        the round-2 acquisition score (so `evolve()` runs unchanged)."""
        return [replace(m, delta_ll=self.adjusted.get((m.position, m.mut_aa), m.delta_ll))
                for m in pool]


def fit_surrogate(
    pool: list,
    measurements: List[Tuple[Sequence[str], float]],
    *,
    kappa: float = 0.4,
    ridge_lambda: float = 1.0,
    prior_lambda: float = 0.1,
) -> Surrogate:
    """Fit the additive surrogate.

    pool:         list of search.Mutation (round-1 single-site pool; each has
                  .position, .mut_aa, .delta_ll).
    measurements: [(mutation_labels, measured_value), …] from the user's bench.
    Returns a Surrogate whose `adjusted` maps (pos, mut_aa) → acquisition score.
    """
    prior = {(m.position, m.mut_aa): float(m.delta_ll) for m in pool}
    index = {key: i for i, key in enumerate(prior.keys())}
    keys = list(prior.keys())
    M = len(keys)

    # Parse + keep only measurements with a numeric value and ≥1 in-pool mutation.
    rows: List[Tuple[List[int], float, float]] = []   # (col_indices, prior_sum, y)
    for labels, value in (measurements or []):
        try:
            y = float(value)
        except (TypeError, ValueError):
            continue
        cols, psum = [], 0.0
        for key in parse_mutations(labels):
            if key in index:
                cols.append(index[key]); psum += prior[key]
        if cols:
            rows.append((cols, psum, y))

    n = len(rows)
    counts = np.zeros(M)
    for cols, _, _ in rows:
        for c in cols:
            counts[c] += 1
    # Exploration bonus: under-measured mutations get a larger nudge.
    unc = 1.0 / np.sqrt(1.0 + counts)

    # Not enough signal → honest fallback to the pure ΔLL prior.
    y_all = np.array([y for _, _, y in rows], dtype=float)
    if n < MIN_MEASUREMENTS or M == 0 or (n and np.std(y_all) < 1e-9):
        return Surrogate(
            adjusted=dict(prior), w_prior=1.0, n_train=n, n_effects=0, learned=False,
            note=(f"{n} measurement(s) logged — need ≥{MIN_MEASUREMENTS} with a spread of "
                  "values to learn; round 2 uses the ESM-2 prior."),
        )

    # Standardize y (assay scale is arbitrary; ranking is scale-invariant).
    y_std = (y_all - y_all.mean()) / (y_all.std() + 1e-9)

    # Design matrix A = [intercept | prior_sum | incidence(M)].
    A = np.zeros((n, 2 + M))
    A[:, 0] = 1.0
    for i, (cols, psum, _) in enumerate(rows):
        A[i, 1] = psum
        for c in cols:
            A[i, 2 + c] = 1.0

    # Ridge: don't regularize the intercept; lightly regularize the prior weight;
    # strongly shrink the per-mutation β toward 0 (→ unseen muts default to prior).
    reg = np.concatenate([[0.0], [prior_lambda], np.full(M, ridge_lambda)])
    try:
        w = np.linalg.solve(A.T @ A + np.diag(reg), A.T @ y_std)
    except np.linalg.LinAlgError:
        w = np.linalg.lstsq(A.T @ A + np.diag(reg), A.T @ y_std, rcond=None)[0]

    w_prior = float(w[1])
    beta = w[2:]
    # Acquisition per pool mutation: prior (re-weighted) + learned correction + explore.
    adjusted = {}
    for key in keys:
        c = index[key]
        adjusted[key] = w_prior * prior[key] + float(beta[c]) + kappa * float(unc[c])
    n_effects = int(np.sum(np.abs(beta) > 1e-6))
    return Surrogate(
        adjusted=adjusted, w_prior=w_prior, n_train=n, n_effects=n_effects, learned=True,
        note=(f"Learned from {n} measured variants (prior weight {w_prior:.2f}; "
              f"{n_effects} mutation effects corrected). Round 2 balances the "
              "learned model with exploration of under-tested positions."),
    )