"""Tests for the cross-user aggregate prior (`dee.core.aggregate`). Pins the privacy + correctness invariants we actually promise: * aggregates by substitution TYPE, standardized within each library; * k-anonymity: a substitution measured by < MIN_USERS distinct users is dropped and never appears in the output; * the execution gate refuses before the policy effective date; * the resulting prior recovers the right sign and folds into the SA pool. Pure numpy — no DB, no torch. """ import datetime as dt import numpy as np import pytest from dee.core import aggregate as agg from dee.core.aggregate import ( EFFECTIVE_DATE, AggregationGateError, GlobalPrior, apply_global_prior, build_priors, parse_sub, ) from dee.optimizer.search import Mutation AFTER = EFFECTIVE_DATE # gate passes on/after the date BEFORE = EFFECTIVE_DATE - dt.timedelta(days=1) # --------------------------------------------------------------------------- # def test_parse_sub(): assert parse_sub("W58L") == ("W", "L") assert parse_sub(" k204r ") == ("K", "R") assert parse_sub("L10L") is None # synonymous assert parse_sub("junk") is None # --------------------------------------------------------------------------- # # the gate # --------------------------------------------------------------------------- # def test_gate_blocks_before_effective_date(): with pytest.raises(AggregationGateError): build_priors([("u1", [("A2V", 1.0), ("A2V,L3I", 2.0)])], now=BEFORE) def test_gate_opens_on_effective_date(): # should not raise (may return empty if no signal, but must not be gated) out = build_priors([("u1", [("A2V", 1.0), ("L3I", 2.0)])], now=AFTER) assert isinstance(out, GlobalPrior) def test_enforce_gate_false_for_tests(): out = build_priors([("u1", [("A2V", 1.0), ("L3I", 2.0)])], now=BEFORE, enforce_gate=False) assert isinstance(out, GlobalPrior) # --------------------------------------------------------------------------- # # k-anonymity floor # --------------------------------------------------------------------------- # def _lib_for_sub(values): """A 2-row library exercising substitutions W>L (good) and D>A (bad).""" # variant 'W10L' high, 'D20A' low -> standardized gives +/- effects return [("W10L", values[0]), ("D20A", values[1])] def test_kanon_drops_substitutions_below_floor(): # W>L and D>A measured by only 2 distinct users -> must be dropped (floor=3) obls = [ ("u1", _lib_for_sub((2.0, 0.0))), ("u2", _lib_for_sub((2.0, 0.0))), ] out = build_priors(obls, min_users=3, now=AFTER) assert out.effects == {} # nothing meets the >=3-user floor # add a third distinct user -> now kept obls.append(("u3", _lib_for_sub((2.0, 0.0)))) out3 = build_priors(obls, min_users=3, now=AFTER) assert ("W", "L") in out3.effects and ("D", "A") in out3.effects assert out3.n_users[("W", "L")] == 3 def test_repeated_user_does_not_inflate_kanon(): # same user across two libraries is still ONE distinct user -> below floor obls = [ ("u1", _lib_for_sub((2.0, 0.0))), ("u1", _lib_for_sub((2.0, 0.0))), ("u2", _lib_for_sub((2.0, 0.0))), ] out = build_priors(obls, min_users=3, now=AFTER) assert out.effects == {} # only 2 distinct users # --------------------------------------------------------------------------- # # correctness: sign recovery + standardization # --------------------------------------------------------------------------- # def test_recovers_substitution_sign_across_users(): # W>L consistently beneficial, D>A consistently deleterious, 4 users, # each on a DIFFERENT arbitrary assay scale (tests within-library z-scoring) scales = [(1.0, 0.0), (100.0, 50.0), (5.0, -5.0), (0.2, 0.0)] # (gain, offset) obls = [] for i, (gain, off) in enumerate(scales): meas = [ ("W10L", off + gain * 1.0), # good ("W30L", off + gain * 0.8), # good (same sub type, other position) ("D20A", off + gain * -1.0), # bad ("W10L,D20A", off + gain * 0.0), # mixed ] obls.append((f"u{i}", meas)) out = build_priors(obls, min_users=3, now=AFTER) assert out.effects[("W", "L")] > 0 # beneficial recovered assert out.effects[("D", "A")] < 0 # deleterious recovered assert out.n_users[("W", "L")] == 4 assert out.n_obs[("W", "L")] >= 4 def test_zero_variance_or_tiny_library_contributes_nothing(): obls = [ ("u1", [("W10L", 5.0)]), # single row -> no effect ("u2", [("W10L", 3.0), ("D20A", 3.0)]), # zero variance -> dropped ] out = build_priors(obls, min_users=1, now=AFTER) assert out.effects == {} # --------------------------------------------------------------------------- # # serialization round-trip + pool blend # --------------------------------------------------------------------------- # def test_to_rows_is_deidentified_and_roundtrips(): obls = [("u%d" % i, _lib_for_sub((2.0, 0.0))) for i in range(3)] out = build_priors(obls, min_users=3, now=AFTER) rows = out.to_rows() assert rows and all(set(r) == {"substitution", "n_users", "n_obs", "mean_effect"} for r in rows) # no user ids, no positions, no raw values anywhere in the serialized form blob = str(rows) assert "u0" not in blob and "u1" not in blob and "10" not in "".join(r["substitution"] for r in rows) back = GlobalPrior.from_rows(rows) assert set(back.effects) == set(out.effects) def test_apply_global_prior_nudges_matching_substitutions(): prior = GlobalPrior(effects={("W", "L"): 2.0}, n_users={("W", "L"): 5}, n_obs={("W", "L"): 9}) pool = [ Mutation(position=57, wt_aa="W", mut_aa="L", delta_ll=1.0), # matches -> nudged Mutation(position=99, wt_aa="D", mut_aa="A", delta_ll=1.0), # no match -> unchanged ] out = apply_global_prior(pool, prior, weight=0.3) assert out[0].delta_ll == pytest.approx(1.0 + 0.3 * 2.0) assert out[1].delta_ll == 1.0 # identity fields preserved assert (out[0].position, out[0].wt_aa, out[0].mut_aa) == (57, "W", "L") def test_apply_global_prior_noop_without_prior(): pool = [Mutation(position=0, wt_aa="A", mut_aa="V", delta_ll=1.0)] assert apply_global_prior(pool, None)[0].delta_ll == 1.0 empty = GlobalPrior(effects={}, n_users={}, n_obs={}) assert apply_global_prior(pool, empty)[0].delta_ll == 1.0 # --------------------------------------------------------------------------- # # Richer aggregate: substitution × ESM ΔLL bin # --------------------------------------------------------------------------- # def test_bin_ll_thresholds(): assert agg.bin_ll(-5) == "lo" and agg.bin_ll(0) == "mid" and agg.bin_ll(5) == "hi" assert agg.bin_ll(None) == "mid" and agg.bin_ll("x") == "mid" def test_binned_build_splits_same_substitution_by_ll(): # W>L at a high-ΔLL position is beneficial; at a low-ΔLL position deleterious units = [(f"u{i}", [("W10L", 2.0), ("W30L", 0.0)]) for i in range(3)] ll = [{"W10L": 5.0, "W30L": -5.0}] * 3 # hi vs lo bin p = agg.build_priors(units, now=AFTER, ll_maps=ll) assert ("W", "L", "hi") in p.effects and ("W", "L", "lo") in p.effects assert p.effects[("W", "L", "hi")] > 0 and p.effects[("W", "L", "lo")] < 0 assert p.n_users[("W", "L", "hi")] == 3 # plain (unbinned) call is unchanged -> 2-tuple keys plain = agg.build_priors(units, now=AFTER) assert ("W", "L") in plain.effects and all(len(k) == 2 for k in plain.effects) def test_binned_serialize_roundtrip(): p = agg.GlobalPrior(effects={("W", "L", "hi"): 1.5}, n_users={("W", "L", "hi"): 4}, n_obs={("W", "L", "hi"): 9}) rows = p.to_rows() assert rows[0]["substitution"] == "W>L@hi" back = agg.GlobalPrior.from_rows(rows) assert ("W", "L", "hi") in back.effects and back.effects[("W", "L", "hi")] == 1.5 def test_binned_blend_matches_mutation_ll_bin(): from dee.optimizer.search import Mutation prior = agg.GlobalPrior( effects={("W", "L", "hi"): 2.0, ("W", "L", "lo"): -2.0}, n_users={("W", "L", "hi"): 5, ("W", "L", "lo"): 5}, n_obs={("W", "L", "hi"): 9, ("W", "L", "lo"): 9}) pool = [ Mutation(position=9, wt_aa="W", mut_aa="L", delta_ll=5.0), # hi bin -> + Mutation(position=29, wt_aa="W", mut_aa="L", delta_ll=-5.0), # lo bin -> - Mutation(position=40, wt_aa="W", mut_aa="L", delta_ll=0.0), # mid -> no matching key -> unchanged ] out = agg.apply_global_prior(pool, prior, weight=0.3) assert out[0].delta_ll == pytest.approx(5.0 + 0.3 * 2.0) assert out[1].delta_ll == pytest.approx(-5.0 + 0.3 * -2.0) assert out[2].delta_ll == 0.0 # mid bin, no fallback key