Spaces:
Running
Running
Tengo Gzirishvili
Claude Sonnet 5
Surface the active-learning surrogate as the Learn-phase "wow" moment
874f438 | """Tests for the Directed-Evolution active-learning surrogate | |
| (`dee.core.active_learning`) — the Design→Build→Test→Learn round-2 core. | |
| These pin the behaviour we actually claim to users: | |
| * mutation-label parsing (1-indexed → 0-indexed), | |
| * honest fallback to the ESM-2 prior below the measurement floor / with no | |
| signal (no overclaiming), | |
| * the headline property: when the zero-shot prior is *wrong*, real bench | |
| measurements override it and the round-2 ranking recovers the truth, | |
| * `adjust_pool` plugs the learned scores back into the SA `Mutation` pool | |
| without disturbing the other fields. | |
| No torch/ESM needed — the surrogate is pure numpy over a supplied pool. | |
| """ | |
| import numpy as np | |
| import pytest | |
| from dee.core.active_learning import ( | |
| MIN_MEASUREMENTS, | |
| Surrogate, | |
| fit_surrogate, | |
| narrate, | |
| parse_label, | |
| parse_mutations, | |
| ) | |
| from dee.optimizer.search import Mutation | |
| # --------------------------------------------------------------------------- # | |
| # label parsing | |
| # --------------------------------------------------------------------------- # | |
| def test_parse_label_basic(): | |
| assert parse_label("W58L") == (57, "L") # 1-indexed -> 0-indexed | |
| assert parse_label(" k204r ") == (203, "R") # whitespace + case | |
| assert parse_label("A1*") == (0, "*") # stop codon, first residue | |
| def test_parse_label_rejects_malformed(bad): | |
| assert parse_label(bad) is None | |
| def test_parse_mutations_string_and_list(): | |
| assert parse_mutations("W58L,K204R") == [(57, "L"), (203, "R")] | |
| assert parse_mutations("W58L K204R; T2A") == [(57, "L"), (203, "R"), (1, "A")] | |
| assert parse_mutations(["W58L", "K204R"]) == [(57, "L"), (203, "R")] | |
| # malformed tokens are silently dropped, valid ones kept | |
| assert parse_mutations("W58L, junk, K204R") == [(57, "L"), (203, "R")] | |
| # --------------------------------------------------------------------------- # | |
| # helpers | |
| # --------------------------------------------------------------------------- # | |
| def _pool(prior_by_key): | |
| """Build a frozen-Mutation pool from {(pos, aa): delta_ll}.""" | |
| return [ | |
| Mutation(position=pos, wt_aa="A", mut_aa=aa, delta_ll=float(dll)) | |
| for (pos, aa), dll in prior_by_key.items() | |
| ] | |
| def _label(pos, aa): | |
| return f"A{pos + 1}{aa}" # inverse of parse_label, wt 'A' is irrelevant | |
| def _spearman(a, b): | |
| a, b = np.asarray(a, float), np.asarray(b, float) | |
| ra = np.argsort(np.argsort(a)) | |
| rb = np.argsort(np.argsort(b)) | |
| return float(np.corrcoef(ra, rb)[0, 1]) | |
| # --------------------------------------------------------------------------- # | |
| # honest fallback | |
| # --------------------------------------------------------------------------- # | |
| def test_fallback_below_measurement_floor(): | |
| pool = _pool({(0, "L"): 1.0, (1, "R"): 0.5}) | |
| meas = [([_label(0, "L")], 3.0), ([_label(1, "R")], 1.0)] # only 2 < floor | |
| s = fit_surrogate(pool, meas) | |
| assert s.learned is False | |
| assert s.n_effects == 0 | |
| assert s.w_prior == 1.0 | |
| # adjusted == the pure prior, so round 2 == a fresh search on ESM-2 | |
| assert s.adjusted == {(0, "L"): 1.0, (1, "R"): 0.5} | |
| assert str(MIN_MEASUREMENTS) in s.note | |
| def test_fallback_when_no_spread_in_values(): | |
| pool = _pool({(i, "L"): 0.1 * i for i in range(5)}) | |
| meas = [([_label(i, "L")], 2.0) for i in range(5)] # ≥ floor but zero variance | |
| s = fit_surrogate(pool, meas) | |
| assert s.learned is False | |
| assert s.adjusted == {(i, "L"): 0.1 * i for i in range(5)} | |
| def test_empty_inputs_dont_crash(): | |
| s = fit_surrogate([], []) | |
| assert s.learned is False | |
| assert s.adjusted == {} | |
| def test_non_numeric_and_out_of_pool_measurements_ignored(): | |
| pool = _pool({(0, "L"): 1.0}) | |
| meas = [ | |
| ([_label(0, "L")], "not-a-number"), # bad value -> dropped | |
| (["Z9Q"], 5.0), # mutation not in pool -> dropped | |
| ] | |
| s = fit_surrogate(pool, meas) | |
| assert s.n_train == 0 | |
| assert s.learned is False | |
| # --------------------------------------------------------------------------- # | |
| # the headline property: data overrides a misleading prior | |
| # --------------------------------------------------------------------------- # | |
| def test_data_overrides_misleading_prior(): | |
| """Ground-truth additive effects exist; the ESM-2 prior is deliberately | |
| *anti-correlated* with them. After logging enough bench measurements, the | |
| round-2 acquisition ranking should recover the truth far better than the | |
| prior did — that's the whole value proposition.""" | |
| rng = np.random.default_rng(7) | |
| keys = [(i, "L") for i in range(8)] | |
| true_effect = {k: v for k, v in zip(keys, [3.0, -2.0, 2.5, -1.0, 1.5, -2.5, 0.5, -0.5])} | |
| # Prior is the NEGATION of the truth (worst case): high-prior = actually bad. | |
| prior = {k: -true_effect[k] for k in keys} | |
| pool = _pool(prior) | |
| # Simulate the user measuring random multi-site variants; value = additive | |
| # truth + small noise (scale is arbitrary; surrogate standardizes it). | |
| measurements = [] | |
| for _ in range(40): | |
| n_sites = int(rng.integers(1, 4)) | |
| chosen = [keys[i] for i in rng.choice(len(keys), size=n_sites, replace=False)] | |
| val = sum(true_effect[k] for k in chosen) + rng.normal(0, 0.25) | |
| measurements.append(([_label(p, a) for (p, a) in chosen], 10.0 + val)) | |
| s = fit_surrogate(pool, measurements) | |
| assert s.learned is True | |
| assert s.n_train == 40 | |
| assert s.n_effects > 0 | |
| true_vec = [true_effect[k] for k in keys] | |
| prior_vec = [prior[k] for k in keys] | |
| learned_vec = [s.adjusted[k] for k in keys] | |
| rho_prior = _spearman(prior_vec, true_vec) | |
| rho_learned = _spearman(learned_vec, true_vec) | |
| assert rho_prior < 0 # prior is anti-correlated (by construction) | |
| assert rho_learned > 0.6 # learned ranking recovers the truth | |
| assert rho_learned > rho_prior + 1.0 # and is a large improvement over the prior | |
| def test_adjust_pool_preserves_fields_and_swaps_score(): | |
| pool = _pool({(0, "L"): 1.0, (3, "R"): -0.5, (5, "K"): 0.2}) | |
| # enough varied data to learn | |
| meas = [ | |
| ([_label(0, "L")], 9.0), | |
| ([_label(3, "R")], 2.0), | |
| ([_label(5, "K")], 5.0), | |
| ([_label(0, "L"), _label(5, "K")], 8.0), | |
| ([_label(3, "R"), _label(5, "K")], 4.0), | |
| ] | |
| s = fit_surrogate(pool, meas) | |
| new_pool = s.adjust_pool(pool) | |
| assert len(new_pool) == len(pool) | |
| for old, new in zip(pool, new_pool): | |
| # identity fields untouched... | |
| assert (new.position, new.wt_aa, new.mut_aa) == (old.position, old.wt_aa, old.mut_aa) | |
| # ...delta_ll replaced by the learned acquisition score | |
| assert new.delta_ll == pytest.approx(s.adjusted[(old.position, old.mut_aa)]) | |
| # frozen dataclass: originals are unchanged | |
| assert pool[0].delta_ll == 1.0 | |
| def test_components_populated_in_fallback(): | |
| """Below the measurement floor, components are still present (prior-only, | |
| no correction, no explore) so a caller never has to special-case the | |
| 'not learned yet' branch when rendering.""" | |
| pool = _pool({(0, "L"): 1.0, (1, "R"): 0.5}) | |
| meas = [([_label(0, "L")], 3.0)] # 1 < floor | |
| s = fit_surrogate(pool, meas) | |
| assert s.learned is False | |
| assert s.components[(0, "L")] == {"prior": 1.0, "learned": 0.0, "explore": 0.0, "n_measured": 1} | |
| assert s.components[(1, "R")] == {"prior": 0.5, "learned": 0.0, "explore": 0.0, "n_measured": 0} | |
| def test_components_sum_to_adjusted_when_learned(): | |
| """prior + learned + explore must reconstruct adjusted[key] exactly — | |
| components is a DECOMPOSITION of the same score, not a separate estimate.""" | |
| pool = _pool({(0, "L"): 1.0, (3, "R"): -0.5, (5, "K"): 0.2}) | |
| meas = [ | |
| ([_label(0, "L")], 9.0), | |
| ([_label(3, "R")], 2.0), | |
| ([_label(5, "K")], 5.0), | |
| ([_label(0, "L"), _label(5, "K")], 8.0), | |
| ([_label(3, "R"), _label(5, "K")], 4.0), | |
| ] | |
| s = fit_surrogate(pool, meas) | |
| assert s.learned is True | |
| for key, comp in s.components.items(): | |
| total = comp["prior"] + comp["learned"] + comp["explore"] | |
| assert total == pytest.approx(s.adjusted[key]) | |
| assert comp["n_measured"] >= 0 | |
| def test_narrate_unmeasured(): | |
| text = narrate(prior=1.0, learned=0.0, explore=0.3, n_measured=0) | |
| assert "not yet tested" in text.lower() | |
| def test_narrate_consistent_with_small_correction(): | |
| text = narrate(prior=1.0, learned=0.01, explore=0.05, n_measured=3) | |
| assert "consistent" in text.lower() | |
| assert "3" in text | |
| def test_narrate_revised_up_and_down(): | |
| up = narrate(prior=1.0, learned=0.5, explore=0.0, n_measured=4) | |
| down = narrate(prior=1.0, learned=-0.5, explore=0.0, n_measured=4) | |
| assert "helped more" in up.lower() | |
| assert "underperformed" in down.lower() | |
| assert up != down | |
| def test_unmeasured_mutation_falls_back_toward_prior(): | |
| """A mutation never measured gets β≈0, so its adjusted score stays driven by | |
| the (re-weighted) prior plus only the exploration bonus — it isn't invented.""" | |
| pool = _pool({(0, "L"): 1.0, (1, "R"): 0.8, (2, "K"): 0.6, (3, "D"): 0.4, (9, "Q"): 5.0}) | |
| # measure everything EXCEPT (9, 'Q') | |
| meas = [ | |
| ([_label(0, "L")], 3.0), | |
| ([_label(1, "R")], 2.0), | |
| ([_label(2, "K")], 1.0), | |
| ([_label(3, "D")], 0.5), | |
| ([_label(0, "L"), _label(1, "R")], 4.5), | |
| ] | |
| s = fit_surrogate(pool, meas) | |
| assert s.learned is True | |
| # the unmeasured high-prior mutation is still present and finite | |
| assert (9, "Q") in s.adjusted | |
| assert np.isfinite(s.adjusted[(9, "Q")]) | |