ignorance.py

"""
INTRODUCING SOMETHING THE BASE CANNOT DO, AND CONTROLLING THE DOSE.

The correction loop failed for a reason worth stating precisely: a
converged base's remaining errors are mostly IRREDUCIBLE. Fitting a member
on held-out mistakes reached only 16.7% of them, and an oracle gate was
worth +0.0211 — the problem was not finding the correctable examples but
that there were few of them.

Ignorance is the opposite case. A base trained on classes 0-6 CANNOT
classify 7-9 at all: its head was never given a reason to raise those
logits, so it does not merely get them wrong, it never names them. The
member supplies the whole of a capability rather than a sliver of a
correction, so the ceiling is the entire thing rather than a few points.

That makes three questions answerable that the correction loop could not
reach.

  HOW MUCH CAPABILITY, AND AT WHAT DOSE? w scales the member at inference.
  At w = 0 the model is incapable by construction; somewhere above it is
  capable. The shape between is the dial, and where it saturates is how
  much of a member is actually needed.

  WHAT DOES IT COST THE REST? The base's own seven classes are what a user
  already had. If introducing three new ones damages them, the dose has a
  price and the curve says what it is.

  DOES THE MODEL KNOW IT IS IGNORANT? A base that has never seen a sandal
  still emits a confident answer about one. If its confidence cannot
  separate the classes it knows from the ones it does not, then no gate
  built on confidence can work — and that is a fact about ignorance rather
  than about this member.

The gates are carried over unchanged, including the LEARNED one supervised
on a held-out split, because the question of when to apply a member is the
same question. Only what the member carries has changed.
"""

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=5000, N=200, lam=1.0, draws=5,
           weights=(0.0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("INTRODUCING SOMETHING THE BASE CANNOT DO")
    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  the base is trained ONLY on classes {CFG['base_classes']}.")
    print(f"  classes {CFG['new_classes']} are not merely hard for it — its")
    print(f"  head was never given a reason to name them at all.")
    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_i = np.where(~old_m)[0][:CFG["n_holdout"]]
    Xtr, Ytr = to_dev(X[tr_i]), to_dev(Y[tr_i])
    Xho, yho = to_dev(X[ho_i]), y[ho_i]
    Yho = to_dev(np.eye(10, dtype=np.float32)[yho])
    Xte_d = to_dev(Xte)
    print(f"\n  base trains on {len(tr_i):,} examples of its seven classes;")
    print(f"  {len(ho_i):,} examples of the unseen three are held back for")
    print(f"  the member", flush=True)

    P, fwd, HEAD, OB, width = train_base(Xtr, Ytr, CFG, CFG["seed"])
    ftr, _ = fwd(Xtr); fho, _ = fwd(Xho); fte, _ = fwd(Xte_d)
    base_te = to_host(fte @ P[HEAD] + P[OB])
    margin = float(to_host((ftr @ P[HEAD] + P[OB]).std()))
    new = np.isin(yte, CFG["new_classes"])
    pred0 = base_te.argmax(1)
    print(f"\n  base: {float((pred0[~new] == yte[~new]).mean()):.4f} on its "
          f"own seven, {float((pred0[new] == yte[new]).mean()):.4f} on the "
          f"three it has never seen")
    print(f"  and it names one of the unseen three "
          f"{float(np.isin(pred0, CFG['new_classes']).mean()):.1%} of the time")

    # does the base know it is ignorant?
    bc = np.exp(base_te - base_te.max(1, keepdims=True))
    bc = (bc/bc.sum(1, keepdims=True)).max(1)
    print(f"\n  ITS CONFIDENCE ON WHAT IT KNOWS: {bc[~new].mean():.4f}")
    print(f"  ITS CONFIDENCE ON WHAT IT CANNOT DO: {bc[new].mean():.4f}")
    sep = (bc[~new].mean() - bc[new].mean())
    r = np.argsort(np.argsort(-bc)); n1 = new.sum(); n0 = (~new).sum()
    auc = float((r[new].sum() - n1*(n1-1)/2)/(n1*n0))
    print(f"  separation {sep:+.4f}, AUC {auc:.4f}   "
          f"(0.5 means it cannot tell at all)")

    print("\n" + "=" * 78)
    print("  THE DOSE")
    print("=" * 78)
    print(f"  {'w':>6s} {'its seven':>10s} {'the three':>10s} "
          f"{'ALL TEN':>9s} {'names new':>10s} {'confidence':>11s}")
    rows = {}
    for dr in range(CFG["draws"]):
        rg = np.random.default_rng(9000 + 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
        W, b = ridge(fho[sd_], margin*Yho[sd_], CFG["lam"])
        delta = to_host(fte @ W + b)
        for w in CFG["weights"]:
            lg = base_te + w*delta
            pr = lg.argmax(1)
            e = np.exp(lg - lg.max(1, keepdims=True))
            cf = (e/e.sum(1, keepdims=True)).max(1)
            rows.setdefault(w, []).append(dict(
                old=float((pr[~new] == yte[~new]).mean()),
                nw=float((pr[new] == yte[new]).mean()),
                all=float((pr == yte).mean()),
                names=float(np.isin(pr, CFG["new_classes"]).mean()),
                conf=float(cf.mean())))
    summ = {}
    for w, rs in rows.items():
        m = {k: float(np.mean([r[k] for r in rs])) for k in rs[0]}
        m["sd"] = float(np.std([r["all"] for r in rs]))
        summ[w] = m
        print(f"  {w:6.2f} {m['old']:10.4f} {m['nw']:10.4f} {m['all']:9.4f} "
              f"{m['names']:9.1%} {m['conf']:11.4f}")
    json.dump({str(k): v for k, v in summ.items()},
              open("ignorance.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    ws = sorted(summ)
    b0 = summ[0.0]
    best_all = max(ws, key=lambda w: summ[w]["all"])
    # where does the new capability saturate, and what has it cost by then
    top = max(summ[w]["nw"] for w in ws)
    sat = next(w for w in ws if summ[w]["nw"] >= top - 0.01)
    print(f"  the base cannot do the three at all: {b0['nw']:.4f} at w = 0\n")
    print(f"  the new capability saturates at w = {sat}: {summ[sat]['nw']:.4f}")
    print(f"  and by then its own seven have gone {b0['old']:.4f} -> "
          f"{summ[sat]['old']:.4f}  ({summ[sat]['old']-b0['old']:+.4f})")
    print(f"\n  best on all ten at w = {best_all}: {summ[best_all]['all']:.4f}")
    print(f"  against {b0['all']:.4f} at w = 0")
    print(f"  seed spread (worst) {max(m['sd'] for m in summ.values()):.4f}")
    print()
    if summ[sat]["old"] > b0["old"] - 0.02:
        print(f"  THE DOSE IS CONTROLLABLE AND CHEAP. A member introduces a")
        print(f"  capability the base did not have, saturating at w = {sat},")
        print(f"  and costs its existing classes "
              f"{abs(summ[sat]['old']-b0['old']):.4f}. Unlike the correction")
        print(f"  loop, there is a great deal to gain and the gain is not")
        print(f"  irreducible — the base was simply never told.")
    else:
        print(f"  THE DOSE IS A TRADE. Reaching the new capability costs")
        print(f"  {abs(summ[sat]['old']-b0['old']):.4f} of the base's own")
        print(f"  classes, so w is a place on a frontier rather than a free")
        print(f"  switch.")
    print()
    if auc < 0.6:
        print(f"  AND THE MODEL DOES NOT KNOW IT IS IGNORANT. Its confidence")
        print(f"  separates what it can do from what it cannot at AUC")
        print(f"  {auc:.4f} — barely better than chance. No gate built on")
        print(f"  confidence can work here, and that is a fact about")
        print(f"  ignorance rather than about this member: a model has no")
        print(f"  representation of a category it was never shown, so")
        print(f"  nothing internal marks it as unfamiliar.")
    else:
        print(f"  AND THE MODEL PARTLY KNOWS: confidence separates the seen")
        print(f"  from the unseen at AUC {auc:.4f}, so a gate has something")
        print(f"  to read after all.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote ignorance.json")


if __name__ == "__main__":
    main()