curriculum.py

"""
STARTING BLIND AND LEARNING TO SEE.

§6 anneals a trained model DOWNWARD -- 4,096 values to 256, splitting
nothing, merging groups as it goes -- and finds the descent beats training
at the small budget directly: a 0.92 solution in eleven seeds of twelve at
48 parameters, where a cold start converges lower. Restarts improved the
direct arm without changing the annealed one, which points at the schedule
avoiding bad basins rather than merely spending compute.

This runs the ladder the other way, which nothing here has tried. Start
heavily folded -- few values, very blind -- and SPLIT groups as training
proceeds. If descending helps by avoiding bad basins, ascending begins in a
good one almost by construction: a model with 64 values has few parameters
and therefore few basins, so it can hardly land badly. Each split then adds
freedom around a solution already found.

The splitting is exact. With a power-of-two ladder and nested partitions,
group j at level K' has parent j >> 1 at level K, so a child initialised to
its parent's value leaves the function COMPLETELY UNCHANGED at the moment
of the split. Nothing is disturbed; the model simply gains the freedom to
differentiate what it previously had to treat alike.

That predicts a particular trace, and it is the thing to look for:

    FLATTEN, SPLIT, RESUME FALLING

If instead accuracy plateaus at every budget and the splits change nothing,
the model was capacity-bound throughout and the curriculum is a slower road
to the same place.

The control is training at the FINAL budget for the same total epochs. If
the curriculum does not beat that, starting constrained bought nothing.
"""

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 nested_ladder(shape, K_top, tag):
    """One assignment at the finest level; every coarser level is a merge.

    With powers of two, group j at level K' has parent j >> 1 at level
    K'/2, so splitting is exactly the inverse of merging and a child can
    inherit its parent's value with no change to the function."""
    import hashlib
    h = int(hashlib.md5(str((tuple(shape), K_top, tag)).encode()).hexdigest()[:8], 16)
    rg = np.random.default_rng(h)
    return rg.integers(0, K_top, int(np.prod(shape))).astype(np.int64)


def at_level(base_idx, K, K_top):
    idx = (base_idx*K)//K_top
    return to_dev(idx, np.int64) if _GPU else idx


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)


def scatter(dW, idx, K):
    g = xp.zeros(K, DT)
    if _GPU:
        import cupyx
        cupyx.scatter_add(g, idx, dW.reshape(-1))
    else:
        np.add.at(g, idx, dW.reshape(-1))
    return g


def split_values(v, K, K2):
    """Each group divides; both children take the parent's value.

    The function is unchanged at this instant — that is the whole point.
    The model gains freedom, it does not gain or lose anything it knew."""
    assert K2 % K == 0
    r = K2//K
    return xp.repeat(v, r)[:K2] if r > 1 else v.copy()


def augment(Xb, side, pad, rg):
    n = Xb.shape[0]
    im = Xb.reshape(n, 3, side, side)
    if pad:
        P = xp.zeros((n, 3, side+2*pad, side+2*pad), dtype=Xb.dtype)
        P[:, :, pad:pad+side, pad:pad+side] = im
        oy = rg.integers(0, 2*pad+1, n); ox = rg.integers(0, 2*pad+1, n)
        oy = xp.asarray(oy) if _GPU else oy
        ox = xp.asarray(ox) if _GPU else ox
        ar = xp.arange(side)
        im = P[xp.arange(n)[:, None, None, None],
               xp.arange(3)[None, :, None, None],
               (oy[:, None]+ar[None, :])[:, None, :, None],
               (ox[:, None]+ar[None, :])[:, None, None, :]]
    fl = rg.random(n) < 0.5
    fl = xp.asarray(fl) if _GPU else fl
    im = xp.where(fl[:, None, None, None], im[:, :, :, ::-1], im)
    return im.reshape(n, -1)


class Folded:
    def __init__(self, D, hid, nc, K, base_idx, K_top, seed, v=None):
        self.D, self.hid, self.K, self.K_top = D, hid, K, K_top
        self.base_idx = base_idx
        self.idx = at_level(base_idx, K, K_top)
        rg = np.random.default_rng(seed)
        if v is None:
            s = np.sqrt(2.0/D)*np.sqrt(max(1.0, D*hid/K))
            v = to_dev(rg.normal(0, s, K))
            v = fold(v[self.idx].reshape(D, hid), self.idx, K)
        self.P = [v, to_dev(rg.normal(0, np.sqrt(2.0/hid), (hid, nc))),
                  xp.zeros(hid, DT), xp.zeros(nc, DT)]
        self.M = [xp.zeros_like(p) for p in self.P]
        self.V = [xp.zeros_like(p) for p in self.P]
        self.t = 0

    def W(self):
        return self.P[0][self.idx].reshape(self.D, self.hid)

    def acc(self, Xte, yte):
        W = self.W(); out = []
        for s in range(0, Xte.shape[0], 4096):
            h = xp.maximum(Xte[s:s+4096] @ W + self.P[2], 0)
            out.append(to_host(h @ self.P[1] + self.P[3]))
        return float((np.concatenate(out).argmax(1) == yte).mean())

    def split_to(self, K2):
        """Grow the resolution. The function is preserved exactly, and the
        Adam moments come along so training resumes rather than restarts."""
        self.P[0] = split_values(self.P[0], self.K, K2)
        self.M[0] = split_values(self.M[0], self.K, K2)
        self.V[0] = split_values(self.V[0], self.K, K2)
        self.K = K2
        self.idx = at_level(self.base_idx, K2, self.K_top)

    def step(self, Xtr, Ytr, cfg, epochs, side, rg):
        n = Xtr.shape[0]
        for ep in range(epochs):
            perm = rg.permutation(n)
            for st in range(0, n, cfg["batch"]):
                b = perm[st:st+cfg["batch"]]
                x = Xtr[b]
                if cfg["augment"]:
                    x = augment(x, side, cfg["aug_pad"], rg)
                y = Ytr[b]
                W = self.W()
                z = x @ W + self.P[2]; h = xp.maximum(z, 0)
                lg = h @ self.P[1] + self.P[3]
                e = xp.exp(lg - lg.max(1, keepdims=True))
                d = (e/e.sum(1, keepdims=True) - y)/len(b)
                d0 = (d @ self.P[1].T)*(z > 0)
                G = [scatter(x.T @ d0, self.idx, self.K), h.T @ d,
                     d0.sum(0), d.sum(0)]
                self.t += 1
                for i, (p_, gr) in enumerate(zip(self.P, G)):
                    self.M[i] = 0.9*self.M[i] + 0.1*gr
                    self.V[i] = 0.999*self.V[i] + 0.001*gr*gr
                    self.P[i] = p_ - cfg["lr"]*(self.M[i]/(1-0.9**self.t)) \
                        / (xp.sqrt(self.V[i]/(1-0.999**self.t))+1e-8)


CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465])
CIFAR_STD = np.array([0.2470, 0.2435, 0.2616])


def _verified_cifar(verbose=True):
    import glob, pickle
    for root in sorted(glob.glob("/kaggle/input/*")) + \
            ["/kaggle/working", "./cifar", "./data", "."]:
        if not os.path.isdir(root):
            continue
        cands = []
        for p in glob.glob(os.path.join(root, "**", "*"), recursive=True):
            if not os.path.isfile(p):
                continue
            try:
                if p.endswith(".npy"):
                    cands.append(np.load(p, allow_pickle=True))
                elif (p.endswith((".pickle", ".pkl", ".p"))
                      or os.path.basename(p).startswith(("data_batch",
                                                         "test_batch"))):
                    with open(p, "rb") as fh:
                        d = pickle.load(fh, encoding="bytes")
                    if isinstance(d, dict):
                        for kk, val in d.items():
                            kk = kk.decode() if isinstance(kk, bytes) else kk
                            a = np.asarray(val)
                            if a.size > 100 and kk in ("data", "labels",
                                                       "fine_labels", "x", "y"):
                                cands.append(a)
            except Exception:
                continue
        imgs = [a for a in cands if a.ndim >= 2 and len(a) >= 1000
                and a.size // len(a) == 3072]
        labs = [np.asarray(a).ravel() for a in cands
                if a.ndim <= 2 and np.issubdtype(np.asarray(a).dtype,
                                                 np.integer)
                and 1000 <= a.size <= 100000]
        if not imgs or not labs:
            continue
        X = np.concatenate([a.reshape(len(a), -1) for a in imgs])
        y = np.concatenate(labs)
        if len(y) != len(X):
            continue
        Xf = X.astype(np.float64)
        if Xf.max() > 1.5:
            Xf = Xf/255.0
        best = None
        for lay, shp in (("HWC", (-1, 32, 32, 3)), ("CHW", (-1, 3, 32, 32))):
            im = Xf[:2000].reshape(shp)
            mu = im.mean((0, 1, 2)) if lay == "HWC" else im.mean((0, 2, 3))
            sd = im.std((0, 1, 2)) if lay == "HWC" else im.std((0, 2, 3))
            e = float(np.abs(mu-CIFAR_MEAN).max() + np.abs(sd-CIFAR_STD).max())
            if best is None or e < best[0]:
                best = (e, lay)
        if best[0] > 0.05:
            continue
        if verbose:
            print(f"  verified CIFAR-10 at {root} ({best[1]}, {len(X):,} "
                  f"images)", flush=True)
        Xr = Xf.astype(np.float32)
        Xr = (Xr.reshape(-1, 32, 32, 3) if best[1] == "HWC"
              else Xr.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1))
        return Xr, y.astype(np.int64)
    if verbose:
        print("  no verified CIFAR-10 found; downloading", flush=True)
    return None


def load(cfg):
    got = _verified_cifar()
    if got is not None:
        X, y = got
    else:
        from tensorflow import keras
        (a, b), (c, d) = keras.datasets.cifar10.load_data()
        X = np.concatenate([a, c]).astype(np.float32)/255.0
        y = np.concatenate([b, d]).ravel().astype(np.int64)
    side = 32
    if cfg["grid"] != side:
        s = side // cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s, 3).mean(axis=(2, 4))
        side = cfg["grid"]
    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 = X[tr].mean((0, 1, 2), keepdims=True)
    sd = X[tr].std((0, 1, 2), keepdims=True) + 1e-8
    f = lambda Z: ((Z-mu)/sd).transpose(0, 3, 1, 2).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, f(X[te]), y[te], side


CFG = dict(grid=16, c_in=3, hidden=2048, n_train=20000, batch=128, lr=1e-3,
           augment=True, aug_pad=2, seeds=(0, 1), stage_epochs=20)
LADDER = [64, 128, 256, 512]      # powers of two, so splitting is exact


def one_seed(Xtr, Ytr, Xte, yte, side, cfg, seed):
    D, hid = Xtr.shape[1], cfg["hidden"]
    K_top = LADDER[-1]
    base = nested_ladder((D, hid), K_top, "curr")
    rg = np.random.default_rng(seed + 4242)
    E = cfg["stage_epochs"]
    trace = []

    # the curriculum: start blind, split, continue
    m = Folded(D, hid, 10, LADDER[0], base, K_top, seed)
    for i, K in enumerate(LADDER):
        if i:
            before = m.acc(Xte, yte)
            m.split_to(K)
            after = m.acc(Xte, yte)
            trace.append(dict(event="split", to=K, before=before,
                              after=after))
        m.step(Xtr, Ytr, cfg, E, side, rg)
        trace.append(dict(event="train", K=K, acc=m.acc(Xte, yte)))
    curr = m.acc(Xte, yte)

    # the control: the final budget, from cold, for the same total epochs
    d = Folded(D, hid, 10, LADDER[-1], base, K_top, seed + 7)
    d.step(Xtr, Ytr, cfg, E*len(LADDER), side, rg)
    direct = d.acc(Xte, yte)

    # and the coarsest budget alone, so the ladder's own gain is visible
    c = Folded(D, hid, 10, LADDER[0], base, K_top, seed + 13)
    c.step(Xtr, Ytr, cfg, E*len(LADDER), side, rg)
    coarse = c.acc(Xte, yte)
    return dict(curriculum=curr, direct=direct, coarse=coarse, trace=trace)


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("STARTING BLIND AND LEARNING TO SEE")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:13s} = {v}")
    print(f"  ladder        = {LADDER}, {CFG['stage_epochs']} epochs a rung")
    print(f"  total         = {CFG['stage_epochs']*len(LADDER)} epochs, "
          f"matched across arms")
    print(f"\n  a split leaves the function EXACTLY unchanged, so the trace")
    print(f"  shows whether new freedom is used: flatten, split, resume")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte, side = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    runs = []
    for s in CFG["seeds"]:
        r = one_seed(Xtr, Ytr, Xte, yte, side, CFG, s)
        runs.append(r)
        print(f"  seed {s}: curriculum {r['curriculum']:.4f}  "
              f"direct {r['direct']:.4f}  coarse-only {r['coarse']:.4f}"
              f"   [{time.time()-t0:.0f}s]", flush=True)
        json.dump(runs, open("curriculum.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  THE TRACE (first seed)")
    print("=" * 78)
    print(f"  {'':>22s} {'accuracy':>9s}")
    for e in runs[0]["trace"]:
        if e["event"] == "train":
            print(f"  {'trained at ' + str(e['K']):>22s} {e['acc']:9.4f}")
        else:
            d = e["after"] - e["before"]
            print(f"  {'SPLIT to ' + str(e['to']):>22s} {e['after']:9.4f}"
                  f"   (function change {d:+.4f}, should be zero)")

    m = {k: float(np.mean([r[k] for r in runs]))
         for k in ("curriculum", "direct", "coarse")}
    sd = {k: float(np.std([r[k] for r in runs]))
          for k in ("curriculum", "direct", "coarse")}
    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    for k in ("curriculum", "direct", "coarse"):
        print(f"  {k:>12s} {m[k]:.4f}  sd {sd[k]:.4f}")
    gap = m["curriculum"] - m["direct"]
    tol = 2*max(sd.values())
    print(f"\n  curriculum against training at the final budget: {gap:+.4f}")
    print(f"  seed spread (worst) {max(sd.values()):.4f}")
    gains = [e["acc"] for e in runs[0]["trace"] if e["event"] == "train"]
    print(f"  the rungs gained "
          f"{'  '.join(f'{b-a:+.4f}' for a, b in zip(gains, gains[1:]))}")
    print()
    if gap > tol:
        print("  STARTING BLIND WINS. A model constrained hard and then")
        print("  progressively freed beats the same budget trained from cold,")
        print("  at matched epochs — which is §6's annealing result running")
        print("  the other way, and says the schedule is doing optimisation")
        print("  work rather than spending compute.")
    elif gap < -tol:
        print("  STARTING BLIND LOSES. The constrained early phase costs more")
        print("  than the good initialisation it buys, so the curriculum is a")
        print("  slower road to a worse place.")
    else:
        print("  NO DIFFERENCE. The curriculum arrives where cold training")
        print("  arrives, so the constrained start neither helps nor hurts —")
        print("  and §6's downward annealing benefit does not have a mirror")
        print("  image on the way up.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote curriculum.json")


if __name__ == "__main__":
    main()