bunyamineymen's picture
Project 01 code (8e118d5)
668f59a verified
Raw History Blame Contribute Delete
5.2 kB
"""Golden-set subsampling simulation: gate decisions on R golden sets vs the population truth."""
import numpy as np
from . import gates, stats
EXACT = ["bm25", "bge", "minilm", "e5"]
def scenarios(cfg):
"""(family, A, B) system pairs; A is the baseline, B the candidate."""
seeds = cfg["hnsw"]["seeds"]
pairs = [("noop", f"hnsw_ef{ef}_s{seeds[0]}", f"hnsw_ef{ef}_s{s}") for ef in (16, 128) for s in seeds[1:]]
pairs += [("natural", "bge", "minilm"), ("natural", "bge", "e5"), ("natural", "bge", "bm25"),
("natural", "hnsw_ef128_s0", "hnsw_ef16_s0")]
pairs += [("natural", "bge", f"del{round(r * 100):02d}") for r in cfg["deletion_rates"]]
pairs += [("controlled", "bge", f"ctl{round(d * 100):02d}") for d in cfg["controlled_deltas"]]
return pairs
def controlled(base, deltas, rng):
"""Zero all metrics of randomly ordered queries until population recall@5 drops by each delta."""
N = len(base["recall@5"])
order = rng.permutation(N)
order = order[base["recall@5"][order] > 0]
cum = np.cumsum(base["recall@5"][order]) / N
out = {}
for delta in deltas:
chosen = order[: np.searchsorted(cum, delta) + 1]
out[f"ctl{round(delta * 100):02d}"] = {m: np.where(np.isin(np.arange(N), chosen), 0.0, v) for m, v in base.items()}
return out
def _contrast(rows, key, x):
mean, lo, hi = stats.mean_ci(x)
rows.append({**key, "mean": mean, "lo": lo, "hi": hi})
def simulate(name, ds_i, systems, cfg, R=None):
sim, seed = cfg["simulation"], cfg["seed"]
R, B, alpha = R or sim["R"], sim["B"], sim["alpha"]
N = len(systems["bge"]["recall@5"])
split = np.random.default_rng([seed, ds_i, 99]).permutation(N)
truths = {"population": (np.arange(N), np.arange(N)), "split_half": (split[: N // 2], split[N // 2:])}
rows_t, rows_p, rows_c = [], [], []
for ti, (truth, (pool, tidx)) in enumerate(truths.items()):
for n in [n for n in sim["n"] if n <= N // 2]:
rng = np.random.default_rng([seed, ds_i, ti, n])
S = rng.permuted(np.tile(pool, (R, 1)), axis=1)[:, :n]
Wa = stats.bootstrap_weights(n, B, rng).astype(np.float32)
Wb = stats.bootstrap_weights(n, B, rng).astype(np.float32)
flips = rng.choice(np.array([-1.0, 1.0], np.float32), size=(B, n))
for metric in sim["metrics"]:
for sysname in EXACT:
v = systems[sysname][metric]
true, est = v[tidx].mean(), v[S].mean(1)
for d in sim["d"]:
rows_t.append(dict(dataset=name, truth=truth, system=sysname, metric=metric, n=n, d=d,
true_value=true, truth_pass=d >= 0,
k_pass=int(gates.g1_threshold(v[S], true - d).sum()), R=R))
family = {}
for fam, A, Bn in scenarios(cfg):
a, b = systems[A][metric], systems[Bn][metric]
XA, XB = a[S].astype(np.float32), b[S].astype(np.float32)
delta, sd_d = (b - a)[tidx].mean(), (b - a)[pool].std(ddof=1)
lb_p = stats.lower_bound((XB - XA) @ Wa.T, alpha)
lb_u = stats.lower_bound(XB @ Wb.T - XA @ Wa.T, alpha)
key = dict(dataset=name, truth=truth, family=fam, A=A, B=Bn, metric=metric, n=n, R=R)
base = {**key, "delta_true": delta, "sd_d": sd_d}
g1 = {h: gates.g1_threshold(XB, a[tidx].mean() - h) for h in sim["headroom"]}
for h, x in g1.items():
rows_p.append({**base, "gate": "G1", "param": h, "k_pass": int(x.sum())})
for m in sim["margins"]:
dec = {"G2": gates.g2_naive_delta(XA, XB, m), "G3": lb_p >= -m, "G3u": lb_u >= -m,
"G3signflip": gates.g3_sign_flip(XA, XB, m, flips, alpha)}
for g, x in dec.items():
rows_p.append({**base, "gate": g, "param": m, "k_pass": int(x.sum()),
"analytic_pass": stats.noninferiority_pass_prob(delta, m, sd_d, n, alpha)
if g == "G3" else np.nan})
h2 = g1[max(sim["headroom"])].astype(float) - dec["G3"]
h3 = (~dec["G3u"]).astype(float) - (~dec["G3"])
_contrast(rows_c, {**key, "param": m, "contrast": "G1pass-G3pass"}, h2)
_contrast(rows_c, {**key, "param": m, "contrast": "G3ufail-G3fail"}, h3)
if fam == "noop":
family[m] = family.get(m, 0) + h3
n_noop = sum(f == "noop" for f, _, _ in scenarios(cfg))
for m, total in family.items():
_contrast(rows_c, dict(dataset=name, truth=truth, family="noop", A="*", B="*", metric=metric,
n=n, R=R, param=m, contrast="G3ufail-G3fail"), total / n_noop)
return rows_t, rows_p, rows_c