Spaces:

WINTER4000
/

syntheogenesis

Running

syntheogenesis / dee /core /active_learning.py

Tengo Gzirishvili

Claude Sonnet 5

Surface the active-learning surrogate as the Learn-phase "wow" moment

874f438 3 months ago

9.8 kB

"""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
    # Per-mutation breakdown of `adjusted[key]` into its three additive terms
    # (prior + learned + explore ≈ adjusted[key]) plus how many of the user's
    # measured variants included this mutation. Exists purely so a caller can
    # NARRATE round 2 ("recommended because your data corrected this position")
    # and VISUALIZE the before/after re-rank — adjusted/adjust_pool are
    # unaffected and existing callers/tests don't need to know this exists.
    components: Dict[Tuple[int, str], Dict[str, float]]

    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. Components
    # still populate (prior-only, no correction, no explore) so a caller can
    # render "not enough data yet" consistently rather than special-casing it.
    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):
        fallback_components = {
            key: {"prior": prior[key], "learned": 0.0, "explore": 0.0,
                  "n_measured": int(counts[index[key]])}
            for key in keys
        }
        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."),
            components=fallback_components,
        )

    # 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, components = {}, {}
    for key in keys:
        c = index[key]
        prior_term = w_prior * prior[key]
        learned_term = float(beta[c])
        explore_term = kappa * float(unc[c])
        adjusted[key] = prior_term + learned_term + explore_term
        components[key] = {"prior": prior_term, "learned": learned_term,
                            "explore": explore_term, "n_measured": int(counts[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."),
        components=components,
    )


# β magnitude below this is "no real correction" for narration purposes — a
# separate, looser threshold than n_effects' 1e-6 (that one's for counting
# whether ridge found ANY signal; this one's for whether it's worth a sentence).
_NARRATE_LEARNED_FLOOR = 0.05


def narrate(prior: float, learned: float, explore: float, n_measured: int) -> str:
    """One short, honest sentence for why a mutation's round-2 score moved (or
    didn't) — the same three components `fit_surrogate` computes, turned into
    plain language. Deterministic and template-only: never invents a reason
    beyond what the numbers actually show."""
    if n_measured == 0:
        return ("Not yet tested — kept close to the ESM-2 prior, nudged up for "
                "exploration since your data hasn't covered it.")
    if abs(learned) < _NARRATE_LEARNED_FLOOR:
        return (f"Consistent with your {n_measured} measurement"
                f"{'s' if n_measured != 1 else ''} so far — no correction needed.")
    direction = "helped more than the ESM-2 prior expected" if learned > 0 else \
                "underperformed what the ESM-2 prior expected"
    return (f"Revised from {n_measured} of your measurement"
            f"{'s' if n_measured != 1 else ''} — this substitution {direction}.")

	"""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
	# Per-mutation breakdown of `adjusted[key]` into its three additive terms
	# (prior + learned + explore ≈ adjusted[key]) plus how many of the user's
	# measured variants included this mutation. Exists purely so a caller can
	# NARRATE round 2 ("recommended because your data corrected this position")
	# and VISUALIZE the before/after re-rank — adjusted/adjust_pool are
	# unaffected and existing callers/tests don't need to know this exists.
	components: Dict[Tuple[int, str], Dict[str, float]]

	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. Components
	# still populate (prior-only, no correction, no explore) so a caller can
	# render "not enough data yet" consistently rather than special-casing it.
	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):
	fallback_components = {
	key: {"prior": prior[key], "learned": 0.0, "explore": 0.0,
	"n_measured": int(counts[index[key]])}
	for key in keys
	}
	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."),
	components=fallback_components,
	)

	# 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, components = {}, {}
	for key in keys:
	c = index[key]
	prior_term = w_prior * prior[key]
	learned_term = float(beta[c])
	explore_term = kappa * float(unc[c])
	adjusted[key] = prior_term + learned_term + explore_term
	components[key] = {"prior": prior_term, "learned": learned_term,
	"explore": explore_term, "n_measured": int(counts[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."),
	components=components,
	)


	# β magnitude below this is "no real correction" for narration purposes — a
	# separate, looser threshold than n_effects' 1e-6 (that one's for counting
	# whether ridge found ANY signal; this one's for whether it's worth a sentence).
	_NARRATE_LEARNED_FLOOR = 0.05


	def narrate(prior: float, learned: float, explore: float, n_measured: int) -> str:
	"""One short, honest sentence for why a mutation's round-2 score moved (or
	didn't) — the same three components `fit_surrogate` computes, turned into
	plain language. Deterministic and template-only: never invents a reason
	beyond what the numbers actually show."""
	if n_measured == 0:
	return ("Not yet tested — kept close to the ESM-2 prior, nudged up for "
	"exploration since your data hasn't covered it.")
	if abs(learned) < _NARRATE_LEARNED_FLOOR:
	return (f"Consistent with your {n_measured} measurement"
	f"{'s' if n_measured != 1 else ''} so far — no correction needed.")
	direction = "helped more than the ESM-2 prior expected" if learned > 0 else \
	"underperformed what the ESM-2 prior expected"
	return (f"Revised from {n_measured} of your measurement"
	f"{'s' if n_measured != 1 else ''} — this substitution {direction}.")