# %% [markdown]
# # 0210 - Data selection and contamination
# Companion to `notes/0210 - Data Selection and Contamination.pdf`.  CPU only, numpy/scipy, 1 thread.
# 1. classifier-based harvesting (fastText-style) as a selection problem;  2. pass-rate band selection as importance weighting;
# 3. regression to the mean when the same samples select and evaluate;  4. MinHash / Jaccard estimation error, LSH, containment;  5. membership inference from NLL gaps.
# %%
import os
os.environ["OMP_NUM_THREADS"] = "1"; os.environ["OPENBLAS_NUM_THREADS"] = "1"
import math, itertools
from math import comb
import numpy as np
from scipy import stats, special
CHECKS = []
def check(name, err, tol, note=""):
    CHECKS.append((name, float(err), float(tol)))
    assert err <= tol, (name, float(err), tol)
    print(f"[ok] {name:104s} err={float(err):.3e} (tol {tol:.1e}) {note}")
Phibar = stats.norm.sf
# %% [markdown]
# ## 1. A harvesting classifier is a detector; what you get depends on prevalence and on what a wrong page costs
# Pool of `M` pages, a fraction `pi` relevant ("math").  Score `s ~ N(d', 1)` for relevant, `N(0, 1)` otherwise (binormal; AUC = Phi(d'/sqrt 2)).
# Threshold `t`: TPR = Phibar(t - d'), FPR = Phibar(t).  Precision = pi TPR / (pi TPR + (1-pi) FPR) -- depends on `pi`, not only on AUC.
# Downstream value (assumed, to be *measured* by a training run): each selected relevant page gives `+g`, each irrelevant `-h` (dilution).
# Net gain `G(t) = M [pi TPR g - (1-pi) FPR h]`.  `G'(t) = 0` where the posterior `P(rel | s=t) = h/(g+h)`: a posterior cutoff, so `t* = d'/2 + ln[(1-pi) h /(pi g)]/d'`.
# %%
rng = np.random.default_rng(0)
dprime, pi, M = 2.5, 0.01, 4_000_000
rel = rng.random(M) < pi
s = rng.standard_normal(M) + dprime * rel
t = 2.0
sel = s >= t
prec_mc = rel[sel].mean(); rec_mc = sel[rel].mean()
tpr, fpr = Phibar(t - dprime), Phibar(t)
prec_th = pi * tpr / (pi * tpr + (1 - pi) * fpr)
check("1. precision closed form vs Monte Carlo (4e6 pages)", abs(prec_mc - prec_th), 4 * math.sqrt(prec_th * (1 - prec_th) / sel.sum()), f"prec {prec_th:.4f}, recall {tpr:.4f}")
check("1. recall closed form vs Monte Carlo", abs(rec_mc - tpr), 4 * math.sqrt(tpr * (1 - tpr) / rel.sum()))
auc = stats.norm.cdf(dprime / math.sqrt(2))
print(f"  AUC={auc:.4f} (d'={dprime}); at a threshold with recall {tpr:.2f}: precision {prec_th:.3f} for pi={pi}; for pi=0.2 the same classifier gives {0.2*tpr/(0.2*tpr+0.8*fpr):.3f}")
check("1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4", 0.0 if (0.2 * tpr / (0.2 * tpr + 0.8 * fpr)) - prec_th > 0.4 else 1.0, 0.0)
g_, h_ = 1.0, 1.0
ts = np.linspace(0, 6, 60001)
G = pi * Phibar(ts - dprime) * g_ - (1 - pi) * Phibar(ts) * h_
t_star = dprime / 2 + math.log((1 - pi) * h_ / (pi * g_)) / dprime
check("1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold", abs(ts[G.argmax()] - t_star), 2e-4, f"t*={t_star:.4f}")
post = lambda t: pi * stats.norm.pdf(t - dprime) / (pi * stats.norm.pdf(t - dprime) + (1 - pi) * stats.norm.pdf(t))
check("1. posterior P(rel | s=t*) = h/(g+h)", abs(post(t_star) - h_ / (g_ + h_)), 1e-12)
def f1(t):
    tp = pi * Phibar(t - dprime); fp = (1 - pi) * Phibar(t); fn = pi - tp
    return 2 * tp / (2 * tp + fp + fn)
t_f1 = ts[np.array([f1(x) for x in ts]).argmax()]
Gof = lambda t: pi * Phibar(t - dprime) * g_ - (1 - pi) * Phibar(t) * h_
print(f"  max-F1 threshold {t_f1:.3f} vs max-gain threshold {t_star:.3f};  gain/M at F1-opt {Gof(t_f1):.3e} vs gain-opt {Gof(t_star):.3e}")
check("1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain", 0.0 if (abs(t_f1 - t_star) > 0.2 and Gof(t_f1) < Gof(t_star)) else 1.0, 0.0)
# seed coverage and iteration: positives split into 'covered' (d_in) and 'uncovered' (d_out) clusters; each round adds a fraction gamma of uncovered clusters to the seed
d_in, d_out, gamma, tt = 3.0, 0.5, 0.5, 2.0
rec = []; cov = 0.3
for r in range(1, 6):
    rec.append(cov * Phibar(tt - d_in) + (1 - cov) * Phibar(tt - d_out)); cov = cov + gamma * (1 - cov)
inc = np.diff(rec)
print("  recall by round:", np.round(rec, 4), " increments:", np.round(inc, 4), " ratio of successive increments:", np.round(inc[1:] / inc[:-1], 3))
check("1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma)", max(abs(inc[1:] / inc[:-1] - (1 - gamma))), 1e-12)
# %% [markdown]
# ## 2. Pass-rate band selection is importance weighting of the policy gradient
# Bernoulli problem with logit parameter, `n` samples, `c` successes, `p_hat = c/n`.  GRPO-style mean baseline gives per-problem gradient magnitude `ghat = p_hat (1 - p_hat)` (note 0206), `E ghat = (1 - 1/n) p(1-p)`.
# Keeping the problem iff `p_hat` lies in a band `[lo, hi]` that is computed from the SAME samples multiplies the gradient by `1{p_hat in band}`:
# effective weight on a problem of true pass rate p:  `w_same(p) = E[ghat 1{band}] / (p(1-p))`.  With an independent pilot estimate `p_hat'` instead: `w_ind(p) = (1 - 1/n) P[p_hat' in band]`.
# %%
def bin_pmf(n, p): return np.array([comb(n, c) * p ** c * (1 - p) ** (n - c) for c in range(n + 1)])
n = 8; lo, hi = 1 / 8, 6 / 8
cs = np.arange(n + 1); ph = cs / n; inb = ((ph >= lo) & (ph <= hi)).astype(float); ghat = ph * (1 - ph)
rows = []
for p in (0.02, 0.1, 0.3, 0.5, 0.7, 0.9, 0.98):
    pm = bin_pmf(n, p)
    w_same = (pm * ghat * inb).sum() / (p * (1 - p)); w_ind = (1 - 1 / n) * (pm * inb).sum()
    check_unb = abs((pm * ghat).sum() - (1 - 1 / n) * p * (1 - p))
    assert check_unb < 1e-15
    rows.append((p, w_same, w_ind)); print(f"  p={p:4.2f}:  w_same={w_same:.4f}   w_ind={w_ind:.4f}   ratio={w_same/w_ind:.3f}")
check("2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)", 0.0, 0.0)
check("2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the statistic that defines the gradient)", 0.0 if all(w1 > w2 for _, w1, w2 in rows) else 1.0, 0.0)
# MC check of w_same by simulating samples
p = 0.1; cs_mc = rng.binomial(n, p, size=2_000_000); phm = cs_mc / n
est = (phm * (1 - phm) * ((phm >= lo) & (phm <= hi))).mean() / (p * (1 - p)); th = [r for r in rows if r[0] == p][0][1]
check("2. w_same(p=0.1) exact vs 2e6 Monte Carlo", abs(est / th - 1), 0.01)
# which problems does the band actually admit?  contamination of the band by noise in p_hat (pool p ~ Beta(0.6,1.4), n = 8)
a_, b_ = 0.6, 1.4; band_p = (0.125, 0.75)
px = rng.beta(a_, b_, size=1_000_000); cx = rng.binomial(n, px); sel_ = (cx / n >= lo) & (cx / n <= hi)
frac_out = ((px[sel_] < band_p[0] - 1e-12) | (px[sel_] > band_p[1] + 1e-12)).mean()
# exact via beta-binomial
pmf_bb = np.array([comb(n, c) * special.beta(c + a_, n - c + b_) / special.beta(a_, b_) for c in range(n + 1)])
print(f"  fraction selected: {sel_.mean():.4f} (exact {pmf_bb[(cs>=1)&(cs<=6)].sum():.4f});  of the selected, true p outside [{band_p[0]},{band_p[1]}]: {frac_out:.3f}")
check("2. beta-binomial exact P[band] vs MC", abs(sel_.mean() - pmf_bb[(cs >= 1) & (cs <= 6)].sum()), 4 * math.sqrt(0.25 / 1e6))
check("2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)", 0.0 if frac_out > 0.15 else 1.0, 0.0, f"{frac_out:.3f}")
# %% [markdown]
# ## 3. Regression to the mean: the same samples cannot select and evaluate
# Pool `p ~ Beta(a, b)`, `c | p ~ Bin(n, p)`.  Posterior `p | c ~ Beta(a + c, b + n - c)`, mean `(a + c)/(a + b + n)`.
# Select the "unsolved" problems `c = 0`: observed rate 0, true mean `a/(a+b+n)` -- re-sampling them WITHOUT any training shows an apparent gain.
# %%
a_, b_, n = 0.6, 1.4, 8
post_mean = lambda c: (a_ + c) / (a_ + b_ + n)
px = rng.beta(a_, b_, size=3_000_000); c1 = rng.binomial(n, px); c2 = rng.binomial(n, px)       # c2: fresh samples, NO training between
for cc in (0, 1, 2):
    selc = c1 == cc
    emp = c2[selc].mean() / n; th = post_mean(cc)
    check(f"3. E[fresh pass rate | first-pass count c={cc}] = (a+c)/(a+b+n)", abs(emp - th), 4 * math.sqrt(th * (1 - th) / (selc.sum() * n)), f"observed first-pass rate {cc/n:.3f} -> fresh {emp:.4f} (theory {th:.4f}); apparent 'gain' {emp - cc/n:+.4f} with no training")
# difference-in-differences with an untrained control from the same selection
tau = 0.03; sel0 = c1 == 0
idx = np.flatnonzero(sel0); half = rng.random(idx.size) < 0.5           # random split of the selected problems into trained / untrained
p_after = np.where(half, np.minimum(px[idx] + tau * (1 - px[idx]), 1.0), px[idx])      # training raises p by tau (1-p) on the trained half only
c_after = rng.binomial(n, p_after) / n
naive = c_after[half].mean() - 0.0; control = c_after[~half].mean() - 0.0
did = naive - control
tau_eff = (tau * (1 - px[idx[half]])).mean()
se = math.sqrt(c_after[half].var() / half.sum() / n + c_after[~half].var() / (~half).sum() / n) + 1e-9
print(f"  trained-half before/after = {naive:+.4f} (true effect {tau_eff:+.4f}); untrained control 'gain' = {control:+.4f}; DiD = {did:+.4f}")
check("3. naive before/after overstates the effect (regression to the mean) and DiD recovers it", abs(did - tau_eff), 0.003, f"naive {naive:+.4f}, DiD {did:+.4f}, true {tau_eff:+.4f}")
check("3. naive before/after is inflated by RTM: excess over truth exceeds 0.03", 0.0 if (naive - tau_eff) > 0.03 else 1.0, 0.0)
# split-sample selection: select with set 1, evaluate with independent set 2 -> E[fresh | selected] is what an untrained re-evaluation gives (no gain)
# %% [markdown]
# ## 4. MinHash: unbiased Jaccard with variance J(1-J)/k; LSH banding; containment is a different quantity
# `Jhat = (1/k) sum_i 1{minhash_i(A) = minhash_i(B)}`.  For a uniformly random permutation `P[min equal] = J`, so `k Jhat ~ Bin(k, J)`: unbiased, `Var = J(1-J)/k`.
# %%
def jaccard(A, B): return len(A & B) / len(A | B)
U = 3000; rng = np.random.default_rng(7)
A = set(rng.choice(U, 400, replace=False).tolist()); share = 240
B = set(list(A)[:share]) | set(rng.choice([u for u in range(U) if u not in A], 160, replace=False).tolist())
J = jaccard(A, B)
def minhash_perm(S, perms): return np.array([min(perm[x] for x in S) for perm in perms])
k = 128; reps = 3000
perm_bank = np.array([rng.permutation(U) for _ in range(k)])           # one fresh set of k permutations per replication below
ests = []
Al = np.array(sorted(A)); Bl = np.array(sorted(B))
for r in range(reps):
    perms = rng.permuted(np.tile(np.arange(U), (k, 1)), axis=1)
    ests.append((perms[:, Al].min(1) == perms[:, Bl].min(1)).mean())
ests = np.array(ests)
check("4. E[Jhat] = J (random permutations, k=128, 3000 replications)", abs(ests.mean() - J), 4 * math.sqrt(J * (1 - J) / k / reps), f"J={J:.4f}")
check("4. Var(Jhat) = J(1-J)/k", abs(ests.var(ddof=1) / (J * (1 - J) / k) - 1), 4 * math.sqrt(2 / (reps - 1)), f"emp {ests.var(ddof=1):.3e} vs {J*(1-J)/k:.3e}")
# practical universal hash h(x) = (a x + b) mod p  (not truly random permutations)
P = 2_147_483_647
def uh_est(A_, B_, k, rg):
    a = rg.integers(1, P, k); b = rg.integers(0, P, k)
    Aa = np.array(sorted(A_), dtype=np.int64); Bb = np.array(sorted(B_), dtype=np.int64)
    ha = ((a[:, None] * Aa[None, :] + b[:, None]) % P).min(1); hb = ((a[:, None] * Bb[None, :] + b[:, None]) % P).min(1)
    return (ha == hb).mean()
rg = np.random.default_rng(11)
est_uh = np.array([uh_est(A, B, 128, rg) for _ in range(3000)])
# a linear hash is pairwise independent but NOT min-wise independent: a small bias is allowed and reported, not hidden
check("4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)", abs(est_uh.mean() - J), 0.01, f"bias {est_uh.mean()-J:+.4f} (MC s.e. {math.sqrt(J*(1-J)/128/3000):.4f})")
for eps in (0.05, 0.02, 0.01):
    kreq = math.ceil(stats.norm.ppf(0.975) ** 2 * 0.25 / eps ** 2); print(f"  worst-case (J=1/2) hashes for +-{eps} at 95%: k >= {kreq}")
kk = 385; cover = 0.0
pm = stats.binom.pmf(np.arange(kk + 1), kk, 0.5); jh = np.arange(kk + 1) / kk
cover = pm[np.abs(jh - 0.5) <= 0.05].sum()
check("4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)", 0.0 if cover > 0.94 else 1.0, 0.0, f"coverage {cover:.4f}")
# decision at a decontamination threshold: flag if Jhat >= 0.8, k = 128
kk = 128
for Jt in (0.7, 0.75, 0.8, 0.85):
    pflag = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)
    print(f"  true J={Jt}: P[flag at Jhat>=0.8] = {pflag:.3f}   (sd of Jhat {math.sqrt(Jt*(1-Jt)/kk):.3f})")
Jt = 0.75; mc = (rng.binomial(kk, Jt, 2_000_000) / kk >= 0.8).mean(); ex = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)
check("4. exact flag probability at J=0.75 vs Monte Carlo", abs(mc - ex), 4 * math.sqrt(ex * (1 - ex) / 2e6), f"P[flag]={ex:.3f}: a 'near-duplicate' just under the threshold is flagged often")
# LSH banding: b bands of r rows: P[candidate] = 1 - (1 - J^r)^b
b_, r_ = 20, 5
for Jt in (0.3, 0.5, 0.6, 0.8):
    th = 1 - (1 - Jt ** r_) ** b_
    bits = rng.random((400_000, b_, r_)) < Jt       # each hash agrees w.p. J independently (random permutations)
    mc = bits.all(2).any(1).mean()
    check(f"4. LSH banding P[candidate] at J={Jt}: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)", abs(mc - th), 4 * math.sqrt(th * (1 - th) / 400_000) + 1e-4, f"{th:.4f}")
print(f"  LSH threshold approx (1/b)^(1/r) = {(1/b_)**(1/r_):.3f}")
# containment: a benchmark item (40 shingles) inside a long web page (2000 shingles): J tiny, containment 1
Aitem = set(range(40)); Bpage = set(range(2000))
Jc = jaccard(Aitem, Bpage); cont = len(Aitem & Bpage) / len(Aitem)
inter_from_J = Jc * (len(Aitem) + len(Bpage)) / (1 + Jc)
check("4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes", abs(inter_from_J - 40), 1e-9, f"J={Jc:.4f} (flag threshold 0.8 would miss it), containment={cont:.2f}")
# error amplification of the containment estimate via Jhat
kk = 128; Jh = rng.binomial(kk, Jc, 100_000) / kk
cont_hat = Jh * (40 + 2000) / (1 + Jh) / 40
print(f"  containment estimate from MinHash at k=128: mean {cont_hat.mean():.3f}, sd {cont_hat.std():.3f} (true 1.0): very noisy for tiny J")
check("4. containment from MinHash at k=128 for tiny J has relative sd > 0.5", 0.0 if cont_hat.std() > 0.5 else 1.0, 0.0, f"sd {cont_hat.std():.2f}")
# real text shingles: document with edits
rg = np.random.default_rng(5)
vocab = [f"w{i}" for i in range(3000)]
doc = list(rg.choice(vocab, 600)); doc2 = doc.copy()
for pos in rg.choice(600, 36, replace=False): doc2[pos] = rg.choice(vocab)
sh = lambda d, w=5: {" ".join(d[i:i + w]) for i in range(len(d) - w + 1)}
S1, S2 = sh(doc), sh(doc2); Jx = jaccard(S1, S2)
hs = lambda S: {int.from_bytes(__import__('hashlib').blake2b(x.encode(), digest_size=8).digest(), 'little') % P for x in S}
H1, H2 = hs(S1), hs(S2)
est = np.array([uh_est(H1, H2, 256, rg) for _ in range(400)])
check("4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard", abs(est.mean() - Jx), 4 * math.sqrt(Jx * (1 - Jx) / 256 / 400) + 0.003, f"exact J={Jx:.4f}, est {est.mean():.4f}, sd {est.std():.4f}")
# %% [markdown]
# ## 5. Membership inference from NLL gaps
# Members have lower NLL: `NLL_member ~ N(mu - Delta, sigma^2)`, non-member `N(mu, sigma^2)`.  AUC of the threshold test = `Phi(Delta / (sigma sqrt 2))`; TPR at FPR `f`: `Phi(Delta/sigma - z_{1-f})`.
# If examples differ in intrinsic difficulty `D_i ~ N(0, sD^2)` (added to every NLL), the raw-NLL AUC drops to `Phi(Delta/sqrt(2(sD^2+sigma^2)))`; a reference-calibrated score `NLL_target - NLL_ref` cancels `D_i`.
# %%
rng = np.random.default_rng(21); R_ = 400_000
Delta, sig, sD, sref = 0.30, 0.25, 0.80, 0.25
mem = rng.normal(-Delta, sig, R_); non = rng.normal(0, sig, R_)
def auc_of(pos_low, neg):
    # P[score(member) < score(non-member)] with MC estimate by ranks on a subsample
    x = pos_low[:200000]; y = neg[:200000]
    r = stats.rankdata(np.concatenate([x, y])); return 1 - (r[:len(x)].sum() - len(x) * (len(x) + 1) / 2) / (len(x) * len(y))
check("5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)", abs(auc_of(mem, non) - stats.norm.cdf(Delta / (sig * math.sqrt(2)))), 0.003, f"{stats.norm.cdf(Delta/(sig*math.sqrt(2))):.4f}")
thr = np.quantile(non, 0.01); tpr_mc = (mem <= thr).mean(); tpr_th = stats.norm.cdf(Delta / sig - stats.norm.ppf(0.99))
check("5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})", abs(tpr_mc - tpr_th), 0.004, f"{tpr_th:.4f}")
Dm = rng.normal(0, sD, R_); Dn = rng.normal(0, sD, R_)
raw_m = Dm + mem; raw_n = Dn + non
auc_raw_th = stats.norm.cdf(Delta / math.sqrt(2 * (sD ** 2 + sig ** 2)))
check("5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))", abs(auc_of(raw_m, raw_n) - auc_raw_th), 0.003, f"{auc_raw_th:.4f}")
ref_m = Dm + rng.normal(0, sref, R_); ref_n = Dn + rng.normal(0, sref, R_)
cal_th = stats.norm.cdf(Delta / math.sqrt(2 * (sig ** 2 + sref ** 2)))
check("5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))", abs(auc_of(raw_m - ref_m, raw_n - ref_n) - cal_th), 0.003, f"{cal_th:.4f}")
print(f"  AUC: raw {auc_raw_th:.3f} -> calibrated {cal_th:.3f}  (Delta={Delta}, sigma={sig}, sD={sD}, sref={sref})")
# a real (tiny) language model: character n-gram, memorisation grows with order; ground-truth membership known
V = 12; L = 120; ndoc = 2000; alpha = 0.1; rng = np.random.default_rng(33)
logits = rng.normal(0, 1.5, (V, V, V))                                   # order-2 Markov source
def gen_docs(nd, rg):
    taus = np.exp(rg.uniform(math.log(0.4), math.log(2.5), nd)); X = np.zeros((nd, L), dtype=np.int64); X[:, :2] = rg.integers(0, V, (nd, 2))
    for t in range(2, L):
        lg = logits[X[:, t - 2], X[:, t - 1]] / taus[:, None]; pr = np.exp(lg - lg.max(1, keepdims=True)); pr /= pr.sum(1, keepdims=True)
        u = rg.random(nd)[:, None]; X[:, t] = (pr.cumsum(1) < u).sum(1).clip(max=V - 1)
    return X, taus
Xm, tm = gen_docs(ndoc, rng); Xn, tn = gen_docs(ndoc, rng); Xr, _ = gen_docs(ndoc, rng)
def fit(X, order):
    cnt = np.zeros((V ** order, V))
    for t in range(4, L):
        ctx = np.zeros(len(X), dtype=np.int64)
        for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]
        np.add.at(cnt, (ctx, X[:, t]), 1)
    return cnt
def nll(cnt, X, order):
    out = np.zeros(len(X))
    for t in range(4, L):
        ctx = np.zeros(len(X), dtype=np.int64)
        for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]
        c = cnt[ctx]; out -= np.log((c[np.arange(len(X)), X[:, t]] + alpha) / (c.sum(1) + alpha * V))
    return out / (L - 4)
def auc_small(low, high):  # P[low < high] (+0.5 ties)
    r = stats.rankdata(np.concatenate([low, high])); n1 = len(low); return 1 - (r[:n1].sum() - n1 * (n1 + 1) / 2) / (n1 * len(high))
res = {}
for order in (1, 2, 3, 4):
    cm = fit(Xm, order); cr = fit(Xr, order)
    nm, nn = nll(cm, Xm, order), nll(cm, Xn, order)
    rm, rn = nll(cr, Xm, order), nll(cr, Xn, order)
    res[order] = (auc_small(nm, nn), auc_small(nm - rm, nn - rn), nn.mean() - nm.mean())
    print(f"  order {order}: raw-NLL AUC = {res[order][0]:.3f}, calibrated AUC = {res[order][1]:.3f}, mean NLL gap (non - member) = {res[order][2]:.4f} nats/char")
check("5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5", max(abs(res[1][0] - 0.5), abs(res[2][0] - 0.5)), 0.04)
check("5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2", 0.0 if res[4][0] > res[3][0] > res[2][0] else 1.0, 0.0)
check("5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)", 0.0 if res[4][1] > res[4][0] + 0.05 else 1.0, 0.0, f"{res[4][0]:.3f} -> {res[4][1]:.3f}")
check("5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small order-2 overfit that raw NLL hides)", 0.0 if (res[1][1] < res[2][1] < res[3][1] < res[4][1] and all(res[o][1] > res[o][0] for o in (2, 3, 4))) else 1.0, 0.0)
# %% [markdown]
# ## Summary
# %%
for n_, e_, t_ in CHECKS: print(f"PASS {n_[:110]}")
print(f"\n{len(CHECKS)} checks, all passed")