injection.py

"""
CAN A BODY SUPPORT A MEMBER IT NEVER SAW?

Everything measured about families so far trains them TOGETHER. The body
sees every member's data, the heads specialise on top, and the storage
economics follow: forty-five members share one body at seventeen times less
storage for a thousandth of a point.

An injection architecture needs something the programme has never tested.
A body is trained, frozen, shipped. Later a member arrives carrying
information the body never saw, and is added by fitting a head alone. If
that works, a family is a corpus one can extend after the fact; if it does
not, the body must have seen the domain and everything downstream needs a
pretraining step first.

The failsafe is what makes it safe to attempt. A member is a DELTA on a
permanent base head:

    output = base(a) + delta_m(a)

so no member injected means delta is zero and the base answers; an
unidentifiable member means ridge shrinks its coefficients toward zero and
the base answers again. There is no branch, no confidence threshold, no
code path that runs only sometimes. Injection is the addition of something
that may be zero, and the mechanism never changes.

The test is deliberately harsh. The family is trained on classes 0-6 ONLY.
The injected member covers 7-9, which the body has never seen in any form.
Four arms:

  FAMILY MEMBER    a member of the original family, for reference
  INJECTED         the new member's head fitted against the FROZEN body
  OWN BODY         the same member trained from scratch, its own ceiling
  RETRAINED        the body retrained with the new member included, which
                   is what one would have to do if injection fails

and a check that injecting changes nothing for the existing members, which
should hold exactly rather than approximately, since the body is frozen and
the heads are separate.
"""

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, K=512, n_train=8000, epochs=40,
           batch=128, lr=1e-3, seeds=(0, 1, 2),
           old_classes=(0, 1, 2, 3, 4, 5, 6), new_classes=(7, 8, 9))


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("CAN A BODY SUPPORT A MEMBER IT NEVER SAW?")
    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"\n  the family is trained on classes {CFG['old_classes']} only.")
    print(f"  the injected member covers {CFG['new_classes']}, which the body")
    print(f"  has never seen in any form.")
    print("=" * 78, flush=True)

    X, y = load(CFG["dataset"])
    runs = []
    for s in CFG["seeds"]:
        r = one_seed(X, y, s, CFG)
        runs.append(r)
        print(f"  seed {s}: " + "  ".join(f"{k} {v:.4f}" for k, v in r.items())
              + f"   [{time.time()-t0:.0f}s]", flush=True)
        json.dump(runs, open("injection.json", "w"), indent=2)

    keys = list(runs[0])
    mean = {k: float(np.mean([r[k] for r in runs])) for k in keys}
    sd = {k: float(np.std([r[k] for r in runs])) for k in keys}
    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    lab = {"family": "existing family member", "injected": "INJECTED, frozen body",
           "own_body": "new member, own body", "retrained": "body retrained with it",
           "family_after": "family after injection"}
    for k in ("family", "family_after", "injected", "own_body", "retrained"):
        print(f"  {lab[k]:>26s} {mean[k]:.4f}  sd {sd[k]:.4f}")
    dist = abs(mean["family"] - mean["family_after"])
    print(f"\n  injecting disturbed the existing members by {dist:.6f}")
    print(f"  (it should be EXACTLY zero: the body is frozen and each")
    print(f"   member's delta is its own, so nothing shared can move)")
    gap = mean["injected"] - mean["own_body"]
    gap2 = mean["injected"] - mean["retrained"]
    print(f"\n  the injected member against its own ceiling: {gap:+.4f}")
    print(f"  against a body retrained to include it:      {gap2:+.4f}")
    print()
    tol = 2*max(sd.values())
    if gap > -tol and gap2 > -tol:
        print("  INJECTION WORKS. A frozen body supports a member whose")
        print("  classes it never saw, as well as a body built for it. A")
        print("  family can be extended after the fact, which is what an")
        print("  injection architecture needs and what nothing here had")
        print("  tested.")
    elif gap2 > -tol:
        print("  INJECTION MATCHES RETRAINING but not a dedicated model, so")
        print("  the shared body costs something a private one would not —")
        print("  and it costs the same whether the member was there from the")
        print("  start or added later, which is the useful half.")
    else:
        print("  INJECTION FALLS SHORT. A frozen body does not support a")
        print("  member whose domain it never saw, so the body must see the")
        print("  domain first and an injection architecture needs a")
        print("  pretraining step rather than an empty one.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote injection.json")


if __name__ == "__main__":
    main()