calibration.py

"""
WHAT HAS TO BE IN A CALIBRATION SET?

A member fitted on one base transfers to another through a linear map
between their features, fitted on inputs both parties can run. Fifty public
examples recover 65% of what a native member is worth, five thousand
recover 92%, and whether the two bases share a partition makes no
difference at all. So the coordination point is not a coordinate system: it
is a set of shared inputs.

Which raises the question that decides how hard that coordination is. WHAT
HAS TO BE IN THAT SET?

If it must resemble the member's task, someone has to curate one per domain
and the standard is a real institution. If any inputs from roughly the
right kind of data will do, a public dataset suffices. And if RANDOM NOISE
works, there is nothing to coordinate at all: two parties agree on a seed
and generate the same inputs from nothing.

Six sources, at four sizes, between bases on DIFFERENT partitions, since
that condition was measured to be no worse and is the conservative one:

  BASE CLASSES     held-out examples of what both bases were trained on.
                   This is what the first measurement used.
  MEMBER CLASSES   examples of what the member is FOR, which neither base
                   was trained on. The task-relevant case.
  MIXED            half of each.
  SHUFFLED PIXELS  real images with their pixels permuted per image. The
                   same intensity statistics and no spatial structure, so
                   this separates "the right kind of picture" from "any
                   picture at all".
  UNIFORM NOISE    values drawn at random. No structure of any kind, and
                   generable from a shared seed by anyone.
  ANOTHER DATASET  handwritten digits, where the bases were trained on
                   clothing. Real images from an unrelated domain.

The map is linear and the features are 3,136 wide, so at fifty examples it
is heavily underdetermined and ridge is doing most of the work. That it
succeeds there at all says the useful part of the map is low-dimensional,
and this measures whether the directions it needs are present in whatever
the calibration set happens to contain.
"""

import numpy as np
import time
import json
import os

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, tag=0):
    """A 3x3 convolution as a tying pattern. tag selects WHICH pattern: the
    same tag is the same partition, which is the coordinate agreement two
    bases either share or do not."""
    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)
    if tag:
        # A GENUINELY DIFFERENT PARTITION, at matched storage. A first
        # version merely PERMUTED which value sat in which slot, which is
        # not a different partition at all: channel labels are
        # interchangeable by construction, and two seeds were already
        # measured to differ by a permutation of 14.7 of 16 channels. That
        # control controlled nothing, and finding no difference between the
        # conditions said nothing. This ties the same number of values
        # ARBITRARILY, so the two bases agree about how much they store and
        # about nothing else.
        rg = np.random.default_rng(9000 + tag)
        idx = rg.integers(0, K, idx.size).reshape(idx.shape)
    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 ridge(A, R, lam):
    """Least squares with the ridge penalty CHOSEN rather than fixed.

    A fixed lambda means the effective regularisation falls as data
    arrives, because the data term grows with n and the penalty does not.
    That is fine while n is far from the feature count and catastrophic
    near it: a fit on 4,000 examples of 3,136 features sits at n/d = 1.28,
    the interpolation peak, where the variance of a weakly regularised
    solution explodes. Measured, a member fitted on 4,000 examples scored
    WORSE than one fitted on 1,000, which is a property of the solver and
    not of members.

    So lambda is selected on a held-out fifth, from a geometric ladder,
    per fit. It costs a handful of solves and removes an artefact that
    reached the paper."""
    n, d = A.shape
    if n < 8:
        return _ridge_at(A, R, lam)
    k = max(2, n//5)
    tr, va = slice(k, None), slice(0, k)
    best, bw = None, None
    for f in (0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 1e4):
        W, b = _ridge_at(A[tr], R[tr], lam*f)
        err = float(to_host(((A[va] @ W + b - R[va])**2).mean()))
        if best is None or err < best:
            best, bw = err, lam*f
    return _ridge_at(A, R, bw)


def _ridge_at(A, R, lam):
    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:
        G = A1 @ A1.T + lam*xp.eye(n, dtype=DT)
        W = A1.T @ xp.linalg.solve(G, R)
    return W[:-1], W[-1]


def train_base(Xtr, Ytr, cfg, seed, tag):
    """One person's base: their own data, their own initialisation, and a
    partition that is either shared with the other party or not."""
    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, tag)
        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


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))
    return ((X - X[p[:20000]].mean())/(X[p[:20000]].std()+1e-8)
            ).reshape(len(X), -1), y, p




CFG = dict(grid=14, c_in=1, chan=16, depth=3, batch=128, lr=1e-3, epochs=25,
           base_classes=(0, 1, 2, 3, 4, 5, 6), new_classes=(7, 8, 9),
           n_each=9000, n_member=400, lam=1.0,
           calib=(50, 200, 1000, 5000), pairs=(0, 1, 2))


def calib_sources(X, y, perm, cfg, pool_start, n_max, rg):
    """The six candidate calibration sets, all the same shape."""
    old = np.isin(y, cfg["base_classes"])
    oi = perm[np.isin(perm, np.where(old)[0])][pool_start:pool_start+n_max]
    ni = perm[np.isin(perm, np.where(~old)[0])][-n_max:]
    half = n_max//2
    src = {"base classes": X[oi],
           "member classes": X[ni],
           "mixed": np.concatenate([X[oi[:half]], X[ni[:half]]])}
    flat = X[oi].copy()
    for r in flat:
        rg.shuffle(r)
    src["shuffled pixels"] = flat
    src["uniform noise"] = rg.normal(0, 1, X[oi].shape).astype(np.float32)
    try:
        from tensorflow import keras
        (a, b), (c, d) = keras.datasets.mnist.load_data()
        M = np.concatenate([a, c]).astype(np.float32)/255.0
        g = cfg["grid"]
        if g != 28:
            sfac = 28//g
            M = M.reshape(-1, g, sfac, g, sfac).mean(axis=(2, 4))
        M = ((M - M.mean())/(M.std()+1e-8)).reshape(len(M), -1)
        src["another dataset"] = M[:n_max]
    except Exception as e:
        print(f"  (digits unavailable: {e})")
    return src


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("WHAT HAS TO BE IN A CALIBRATION SET?")
    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 two bases use DIFFERENT partitions, which was measured to")
    print(f"  cost nothing and is the conservative choice")
    print("=" * 78, flush=True)

    X, y, perm = load(CFG)
    old = np.isin(y, CFG["base_classes"])
    oi = perm[np.isin(perm, np.where(old)[0])]
    ni = perm[np.isin(perm, np.where(~old)[0])]
    n = CFG["n_each"]
    Adata, Bdata = oi[:n], oi[n:2*n]
    memb_pool = ni[:CFG["n_member"]*4]
    test = np.concatenate([oi[2*n+6000:2*n+9000],
                           ni[CFG["n_member"]*4:][:2000]])
    yte = y[test]
    nm = np.array(CFG["new_classes"])
    new_te = np.isin(yte, nm)
    te3 = np.where(new_te)[0]
    y3 = np.searchsorted(nm, yte[te3])
    oh = lambda i: to_dev(np.eye(10, dtype=np.float32)[y[i]])
    oh3 = lambda i: to_dev((y[i][:, None] == nm[None, :]).astype(np.float32))
    Xd = lambda i: to_dev(X[i])

    rg0 = np.random.default_rng(11)
    sources = calib_sources(X, y, perm, CFG, 2*n, max(CFG["calib"]), rg0)
    print(f"\n  {len(sources)} calibration sources, "
          f"{max(CFG['calib']):,} examples each", flush=True)

    res, ceil, home = {}, [], []
    for pair in CFG["pairs"]:
        A_P, A_f, A_H, A_O = train_base(Xd(Adata), oh(Adata), CFG, 100+pair, 0)
        B_P, B_f, B_H, B_O = train_base(Xd(Bdata), oh(Bdata), CFG, 500+pair, 1)
        fA_te, _ = A_f(Xd(test)); fB_te, _ = B_f(Xd(test))
        fA_mb, _ = A_f(Xd(memb_pool)); fB_mb, _ = B_f(Xd(memb_pool))
        sub = np.random.default_rng(pair).choice(
            len(memb_pool), CFG["n_member"], replace=False)
        sd_ = to_dev(sub, np.int64) if _GPU else sub
        WA, bA = ridge(fA_mb[sd_], oh3(memb_pool[sub]), CFG["lam"])
        WB, bB = ridge(fB_mb[sd_], oh3(memb_pool[sub]), CFG["lam"])
        sc = lambda F, W, b: float(
            (to_host(F[te3] @ W + b).argmax(1) == y3).mean())
        home.append(sc(fA_te, WA, bA)); ceil.append(sc(fB_te, WB, bB))
        res.setdefault(("raw", 0), []).append(sc(fB_te, WA, bA))
        # THE CONTROL THE MARKETPLACE RESTS ON. If a borrower holds N
        # examples of the domain, the alternative to borrowing is fitting
        # their OWN member on those same N. Borrowing is only worth
        # anything if it beats that, and the whole value proposition is
        # the gap: A's member was fitted on far more data than B holds, so
        # borrowing amplifies a small local sample by importing someone
        # else's larger one. Or it does not, and everyone fits their own.
        mp = ni[:CFG["n_member"]*4]
        for nc in CFG["calib"]:
            k = np.random.default_rng(300+pair).choice(
                len(mp), min(nc, len(mp)), replace=False)
            kd = to_dev(k, np.int64) if _GPU else k
            Wl, bl = ridge(fB_mb[kd], oh3(mp[k]), CFG["lam"])
            res.setdefault(("B fits its own on N", nc), []).append(
                sc(fB_te, Wl, bl))
        for name, Xc in sources.items():
            Xc_d = to_dev(Xc)
            fA_c, _ = A_f(Xc_d); fB_c, _ = B_f(Xc_d)
            for nc in CFG["calib"]:
                k = to_dev(np.arange(min(nc, len(Xc))), np.int64) \
                    if _GPU else np.arange(min(nc, len(Xc)))
                T, tb = ridge(fB_c[k], fA_c[k], CFG["lam"])
                res.setdefault((name, nc), []).append(float(
                    (to_host((fB_te[te3] @ T + tb) @ WA + bA).argmax(1)
                     == y3).mean()))
        print(f"  pair {pair}: A at home {home[-1]:.4f}, B's own "
              f"{ceil[-1]:.4f}   [{time.time()-t0:.0f}s]", flush=True)
        _FIXED.clear()
        if _GPU:
            _cp.get_default_memory_pool().free_all_blocks()

    m = {k: (float(np.mean(v)), float(np.std(v))) for k, v in res.items()}
    C, H, F = float(np.mean(ceil)), float(np.mean(home)), 1.0/len(nm)
    json.dump({f"{a}/{b}": list(v) for (a, b), v in m.items()},
              open("calibration.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  RECOVERED, AS A SHARE OF WHAT A NATIVE MEMBER IS WORTH")
    print("=" * 78)
    print(f"  chance {F:.4f}; B's own member {C:.4f}; A's at home {H:.4f}\n")
    share = lambda v: (v - F)/max(C - F, 1e-9)
    print(f"  {'source':>17s} " + " ".join(f"{n:>9,}" for n in CFG["calib"]))
    order = []
    for name in list(sources) :
        row = [m[(name, nc)][0] for nc in CFG["calib"]]
        order.append((max(row), name))
        print(f"  {name:>17s} " + " ".join(f"{share(v):8.0%} " for v in row))
    print(f"  {'(applied raw)':>17s} {share(m[('raw', 0)][0]):8.0%}")
    print()
    row = [m[("B fits its own on N", nc)][0] for nc in CFG["calib"]]
    print(f"  {'B fits its own':>17s} " + " ".join(f"{share(v):8.0%} "
                                                  for v in row))
    print(f"  {'':>17s} " + " ".join(
        f"{share(m[('member classes', nc)][0]) - share(v):+8.0%} "
        for nc, v in zip(CFG["calib"], row)) + "  <- borrowing's advantage")
    sd = max(v[1] for v in m.values())
    print(f"\n  worst spread across pairs {sd:.4f}")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    order.sort(reverse=True)
    best_v, best_n = order[0]
    worst_v, worst_n = order[-1]
    noise = max(m[("uniform noise", nc)][0] for nc in CFG["calib"])
    print(f"  best source: {best_n} at {share(best_v):.0%} of a native member")
    print(f"  worst: {worst_n} at {share(worst_v):.0%}\n")
    if share(noise) > 0.5:
        print(f"  RANDOM NOISE WORKS, at {share(noise):.0%}. There is nothing to")
        print(f"  coordinate: two parties agree on a seed, generate the same")
        print(f"  inputs from nothing, and fit the map. No dataset has to be")
        print(f"  published, curated or agreed, and the standard the")
        print(f"  possibilities document proposed reduces to a random number")
        print(f"  generator.")
    elif share(noise) > 0.2:
        print(f"  RANDOM NOISE PARTLY WORKS, at {share(noise):.0%} against")
        print(f"  {share(best_v):.0%} for the best real source. Structure helps")
        print(f"  and is not required, so a calibration set should be real")
        print(f"  data and need not be curated.")
    else:
        print(f"  RANDOM NOISE FAILS, at {share(noise):.0%}. The map needs")
        print(f"  inputs that put the features where real data puts them, so")
        print(f"  a calibration set has to be REAL and someone has to publish")
        print(f"  one.")
    rel = max(m[("member classes", nc)][0] for nc in CFG["calib"])
    base = max(m[("base classes", nc)][0] for nc in CFG["calib"])
    print()
    if rel > base + 2*sd:
        print(f"  AND IT HAS TO BE TASK-RELEVANT: examples of what the member")
        print(f"  is FOR give {share(rel):.0%} against {share(base):.0%} for the")
        print(f"  bases' own classes. A calibration set is per DOMAIN, which")
        print(f"  is a real institution to build.")
    elif base > rel + 2*sd:
        print(f"  AND TASK-RELEVANCE HURTS: the bases' own classes give")
        print(f"  {share(base):.0%} against {share(rel):.0%} for the member's. The")
        print(f"  map is fitted where both bases have LEARNED something, not")
        print(f"  where the member operates.")
    else:
        print(f"  AND TASK-RELEVANCE DOES NOT MATTER: {share(rel):.0%} against")
        print(f"  {share(base):.0%}. One calibration set serves any member")
        print(f"  between a pair of bases, whatever it carries.")
    own = [m[("B fits its own on N", nc)][0] for nc in CFG["calib"]]
    bor = [m[("member classes", nc)][0] for nc in CFG["calib"]]
    adv = [share(b) - share(o) for b, o in zip(bor, own)]
    print(f"\n  IS BORROWING WORTH ANYTHING? At each N, against a borrower")
    print(f"  who simply fits their own member on the same N examples:")
    for nc, a in zip(CFG["calib"], adv):
        print(f"    {nc:6,}: {a:+.0%}")
    if max(adv) > 0.05:
        i = int(np.argmax(adv))
        print(f"\n  YES, AND MOST AT SMALL N: {adv[i]:+.0%} at "
              f"{CFG['calib'][i]:,} examples. Borrowing AMPLIFIES a small")
        print(f"  local sample by importing someone else's larger one, which")
        print(f"  is exactly what a marketplace is for. The advantage should")
        print(f"  close as the borrower's own data grows, and it does.")
    elif max(adv) > -0.05:
        print(f"\n  BARELY. Borrowing and fitting locally land within five")
        print(f"  points at every N, so a marketplace saves the fitting and")
        print(f"  not the data, which is a much weaker proposition.")
    else:
        print(f"\n  NO. A borrower holding N examples does better fitting")
        print(f"  their own member on them than importing one and")
        print(f"  translating it. The marketplace solves a problem that does")
        print(f"  not exist at this scale.")
    small = [nc for nc in CFG["calib"]
             if share(max(m[(n2, nc)][0] for _, n2 in order)) > 0.5]
    if small:
        print(f"\n  and the price of entry is {min(small):,} examples for half")
        print(f"  of a native member.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote calibration.json")


if __name__ == "__main__":
    main()