provisioning.py

"""
HOW MUCH RESERVE DOES A BODY NEED?

A member injected into a frozen body reached 96.7% of what the same member
achieves with a body built for it, and disturbed the existing members by
exactly nothing. But it failed badly in a smaller configuration -- 0.208
against a 0.775 ceiling -- and the difference was not the domain gap. Both
tests crossed a comparable gap; what differed was the BODY. A starved body
has no spare features for a new head to recruit, and a richer one does.

That turns injection into a provisioning question. If a family base is
built to fit exactly the members it has, a member arriving later finds
nothing to work with. If it is built with reserve, it does.

And there is a reason to think the reserve is invisible to anyone measuring
the family. A single folded body SATURATES: past roughly a thousand stored
values, more storage stops improving the members it serves. If injection
keeps improving past that point, then the reserve is capacity that does
nothing for the members present and everything for the ones not yet built
-- precisely a cap one would have to decide in advance and could not
discover by watching the family.

So this sweeps the body's storage and measures both at every point. The gap
between an injected member and its own ceiling, as a function of body
storage, is the provisioning curve. Where it closes is the cap.
"""

import numpy as np
import time
import json
from itertools import combinations

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)


_MAPS = {}


def tie(shape, K, tag):
    key = (tuple(shape), K, tag)
    if key not in _MAPS:
        import hashlib
        h = int(hashlib.md5(str(key).encode()).hexdigest()[:8], 16)
        rg = np.random.default_rng(h)
        idx = rg.integers(0, K, int(np.prod(shape)))
        _MAPS[key] = to_dev(idx, np.int64) if _GPU else idx.astype(np.int64)
    return _MAPS[key]


def fold(W, idx, K):
    flat = W.reshape(-1)
    s = xp.zeros(K, DT); c = xp.zeros(K, DT)
    if _GPU:
        import cupyx
        cupyx.scatter_add(s, idx, flat); cupyx.scatter_add(c, idx, xp.ones_like(flat))
    else:
        np.add.at(s, idx, flat); np.add.at(c, idx, np.ones_like(flat))
    return (s / xp.maximum(c, 1.0))[idx].reshape(W.shape)


def folded_init(rg, fan_in, shape, K):
    """Draw so the weights carry the intended scale AFTER folding: the mean
    of n samples has standard deviation sigma/sqrt(n), so projecting a
    fresh matrix shrinks it by the square root of the tying ratio."""
    s = np.sqrt(2.0/fan_in)*np.sqrt(max(1.0, np.prod(shape)/K))
    return rg.normal(0, s, shape)


class Hive:
    """A shared body, a permanent BASE head, and a delta head per member.

    The base head is what answers when no member is active, and a member is
    a correction on top of it rather than a replacement for it. That is the
    whole failsafe: a missing or unidentifiable member contributes zero."""

    def __init__(self, D, W, nc, S, K, seed=0, tag="inj"):
        rg = np.random.default_rng(seed)
        self.D, self.W, self.nc, self.S, self.K = D, W, nc, S, K
        self.i0 = tie((D, W), K, tag+"0")
        self.i1 = tie((W, W), K, tag+"1")
        self.W0 = to_dev(folded_init(rg, D, (D, W), K))
        self.W1 = to_dev(folded_init(rg, W, (W, W), K))
        self.base = to_dev(rg.normal(0, np.sqrt(2.0/W), (W, nc)))
        self.delta = xp.zeros((S, W, nc), DT)      # members start at ZERO
        self.b0 = xp.zeros(W, DT); self.b1 = xp.zeros(W, DT)
        self.bb = xp.zeros(nc, DT); self.bd = xp.zeros((S, nc), DT)
        self.project()

    def project(self):
        self.W0 = fold(self.W0, self.i0, self.K)
        self.W1 = fold(self.W1, self.i1, self.K)

    def body(self, X):
        z1 = X @ self.W0 + self.b0; a1 = xp.maximum(z1, 0)
        z2 = a1 @ self.W1 + self.b1; a2 = xp.maximum(z2, 0)
        return z1, a1, z2, a2

    def logits(self, a2, members=None):
        """base + delta. With members=None only the base answers, which is
        exactly what happens when nothing is injected."""
        lg = a2 @ self.base + self.bb
        if members is None:
            return lg[None]
        return lg[None] + xp.einsum('ni,sio->sno', a2,
                                    self.delta[members]) + self.bd[members][:, None, :]


def train(m, X, Y, masks, epochs, batch, lr, seed, freeze_body=False,
          only=None):
    """Fit the hive. freeze_body leaves the body untouched, which is what
    injection does; only restricts which member deltas move."""
    rg = np.random.default_rng(seed)
    n = X.shape[0]; S = m.S
    sel = np.arange(S) if only is None else np.asarray(only)
    P = [m.W0, m.W1, m.base, m.delta, m.b0, m.b1, m.bb, m.bd]
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    t = 0
    for ep in range(epochs):
        perm = rg.permutation(n)
        for st in range(0, n, batch):
            b = perm[st:st+batch]
            bi = to_dev(b, np.int64) if _GPU else b
            x = X[bi]; y = Y[bi]; mk = masks[:, bi][sel]
            z1, a1, z2, a2 = m.body(x)
            lg = m.logits(a2, sel)
            e = xp.exp(lg - lg.max(-1, keepdims=True))
            d = (e/e.sum(-1, keepdims=True) - y[None])*mk[:, :, None]
            d = d/max(1, len(b))
            gd = xp.einsum('ni,sno->sio', a2, d)
            gbd = d.sum(1)
            gbase = xp.einsum('ni,sno->io', a2, d)
            gbb = d.sum((0, 1))
            da2 = (xp.einsum('sno,io->ni', d, m.base)
                   + xp.einsum('sno,sio->ni', d, m.delta[sel]))
            G = [None]*8
            G[2] = gbase; G[6] = gbb
            G[3] = xp.zeros_like(m.delta); G[3][sel] = gd
            G[7] = xp.zeros_like(m.bd); G[7][sel] = gbd
            if freeze_body:
                G[0] = xp.zeros_like(m.W0); G[1] = xp.zeros_like(m.W1)
                G[4] = xp.zeros_like(m.b0); G[5] = xp.zeros_like(m.b1)
                G[2] = xp.zeros_like(m.base); G[6] = xp.zeros_like(m.bb)
            else:
                d2 = da2*(z2 > 0)
                G[1] = a1.T @ d2; G[5] = d2.sum(0)
                d1 = (d2 @ m.W1.T)*(z1 > 0)
                G[0] = x.T @ d1; G[4] = d1.sum(0)
            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_ - lr*(M[i]/(1-0.9**t))/(xp.sqrt(V[i]/(1-0.999**t))+1e-8)
            m.W0, m.W1, m.base, m.delta, m.b0, m.b1, m.bb, m.bd = P
            if not freeze_body:
                m.project()
                P[0], P[1] = m.W0, m.W1
    return m


def evaluate(m, X, y, masks, members):
    _, _, _, a2 = m.body(X)
    lg = to_host(m.logits(a2, np.asarray(members)))
    mk = to_host(masks[members]).astype(bool)
    out = []
    for i in range(len(members)):
        s = mk[i]
        out.append(float((lg[i][s].argmax(1) == y[s]).mean()) if s.any() else np.nan)
    return np.array(out)


def load(name="fashion"):
    if name == "fashion":
        from tensorflow import keras
        (a, b), (c, d) = keras.datasets.fashion_mnist.load_data()
        X = np.concatenate([a, c]).reshape(-1, 784).astype(np.float32)/255.0
        y = np.concatenate([b, d]).ravel().astype(np.int64)
    else:
        from sklearn.datasets import load_digits
        dd = load_digits()
        X = dd.data.astype(np.float32)/16.0; y = dd.target.astype(np.int64)
    return X, y


CFG = dict(dataset="fashion", width=64, n_train=8000, epochs=40,
           batch=128, lr=1e-3, seeds=(0, 1),
           old_classes=(0, 1, 2, 3, 4, 5, 6), new_classes=(7, 8, 9))
KS = [128, 256, 512, 1024, 2048, 4096]


def one_seed(X, y, seed, cfg):
    rg = np.random.default_rng(seed)
    old, new = list(cfg["old_classes"]), list(cfg["new_classes"])
    subs = [tuple(sorted(c)) for c in combinations(old, 4)]
    subs = [subs[i] for i in rg.choice(len(subs), 8, replace=False)]
    newsub = tuple(sorted(new))

    keep_old = np.isin(y, old); keep_new = np.isin(y, new)
    io = np.where(keep_old)[0]; inew = np.where(keep_new)[0]
    rg.shuffle(io); rg.shuffle(inew)
    tro, teo = io[:cfg["n_train"]], io[cfg["n_train"]:cfg["n_train"]+3000]
    trn, ten = inew[:cfg["n_train"]//3], inew[cfg["n_train"]//3:][:2000]

    D = X.shape[1]
    def mk(idx, ss):
        return to_dev(np.stack([np.isin(y[idx], s) for s in ss]).astype(np.float32))
    Y = lambda idx: to_dev(np.eye(10, dtype=np.float32)[y[idx]])

    Xtro, Xteo = to_dev(X[tro]), to_dev(X[teo])
    Xtrn, Xten = to_dev(X[trn]), to_dev(X[ten])
    S = len(subs) + 1                      # the new member is the last slot
    allsubs = subs + [newsub]

    # 1. the family, trained on the OLD classes only
    m = Hive(D, cfg["width"], 10, S, cfg["K"], seed)
    train(m, Xtro, Y(tro), mk(tro, allsubs), cfg["epochs"], cfg["batch"],
          cfg["lr"], seed, only=list(range(len(subs))))
    fam_before = evaluate(m, Xteo, y[teo], mk(teo, allsubs),
                          list(range(len(subs))))

    # 2. INJECTION: fit the new member's head against the FROZEN body
    train(m, Xtrn, Y(trn), mk(trn, allsubs), cfg["epochs"], cfg["batch"],
          cfg["lr"], seed+1, freeze_body=True, only=[S-1])
    inj = evaluate(m, Xten, y[ten], mk(ten, allsubs), [S-1])[0]
    fam_after = evaluate(m, Xteo, y[teo], mk(teo, allsubs),
                         list(range(len(subs))))

    # 3. the new member with a body of its own: its ceiling
    m2 = Hive(D, cfg["width"], 10, 1, cfg["K"], seed+2, tag="own")
    train(m2, Xtrn, Y(trn), mk(trn, [newsub]), cfg["epochs"], cfg["batch"],
          cfg["lr"], seed+2)
    own = evaluate(m2, Xten, y[ten], mk(ten, [newsub]), [0])[0]

    # 4. the body retrained with the new member included
    m3 = Hive(D, cfg["width"], 10, S, cfg["K"], seed+3, tag="re")
    Xb = xp.concatenate([Xtro, Xtrn])
    yb = np.concatenate([y[tro], y[trn]])
    Yb = to_dev(np.eye(10, dtype=np.float32)[yb])
    mkb = to_dev(np.stack([np.isin(yb, s) for s in allsubs]).astype(np.float32))
    train(m3, Xb, Yb, mkb, cfg["epochs"], cfg["batch"], cfg["lr"], seed+3)
    retr = evaluate(m3, Xten, y[ten], mk(ten, allsubs), [S-1])[0]

    return dict(family=float(np.nanmean(fam_before)), injected=float(inj),
                own_body=float(own), retrained=float(retr),
                family_after=float(np.nanmean(fam_after)))



def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("HOW MUCH RESERVE DOES A BODY NEED?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:12s} = {v}")
    print(f"  body storage sweep = {KS}")
    print(f"\n  the family is trained on classes {CFG['old_classes']}; the")
    print(f"  injected member covers {CFG['new_classes']}, which the body has")
    print(f"  never seen. Both are measured at every body size.")
    print("=" * 78, flush=True)

    X, y = load(CFG["dataset"])
    res = {}
    print(f"\n  {'stored':>8s} {'family':>8s} {'injected':>9s} "
          f"{'ceiling':>8s} {'shortfall':>10s}")
    for K in KS:
        cfg = dict(CFG); cfg["K"] = K
        runs = [one_seed(X, y, s, cfg) for s in cfg["seeds"]]
        m = {k: float(np.mean([r[k] for r in runs])) for k in runs[0]}
        # paired, because the arms share a seed
        gap = np.array([r["injected"] - r["own_body"] for r in runs])
        res[K] = dict(family=m["family"], injected=m["injected"],
                      own=m["own_body"], gap=float(gap.mean()),
                      gap_sd=float(gap.std()),
                      disturb=abs(m["family"] - m["family_after"]))
        print(f"  {K:8,} {m['family']:8.4f} {m['injected']:9.4f} "
              f"{m['own_body']:8.4f} {gap.mean():+10.4f}"
              f"   [{time.time()-t0:.0f}s]", flush=True)
        json.dump({str(k): v for k, v in res.items()},
                  open("provisioning.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    ks = sorted(res)
    print(f"  {'stored':>8s} {'family':>8s} {'family gain':>12s} "
          f"{'shortfall':>10s} {'closing by':>11s}")
    for a, b in zip([None]+ks, ks):
        r = res[b]
        fg = "" if a is None else f"{r['family']-res[a]['family']:+12.4f}"
        cl = "" if a is None else f"{r['gap']-res[a]['gap']:+11.4f}"
        print(f"  {b:8,} {r['family']:8.4f} {fg:>12s} {r['gap']:+10.4f} "
              f"{cl:>11s}")
    fam = np.array([res[k]["family"] for k in ks])
    gap = np.array([res[k]["gap"] for k in ks])
    fam_step = np.diff(fam); gap_step = np.diff(gap)
    sd = max(r["gap_sd"] for r in res.values())
    # where does the family stop improving, and does the shortfall keep
    # closing past that point?
    sat = next((ks[i+1] for i, s in enumerate(fam_step) if s < 2*sd), None)
    print(f"\n  disturbance to existing members, every size: "
          f"{max(r['disturb'] for r in res.values()):.6f}")
    if sat:
        i = ks.index(sat)
        after = gap[-1] - gap[i]
        print(f"\n  the family saturates at {sat:,} values — past that point")
        print(f"  more storage buys the members present "
              f"{fam[-1]-fam[i]:+.4f}")
        print(f"  and closes the injection shortfall by {after:+.4f}")
        print()
        if after > 2*sd:
            print("  THERE IS RESERVE, AND IT IS INVISIBLE FROM THE FAMILY.")
            print("  Storage past saturation does nothing for the members the")
            print("  body was built for and measurably helps one added later,")
            print("  so a base must be provisioned for members that do not")
            print("  exist yet — and no amount of watching the family would")
            print("  tell you how much.")
        else:
            print("  NO SEPARATE RESERVE. The shortfall closes with the same")
            print("  storage that helps the family, so provisioning for")
            print("  injection is provisioning for accuracy and there is no")
            print("  second number to decide.")
    else:
        print("\n  the family had not saturated within this sweep, so the")
        print("  question of what lies PAST saturation is not yet answerable;")
        print("  extend KS upward.")
    closed = [k for k in ks if res[k]["gap"] > -2*sd]
    if closed:
        print(f"\n  the shortfall first closes to within noise at "
              f"{closed[0]:,} values")
    else:
        print(f"\n  the shortfall never closes within this sweep; the best is"
              f" {max(res[k]['gap'] for k in ks):+.4f} at "
              f"{max(ks, key=lambda k: res[k]['gap']):,}")
    print(f"\n  total {time.time()-t0:.0f}s; wrote provisioning.json")


if __name__ == "__main__":
    main()