gate.py

"""
WHEN SHOULD A MEMBER BE APPLIED AT ALL?

A member fitted on four of ten classes reaches 0.9295 on its own from
twenty examples, against the base's 0.9071 — but it drags the other six
from 0.8440 to 0.7408, so applied to everything it is a NET LOSS. It is a
specialist, and specialists need to know when to speak.

An oracle that applied it only on its own classes would give
0.4 x 0.9295 + 0.6 x 0.8440 = 0.8782, against the base's 0.8690. So the
whole correction loop is worth about +0.009, and every point of it depends
on the gate.

FIVE ROUTING ATTEMPTS IN THIS PROGRAMME FAILED and one succeeded. The one
that worked — the cascade — worked because its payoff was ASYMMETRIC: a
wrong escalation cost arithmetic, not accuracy, so a mediocre signal was a
perfectly good throttle. This gate has the same shape. Applying a member
wrongly costs accuracy on one example; withholding it costs the correction
on one example. Neither is catastrophic, so a gate does not have to be good
— only better than always or never.

FOUR SIGNALS, all computable from what is already being calculated:

  BASE ARGMAX     apply when the base already predicts one of the member's
                  classes. Free, and uses no member information at all.
  DELTA SIZE      apply when the member has a strong opinion about this
                  example — |delta(a)| above a threshold.
  DELTA MARGIN    apply when the member's top-two gap is wide, which was
                  the best of the three confidence signals in the cascade.
  AGREEMENT       apply when base and member agree on the answer, which
                  makes the member a confirmer rather than an overruler.

and two continuous variants, because w is a dial and a gate need not be
binary: w scaled by the signal, and w chosen per example.

The oracle gate is the ceiling. The gap between it and the best rule is
what a better signal would be worth — the same reading the cascade's oracle
gave.
"""

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, member_classes=(0, 1, 2, 3),
           fit_on="errors", N=20, lam=1.0, draws=5,
           n_holdout=5000, global_w=(0.1, 0.25, 0.5))


def gated(base_lg, delta, keep, w=1.0):
    """Apply the member only where keep is true."""
    k = keep[:, None].astype(np.float32)
    return base_lg + (w*k)*delta


def report(name, lg, yte, own, base_all):
    pred = lg.argmax(1)
    acc = float((pred == yte).mean())
    return dict(name=name, acc=acc, gain=acc-base_all,
                own=float((pred[own] == yte[own]).mean()),
                rest=float((pred[~own] == yte[~own]).mean()))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("WHEN SHOULD A MEMBER BE APPLIED AT ALL?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:15s} = {v}")
    if CFG.get("fit_on") == "errors":
        print(f"\n  the member is fitted on {CFG['N']} examples the base "
              f"GETS WRONG")
    else:
        print(f"\n  the member is fitted on {CFG['N']} examples of classes "
              f"{CFG['member_classes']}")
    print(f"  applying it everywhere is a NET LOSS; the question is the gate")
    print("=" * 78, flush=True)

    Xtr0, Ytr0, ytr0, Xte, yte = load(CFG)
    nh = CFG.get("n_holdout", 0)
    if nh:
        Xho, yho = Xtr0[-nh:], ytr0[-nh:]
        Xtr, Ytr, ytr = Xtr0[:-nh], Ytr0[:-nh], ytr0[:-nh]
        Yho = np.eye(10, dtype=np.float32)[yho]
    else:
        Xtr, Ytr, ytr = Xtr0, Ytr0, ytr0
        Xho, yho, Yho = Xtr0[:1], ytr0[:1], np.eye(10, np.float32)[ytr0[:1]]
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    Xho_d, Yho_d = to_dev(Xho), to_dev(Yho)
    P, fwd, HEAD, OB, width = train_base(Xtr, Ytr, CFG, CFG["seed"])
    ftr, _ = fwd(Xtr); fte, _ = fwd(Xte); fho, _ = fwd(Xho_d)
    base_tr = ftr @ P[HEAD] + P[OB]
    base_hold = fho @ P[HEAD] + P[OB]
    base_te_d = fte @ P[HEAD] + P[OB]
    base_te = to_host(base_te_d)
    own = np.isin(yte, CFG["member_classes"])
    if CFG.get("fit_on") == "errors":
        own = base_te.argmax(1) != yte      # the region the member is for
    base_all = float((base_te.argmax(1) == yte).mean())
    margin = float(to_host(base_tr.std()))
    if CFG.get("fit_on") == "errors":
        # THE ERROR POOL MUST BE HELD OUT. The base memorises its training
        # set — 2.3% errors there against 13.1% on test — so training
        # errors are the weird few it could not memorise, not the ordinary
        # hard ones a member has to fix. The first run of this fitted on
        # those and the member only reached 32% of the errors.
        hp = to_host(base_hold).argmax(1)
        wrong = hp != yho
        pool_idx = np.where(wrong)[0]
        print(f"  the base errs on {wrong.mean():.1%} of the HELD-OUT split "
              f"({len(pool_idx):,} of {len(yho):,}), against "
              f"{float((to_host(base_tr).argmax(1) != ytr).mean()):.1%} of "
              f"what it trained on")
    else:
        pool_idx = np.where(np.isin(ytr, CFG["member_classes"]))[0]
    print(f"\n  base {base_all:.4f} overall, "
          f"{float((base_te.argmax(1)[own] == yte[own]).mean()):.4f} on its "
          f"four, {float((base_te.argmax(1)[~own] == yte[~own]).mean()):.4f} "
          f"on the rest   [{time.time()-t0:.0f}s]", flush=True)

    rows = {}
    for dr in range(CFG["draws"]):
        rg = np.random.default_rng(7000 + dr)
        sub = rg.choice(pool_idx, CFG["N"], replace=False)
        sd_ = to_dev(sub, np.int64) if _GPU else sub
        src_f, src_Y = ((fho, Yho_d) if CFG.get("fit_on") == "errors"
                        else (ftr, Ytr))
        W, b = ridge(src_f[sd_], margin*src_Y[sd_], CFG["lam"])
        delta = to_host(fte @ W + b)
        full = base_te + delta

        # the signals
        sig = {}
        if CFG.get("fit_on") == "errors":
            # the base cannot flag its own errors by naming a class, so the
            # cheap gate becomes its own confidence
            bc = np.exp(base_te - base_te.max(1, keepdims=True))
            bc = (bc/bc.sum(1, keepdims=True)).max(1)
            sig["base unsure"] = bc < np.quantile(bc, 0.30)
        else:
            sig["base argmax"] = np.isin(base_te.argmax(1),
                                         CFG["member_classes"])
        dn = np.linalg.norm(delta, axis=1)
        sig["delta size"] = dn > np.quantile(dn, 0.55)
        d2 = np.partition(delta, -2, axis=1)[:, -2:]
        dm = d2[:, 1] - d2[:, 0]
        sig["delta margin"] = dm > np.quantile(dm, 0.55)
        sig["agreement"] = base_te.argmax(1) == full.argmax(1)
        if CFG.get("fit_on") == "errors" and nh:
            # THE ORACLE CANNOT GATE — it reads the test labels. But on
            # HELD-OUT data the labels are known, so it can SUPERVISE a
            # detector: features -> "will the base be wrong here". Fitted
            # by the same closed-form solve, applied to unlabelled input.
            # Every gate before this was a heuristic; none was fitted to
            # predict error.
            tgt = to_dev((to_host(base_hold).argmax(1) != yho)
                         .astype(np.float32)[:, None])
            Wg, bg = ridge(fho, tgt - 0.5, CFG["lam"])
            score = to_host(fte @ Wg + bg).ravel()
            sig["LEARNED gate"] = score > np.quantile(score, 0.70)
            learned_soft = np.clip((score - score.min())
                                   / max(score.max()-score.min(), 1e-9), 0, 1)
        sig["ORACLE"] = own

        got = [report("never (base)", base_te, yte, own, base_all),
               report("always (w=1)", full, yte, own, base_all)]
        for nm, keep in sig.items():
            r = report(nm, gated(base_te, delta, keep), yte, own, base_all)
            r["fired"] = float(keep.mean())
            r["precision"] = float(own[keep].mean()) if keep.any() else 0.0
            r["recall"] = float(keep[own].mean())
            got.append(r)
        # soft: w proportional to the member's own margin, scaled to [0,1]
        soft = (dm - dm.min())/max(dm.max()-dm.min(), 1e-9)
        got.append(report("soft w by margin",
                          base_te + soft[:, None]*delta, yte, own, base_all))
        if CFG.get("fit_on") == "errors" and nh:
            got.append(report("LEARNED soft w",
                              base_te + learned_soft[:, None]*delta,
                              yte, own, base_all))
        for gw in CFG.get("global_w", ()):
            got.append(report(f"global w = {gw}", base_te + gw*delta,
                              yte, own, base_all))
        for r in got:
            rows.setdefault(r["name"], []).append(r)

    print(f"\n  {'gate':>17s} {'overall':>8s} {'vs base':>9s} {'region':>8s} "
          f"{'rest':>8s} {'fires':>7s} {'precision':>10s} {'recall':>8s}")
    summ = {}
    for nm, rs in rows.items():
        m = {k: float(np.mean([r[k] for r in rs]))
             for k in rs[0] if k != "name"}
        m["sd"] = float(np.std([r["acc"] for r in rs]))
        summ[nm] = m
        f = f"{m['fired']:6.1%}" if "fired" in m else ""
        p = f"{m['precision']:9.1%}" if "precision" in m else ""
        rc = f"{m['recall']:7.1%}" if "recall" in m else ""
        print(f"  {nm:>17s} {m['acc']:8.4f} {m['gain']:+9.4f} {m['own']:8.4f} "
              f"{m['rest']:8.4f} {f:>7s} {p:>10s} {rc:>8s}")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    orac = summ["ORACLE"]["gain"]; sd = max(m["sd"] for m in summ.values())
    cand = {k: v for k, v in summ.items()
            if k not in ("ORACLE", "never (base)", "always (w=1)")}
    best = max(cand.items(), key=lambda kv: kv[1]["gain"])
    print(f"  an ORACLE gate is worth {orac:+.4f} — that is the ceiling on")
    print(f"  the whole correction loop, and every point of it is the gate\n")
    print(f"  applying it always: {summ['always (w=1)']['gain']:+.4f}")
    print(f"  the best real gate: {best[0]} at {best[1]['gain']:+.4f}")
    print(f"  seed spread (worst) {sd:.4f}\n")
    if best[1]["gain"] > 2*sd:
        print(f"  THE GATE WORKS. {best[0]} captures "
              f"{best[1]['gain']/orac:.0%} of what an oracle would give,")
        print(f"  which turns a member from a net loss into a net gain. The")
        print(f"  cascade's lesson holds again: the payoff is asymmetric, so")
        print(f"  the signal does not have to be good.")
    elif orac > 2*sd:
        print(f"  NO GATE CAPTURES THE GAIN. The oracle says {orac:+.4f} is")
        print(f"  available and the best rule reaches {best[1]['gain']:+.4f},")
        print(f"  so this is the sixth routing result of the same shape —")
        print(f"  real complementarity, invisible signal.")
    else:
        print(f"  THERE IS NOTHING TO GATE. Even a perfect gate is worth only")
        print(f"  {orac:+.4f} against a spread of {sd:.4f}, so the member is")
        print(f"  not adding enough on its own classes to be worth applying")
        print(f"  selectively.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote gate.json")
    json.dump(summ, open("gate.json", "w"), indent=2)


if __name__ == "__main__":
    main()