qualify.py

"""
QUALIFYING A MEMBER: CONCEDE, OR OVERRIDE.

A confidence member asserts that the base is out of its depth, and the base
CANNOT CHECK — it has no representation of unfamiliarity, so there is
nothing to audit the claim against. That is why the member is needed and
also why it cannot simply be believed. Lee's protocol: build the member,
TEST it against the base's confidence, and let it concede or override on
the result rather than on its author's word.

The test splits into two tiers, and they differ in who can run them.

  NON-INTERFERENCE is verifiable by the base owner ALONE. Run the base's
  own evaluation with the member active and check that nothing it already
  did got worse. No data from the member's domain is needed. For a
  confidence member this is structural rather than tested — a positive
  scale cannot move the argmax — and for a capability member it is what
  the exact-zero isolation already guarantees, provided the member has its
  own output.

  COMPETENCE is NOT verifiable that way. To confirm a member is RIGHT
  about the territory it claims you need labelled examples of that
  territory — and the case where a member is most useful is exactly the
  case where you have none. So competence is a trust decision with a name
  attached, not a check.

WHICH GIVES A TWO-TIER ADMISSION. A member is ADMITTED on
non-interference, which anyone can verify and which the failsafe makes
cheap: an unqualified member costs nothing because dropping to the base is
exact and instant. It is PROMOTED TO OVERRIDE only on domain evidence
somebody supplies and vouches for.

This measures where the boundary actually falls. The jurisdiction score is
swept, and at each threshold three things are read: how much territory the
member claims, whether it beats the base THERE, and whether the base's own
work outside the claim is untouched. The qualified threshold is the widest
claim that passes both.

AND IT CHECKS SOMETHING NEITHER MEMBER CAN CHECK ALONE: whether the region
the CONFIDENCE member claims is the region the CAPABILITY member can
actually serve. Two independently fitted signals about the same boundary,
and they need not agree.
"""

import numpy as np
import time
import json

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def windowed(g, c_in, k, c_out):
    ni, no = c_in*g*g, c_out*g*g
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g*g), ii % (g*g)
    co, po = jj // (g*g), jj % (g*g)
    dr = pi // g - (po // g - k//2)
    dc = pi % g - (po % g - k//2)
    inside = (dr >= 0) & (dr < k) & (dc >= 0) & (dc < k)
    K = c_in*c_out*k*k + 1
    idx = np.where(inside, (ci*c_out + co)*k*k + dr*k + dc, K-1)
    return idx.ravel().astype(np.int32), K, no


_FIXED = {}


class FixedScatter:
    def __init__(self, idx, K, cap=8192):
        h = to_host(idx).astype(np.int64).reshape(-1)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        big = np.where(counts > cap)[0]
        small = np.where(counts <= cap)[0]
        self.K = K
        self.order = to_dev(order, np.int64) if _GPU else order
        self.big = [(int(b), int(starts[b]), int(starts[b]+counts[b]))
                    for b in big]
        self.small = to_dev(small, np.int64) if _GPU else small
        self.width = int(counts[small].max()) if len(small) else 0
        if self.width:
            pos = np.concatenate([np.arange(counts[s]) for s in small])
            src = np.concatenate([np.arange(starts[s], starts[s]+counts[s])
                                  for s in small])
            row = np.repeat(np.arange(len(small)), counts[small])
            self.src = to_dev(src, np.int64) if _GPU else src
            sl = row*self.width + pos
            self.slot = to_dev(sl, np.int64) if _GPU else sl
            self.buf = xp.zeros(len(small)*self.width, DT)
        self._keep = idx

    def __call__(self, g):
        gs = g.reshape(-1)[self.order]
        out = xp.zeros(self.K, DT)
        if self.width:
            self.buf[:] = 0
            self.buf[self.slot] = gs[self.src]
            out[self.small] = self.buf.reshape(-1, self.width).sum(1)
        for b, a, z in self.big:
            out[b] = gs[a:z].sum()
        return out


def scatter(dW, idx, K):
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


def train_base(Xtr, Ytr, cfg, seed):
    D, g, ch = Xtr.shape[1], cfg["grid"], cfg["chan"]
    rg = np.random.default_rng(seed)
    layers, cin = [], cfg["c_in"]
    for l in range(cfg["depth"]):
        idx, K, no = windowed(g, cin, 3, ch)
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, out=no, taps=cin*9,
                           ins=D if l == 0 else layers[-1]["out"]))
        cin = ch
    L = cfg["depth"]
    P = []
    for l in layers:
        v = rg.normal(0, np.sqrt(2.0/l["taps"]), l["K"]).astype(np.float32)
        v[-1] = 0.0
        P.append(to_dev(v))
    P += [xp.ones(l["out"], DT) for l in layers]
    P += [xp.zeros(l["out"], DT) for l in layers]
    P += [to_dev(rg.normal(0, np.sqrt(2.0/layers[-1]["out"]),
                           (layers[-1]["out"], 10))), xp.zeros(10, DT)]
    HEAD, OB = 3*L, 3*L+1
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = Xtr.shape[0]; t = 0
    ag = np.random.default_rng(seed + 991)

    def fwd(x):
        cache = []; h = x
        for li, l in enumerate(layers):
            W = P[li][l["idx"]].reshape(l["ins"], l["out"])
            z = h @ W
            var = z.var(1, keepdims=True) + 1e-5
            zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
            zs = zn*P[L+li] + P[2*L+li]
            a = xp.maximum(zs, 0)
            cache.append((h, W, var, zn, zs))
            h = a
        return h, cache

    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]; y = Ytr[b]
            h, cache = fwd(x)
            lg = h @ P[HEAD] + P[OB]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - y)/len(b)
            G = [xp.zeros_like(p) for p in P]
            G[HEAD] = h.T @ d; G[OB] = d.sum(0)
            dh = d @ P[HEAD].T
            for li in range(L-1, -1, -1):
                hin, W, var, zn, zs = cache[li]
                dzs = dh*(zs > 0)
                G[L+li] = (dzs*zn).sum(0); G[2*L+li] = dzs.sum(0)
                dzn = dzs*P[L+li]
                dz = (dzn - dzn.mean(1, keepdims=True)
                      - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var)
                G[li] = scatter(hin.T @ dz, layers[li]["idx"], layers[li]["K"])
                if li > 0:
                    dh = dz @ W.T
            t += 1
            for i, (p_, gr) in enumerate(zip(P, G)):
                M[i] = 0.9*M[i] + 0.1*gr
                V[i] = 0.999*V[i] + 0.001*gr*gr
                P[i] = p_ - cfg["lr"]*(M[i]/(1-0.9**t)) \
                    / (xp.sqrt(V[i]/(1-0.999**t))+1e-8)
    return P, fwd, HEAD, OB, layers[-1]["out"]


def ridge(A, R, lam):
    """The closed-form member: one solve, no steps.

    Fitted on the RESIDUAL between the target and what the base already
    says, so a member that explains nothing contributes nothing — and as
    lambda dominates the solution shrinks toward zero and the base answers.
    That is the failsafe as arithmetic rather than as a rule."""
    n, d = A.shape
    A1 = xp.concatenate([A, xp.ones((n, 1), DT)], 1)
    if d + 1 <= n:
        G = A1.T @ A1 + lam*xp.eye(d+1, dtype=DT)
        W = xp.linalg.solve(G, A1.T @ R)
    else:
        # the wide case: solve in the sample space instead, which is the
        # only tractable form when there are five examples and 3,136
        # features
        G = A1 @ A1.T + lam*xp.eye(n, dtype=DT)
        W = A1.T @ xp.linalg.solve(G, R)
    return W[:-1], W[-1]


def descent(A, R, lam, steps, lr):
    """The same fit by gradient descent, for comparison. The closed form is
    only worth having if it matches."""
    n, d = A.shape
    W = xp.zeros((d, R.shape[1]), DT); b = xp.zeros(R.shape[1], DT)
    M = [xp.zeros_like(W), xp.zeros_like(b)]
    V = [xp.zeros_like(W), xp.zeros_like(b)]
    for t in range(1, steps+1):
        E = A @ W + b - R
        G = [A.T @ E/n + lam*W/n, E.mean(0)]
        for i, (p_, gr) in enumerate(zip([W, b], G)):
            M[i] = 0.9*M[i] + 0.1*gr
            V[i] = 0.999*V[i] + 0.001*gr*gr
            upd = p_ - lr*(M[i]/(1-0.9**t))/(xp.sqrt(V[i]/(1-0.999**t))+1e-8)
            if i == 0:
                W = upd
            else:
                b = upd
    return W, b


def load(cfg):
    from tensorflow import keras
    (a, b), (c, d) = keras.datasets.fashion_mnist.load_data()
    X = np.concatenate([a, c]).astype(np.float32)/255.0
    y = np.concatenate([b, d]).ravel().astype(np.int64)
    if cfg["grid"] != 28:
        s = 28//cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s).mean(axis=(2, 4))
    rg = np.random.default_rng(0); p = rg.permutation(len(X))
    tr, te = p[:cfg["n_train"]], p[cfg["n_train"]:cfg["n_train"]+10000]
    mu, sd = X[tr].mean(), X[tr].std()+1e-8
    f = lambda Z: ((Z-mu)/sd).reshape(len(Z), -1)
    Y = np.zeros((len(tr), 10), np.float32); Y[np.arange(len(tr)), y[tr]] = 1
    return f(X[tr]), Y, y[tr], f(X[te]), y[te]




# fit_on: "classes" fits a specialist on a class subset (measured: an
# oracle gate is worth only +0.0070, because the base was already good
# there). "errors" fits on the base's own mistakes, which is the correction
# loop and has 13% of the test set to work with rather than a 1.8-point
# margin on 40% of it.






CFG = dict(grid=14, c_in=1, chan=16, depth=3, n_train=20000, batch=128,
           lr=1e-3, epochs=30, seed=0,
           base_classes=(0, 1, 2, 3, 4, 5, 6),
           new_classes=(7, 8, 9),
           n_holdout=4000, N=200, lam=1.0, draws=5,
           # THE BASE SEED IS SWEPT TOO. A first version trained ONE base
           # and varied only the member's 200 examples, so every figure
           # rested on a single training run of the thing being extended.
           base_seeds=(0, 1, 2),
           claims=(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0),
           # THE MEMBER SHARES THE BASE'S OUTPUT HERE, so it has to
           # overcome the suppression the head learned: classes 7-9 were
           # always the wrong answer, so their logits were pushed down.
           # Measured separately, a member does NOTHING from w=0 to w=1
           # (identical to four decimals), then 0.0030 at 1.5, 0.0653 at
           # 2.0, 0.3249 at 3.0, 0.8723 at 5.0. A first version of this
           # test ran at w=1 and therefore tested a member that could not
           # speak — it never qualified anywhere, which said nothing.
           member_w=(1.5, 2.0, 3.0, 5.0, 8.0))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("QUALIFYING A MEMBER: CONCEDE, OR OVERRIDE")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:14s} = {v}")
    print(f"\n  a CONFIDENCE member claims territory; a CAPABILITY member")
    print(f"  answers on it. Neither can verify the other, and the base can")
    print(f"  verify neither — only that nothing it already did got worse.")
    print("=" * 78, flush=True)

    X, Y, y, Xte, yte = load(CFG)
    old_m = np.isin(y, CFG["base_classes"])
    tr_i = np.where(old_m)[0]
    ho_new = np.where(~old_m)[0][:CFG["n_holdout"]//2]
    ho_old = tr_i[-(CFG["n_holdout"]//2):]
    tr_i = tr_i[:-(CFG["n_holdout"]//2)]
    ho_i = np.concatenate([ho_old, ho_new])
    Xtr, Ytr = to_dev(X[tr_i]), to_dev(Y[tr_i])
    Xho = to_dev(X[ho_i]); yho = y[ho_i]
    Yho = to_dev(np.eye(10, dtype=np.float32)[yho])
    unseen_ho = np.isin(yho, CFG["new_classes"])
    Xte_d = to_dev(Xte)

    unseen = np.isin(yte, CFG["new_classes"])
    rows, agree, bases = {}, {}, []
    for bs in CFG["base_seeds"]:
        P, fwd, HEAD, OB, width = train_base(Xtr, Ytr, CFG, bs)
        fho, _ = fwd(Xho); fte, _ = fwd(Xte_d)
        base_te = to_host(fte @ P[HEAD] + P[OB])
        margin = float(to_host((fho @ P[HEAD] + P[OB]).std()))
        base_pred = base_te.argmax(1); base_ok = base_pred == yte
        bases.append(float(base_ok.mean()))
        print(f"  base seed {bs}: {float(base_ok[~unseen].mean()):.4f} on its "
              f"seven, {float(base_ok[unseen].mean()):.4f} on the three"
              f"   [{time.time()-t0:.0f}s]", flush=True)
        for dr in range(CFG["draws"]):
            rg = np.random.default_rng(2200 + 97*bs + dr)
            sub = rg.choice(len(ho_i), min(CFG["N"], len(ho_i)), replace=False)
            sd_ = to_dev(sub, np.int64) if _GPU else sub
            jt = to_dev(np.where(unseen_ho[sub], 1.0, -1.0)
                        .astype(np.float32)[:, None])
            Wj, bj = ridge(fho[sd_], jt, CFG["lam"])
            claim_score = to_host(fte @ Wj + bj).ravel()
            Wc, bc = ridge(fho[sd_], margin*Yho[sd_], CFG["lam"])
            delta = to_host(fte @ Wc + bc)
            for w in CFG["member_w"]:
                member_pred = (base_te + w*delta).argmax(1)
                member_ok = member_pred == yte
                wins = member_ok & ~base_ok
                r = np.argsort(np.argsort(claim_score))
                n1, n0 = wins.sum(), (~wins).sum()
                agree.setdefault(w, []).append(
                    float((r[wins].sum() - n1*(n1-1)/2)/(n1*n0))
                    if n1 and n0 else float("nan"))
                for c in CFG["claims"]:
                    keep = (np.ones(len(yte), bool) if c >= 1
                            else claim_score > np.quantile(claim_score, 1-c))
                    pred = np.where(keep, member_pred, base_pred)
                    ok = pred == yte
                    rows.setdefault((w, c), []).append(dict(
                        overall=float(ok.mean()),
                        on_claim=float(ok[keep].mean()) if keep.any() else np.nan,
                        base_on_claim=float(base_ok[keep].mean()) if keep.any()
                        else np.nan,
                        outside=float(ok[~keep].mean()) if (~keep).any() else np.nan,
                        base_outside=float(base_ok[~keep].mean()) if (~keep).any()
                        else np.nan,
                        purity=float(unseen[keep].mean()) if keep.any() else np.nan))
    print(f"\n  {len(bases)} bases, spanning {min(bases):.4f} to "
          f"{max(bases):.4f} overall")

    summ = {}
    for key, rs in rows.items():
        m = {k: float(np.nanmean([r[k] for r in rs])) for k in rs[0]}
        m["sd"] = float(np.std([r["overall"] for r in rs]))
        summ[key] = m
    base_all = float(np.mean(bases))
    for w in CFG["member_w"]:
        print(f"\n  member at w = {w}"
              f"   (it reaches {summ[(w, 1.0)]['on_claim']:.4f} applied "
              f"everywhere, against the base's {base_all:.4f})")
        print(f"  {'claims':>7s} {'overall':>8s} {'member here':>12s} "
              f"{'base here':>10s} {'COMPETENT':>10s} {'outside':>8s} "
              f"{'unseen share':>13s}")
        for c in CFG["claims"]:
            m = summ[(w, c)]
            comp = m["on_claim"] - m["base_on_claim"]
            ni = (abs(m["outside"] - m["base_outside"]) < 1e-9
                  if not np.isnan(m["outside"]) else True)
            print(f"  {c:6.0%} {m['overall']:8.4f} {m['on_claim']:12.4f} "
                  f"{m['base_on_claim']:10.4f} {comp:+10.4f} "
                  f"{'exact' if ni else 'MOVED':>8s} {m['purity']:12.1%}")
    json.dump({f"{k[0]}/{k[1]}": v for k, v in summ.items()},
              open("qualify.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    print(f"  NON-INTERFERENCE is exact at every threshold and every w by")
    print(f"  construction: outside the claim the base answers untouched,")
    print(f"  and dropping the member costs nothing. That tier ADMITS")
    print(f"  without evidence.\n")
    print(f"  {'w':>5s} {'widest qualified claim':>24s} {'competence':>11s} "
          f"{'overall':>9s} {'vs base':>9s} {'boundaries':>11s}")
    best = None
    for w in CFG["member_w"]:
        ok = [c for c in CFG["claims"]
              if summ[(w, c)]["on_claim"] > summ[(w, c)]["base_on_claim"]]
        a = float(np.nanmean(agree[w]))
        if ok:
            c = max(ok); m = summ[(w, c)]
            print(f"  {w:5.1f} {c:23.0%} "
                  f"{m['on_claim']-m['base_on_claim']:+11.4f} "
                  f"{m['overall']:9.4f} {m['overall']-base_all:+9.4f} "
                  f"{a:11.4f}")
            if best is None or m["overall"] > best[1]["overall"]:
                best = (w, m, c, a)
        else:
            print(f"  {w:5.1f} {'never':>23s} {'':>11s} {'':>9s} {'':>9s} "
                  f"{a:11.4f}")
    print()
    if best:
        w, m, c, a = best
        print(f"  QUALIFIED. At w = {w} the member beats the base on the")
        print(f"  widest {c:.0%} it claims ("
              f"{m['on_claim']-m['base_on_claim']:+.4f}) and lifts the whole")
        print(f"  task to {m['overall']:.4f} against {base_all:.4f} "
              f"({m['overall']-base_all:+.4f}).")
        print(f"  Beyond that claim it should CONCEDE, and the concession is")
        print(f"  free because the base is exact outside it.")
        print(f"\n  DO THE TWO BOUNDARIES AGREE? The confidence member's claim")
        print(f"  ranks where the capability member actually wins at AUC")
        print(f"  {a:.4f}.")
        if a > 0.75:
            print(f"  THEY DO. One member says 'mine' and the other can serve")
            print(f"  it, fitted independently on the same budget — the")
            print(f"  handover is coherent rather than two assertions.")
        elif a > 0.6:
            print(f"  PARTLY. The claim points at the right region without")
            print(f"  matching it, so an override takes on territory the")
            print(f"  capability member cannot serve.")
        else:
            print(f"  THEY DO NOT. Claimed and servable are different sets,")
            print(f"  so a confidence member cannot AUTHORISE a capability")
            print(f"  member — each qualifies on its own evidence.")
    else:
        print(f"  NEVER QUALIFIED AT ANY w. The member should concede")
        print(f"  everywhere and be admitted on non-interference alone.")
    print(f"\n  seed spread (worst) "
          f"{max(m2['sd'] for m2 in summ.values()):.4f}")
    print(f"  total {time.time()-t0:.0f}s; wrote qualify.json")


if __name__ == "__main__":
    main()