Harness for knowledge-as-dynamics issue #1: read-carrier (2D union/lineage/onion), k-NN + Wiener/GP carriers on the MNIST tier, basin(N) curve, d-scaling 8-32, orderability frontier, uncertainty-exponent sweep. Self-contained (mlx+numpy).

spike_2d.py

"""Read-carrier vs fit-carrier on the knowledge-as-dynamics union test.

Claim under test: the losses reported in Helou 2026 (knowledge-as-dynamics) —
union seam loss (0.82 vs ceiling 0.98), lineage decay, and the "onion filter"
(disordered knowledge cannot be inherited) — are properties of the CARRIER
(SGD regression into a parametric AE), not of the knowledge channel
(the latent vector field V(z) = Enc(Dec(z)) - z).

Carrier here: a lattice snapshot of the field (512x512 grid, 262k samples =
7.8x FEWER field evaluations than their student's 8000 steps x 256 batch),
read back by nearest-cell lookup. No gradients, no training. Composition =
the same gated union, read at lookup time (adelic shape: local-per-territory
tables + gate at read). Inheritance = resample the lookup field on a jittered
lattice.

World/teacher/oracle logic adapted line-for-line from the reviewed source of
github.com/Veso-AI-Open-Source/knowledge-as-dynamics (MIT), files
src/m5_fielddistill/{ae,train,gi}.py, so numbers are comparable to their
runs/gi.json (ceiling 0.98, union student 0.82, lineage 0.91->0.83,
contaminated B-side 0.24->0.19).

Run: uv run --no-project --with mlx --with numpy python spike_2d.py
"""

import json
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np

BASIN_EPS = 0.3
ORACLE_STEPS = 400
GENS = 5
GRID_N = 512  # lattice resolution -> 262,144 field samples

OUT = Path(__file__).parent / "spike_2d_results.json"


# ---------- model + field (adapted from their ae.py) ----------

class MLP(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.layers = [nn.Linear(a, b) for a, b in zip(dims[:-1], dims[1:])]

    def __call__(self, x):
        n = len(self.layers)
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < n - 1:
                x = nn.gelu(x)
        return x


class AE(nn.Module):
    def __init__(self, d_in=16, d_lat=2, width=64):
        super().__init__()
        self.enc = MLP([d_in, width, width, d_lat])
        self.dec = MLP([d_lat, width, width, d_in])

    def __call__(self, x):
        return self.dec(self.enc(x))


def field(model, z):
    return model.enc(model.dec(z)) - z


def np_field(model, z):
    return np.array(field(model, mx.array(z.astype(np.float32))))


def np_step(model, z):
    zm = mx.array(z.astype(np.float32))
    return np.array(model.enc(model.dec(zm)))


def find_attractors(endpoints, vnorms, vtol=5e-3, merge_r=0.15):
    pts = endpoints[vnorms < vtol]
    centers = []
    for p in pts:
        if all(np.linalg.norm(p - c) > merge_r for c in centers):
            centers.append(p)
    return np.array(centers) if centers else np.zeros((0, endpoints.shape[1]))


def train_teacher(x_train, steps=4000, bs=256, lr=1e-3, seed=0, noise=0.1):
    mx.random.seed(seed)
    model = AE(d_in=x_train.shape[1])
    mx.eval(model.parameters())
    opt = optim.Adam(learning_rate=lr)
    X = mx.array(x_train)
    rng = np.random.default_rng(seed)
    d = x_train.shape[1]

    def loss_fn(m, xb, xn):
        return mx.mean((m(xn) - xb) ** 2)

    lg = nn.value_and_grad(model, loss_fn)
    for _ in range(steps):
        idx = mx.array(rng.integers(0, len(x_train), bs))
        xb = X[idx]
        xn = xb + noise * mx.array(rng.standard_normal((bs, d)).astype(np.float32))
        _, grads = lg(model, xb, xn)
        opt.update(model, grads)
        mx.eval(model.parameters(), opt.state)
    return model


# ---------- world + union frame (adapted from their gi.py) ----------

def make_world8(n_per=500, radius=2.2, std=0.32, d_high=16):
    rng = np.random.default_rng(0)
    angles = np.arange(8) / 8 * 2 * np.pi
    centers = radius * np.stack([np.cos(angles), np.sin(angles)], axis=1)
    labels = np.repeat(np.arange(8), n_per)
    x2 = centers[labels] + std * rng.normal(size=(len(labels), 2))
    W1 = rng.normal(size=(2, d_high)) / np.sqrt(2)
    b1 = 0.1 * rng.normal(size=(d_high,))
    W2 = rng.normal(size=(d_high, d_high)) / np.sqrt(d_high)
    xh = (np.tanh(x2 @ W1 + b1) @ W2).astype(np.float32)
    return xh, labels


class UnionFrame:
    """Isometric graft only (their mode='iso'): B's chart translated into empty
    frame territory; hard nearest-centroid gate between V_A and grafted V_B."""

    def __init__(self, A, B, x_all, labels):
        self.A, self.B = A, B
        xA, xB = x_all[labels < 4], x_all[labels >= 4]
        zA_own = np.array(A.enc(mx.array(xA)))
        zB_own = np.array(B.enc(mx.array(xB)))
        gap = 1.5
        d = np.array([zA_own[:, 0].max() - zB_own[:, 0].min() + gap,
                      zA_own[:, 1].mean() - zB_own[:, 1].mean()], dtype=np.float32)
        eye = np.eye(2, dtype=np.float32)
        self.T = (eye, d, eye)
        W, b, _ = self.T
        graft = zB_own @ W + b
        labA, labB = labels[labels < 4], labels[labels >= 4]
        cents = [zA_own[labA == c].mean(0) for c in range(4)]
        cents += [graft[labB == c].mean(0) for c in range(4, 8)]
        self.centroids = np.stack(cents)
        allz = np.concatenate([zA_own, graft])
        lo, hi = allz.min(0), allz.max(0)
        span = hi - lo
        self.lo, self.hi = (lo - 0.35 * span).astype(np.float32), (hi + 0.35 * span).astype(np.float32)
        self.clip_lo, self.clip_hi = self.lo - 2 * span, self.hi + 2 * span

    def gate(self, z):
        d = np.linalg.norm(z[:, None, :] - self.centroids[None], axis=-1)
        return d.argmin(axis=1) >= 4

    def field_np(self, z):
        vA = np_field(self.A, z)
        W, b, Winv = self.T
        zB = (z - b) @ Winv
        fB = np_step(self.B, zB)
        vB = (fB @ W + b) - z
        return np.where(self.gate(z)[:, None], vB, vA)

    def iterate_np(self, z, steps=ORACLE_STEPS):
        z = z.copy()
        for _ in range(steps):
            z = np.clip(z + self.field_np(z), self.clip_lo, self.clip_hi)
        return z

    def oracle(self, rng, n_probes=400, grid_n=30):
        probes = (self.lo + rng.random((n_probes, 2)).astype(np.float32) * (self.hi - self.lo))
        ep = self.iterate_np(probes)
        ep_pert = self.iterate_np(
            probes + 0.05 * rng.standard_normal(probes.shape).astype(np.float32))
        d = np.linalg.norm(ep - ep_pert, axis=1)
        side = self.gate(probes)
        xs = np.linspace(self.lo[0], self.hi[0], grid_n, dtype=np.float32)
        ys = np.linspace(self.lo[1], self.hi[1], grid_n, dtype=np.float32)
        gx, gy = np.meshgrid(xs, ys)
        G = np.stack([gx.ravel(), gy.ravel()], axis=1)
        epg = self.iterate_np(G)
        vn = np.linalg.norm(self.field_np(epg), axis=1)
        attractors = find_attractors(epg, vn)
        return dict(probes=probes, ep=ep, side=side,
                    ceil_joint=float((d < BASIN_EPS).mean()),
                    ceil_A=float((d[~side] < BASIN_EPS).mean()),
                    ceil_B=float((d[side] < BASIN_EPS).mean()),
                    attractors=attractors, G=G)


# ---------- the read carrier ----------

class LatticeCarrier:
    """Piecewise-constant snapshot of a field on a regular lattice.
    Knowledge = the table. Read = nearest-cell lookup. No training."""

    def __init__(self, field_fn, lo, hi, n=GRID_N, chunk=8192, origin_jitter=None):
        self.lo, self.hi, self.n = lo.copy(), hi.copy(), n
        self.h = (hi - lo) / (n - 1)
        if origin_jitter is not None:
            self.lo = self.lo + origin_jitter * self.h  # shifted lattice for lineage resampling
        xs = self.lo[0] + np.arange(n, dtype=np.float32) * self.h[0]
        ys = self.lo[1] + np.arange(n, dtype=np.float32) * self.h[1]
        gx, gy = np.meshgrid(xs, ys, indexing="ij")
        pts = np.stack([gx.ravel(), gy.ravel()], axis=1)
        V = np.empty_like(pts)
        for i in range(0, len(pts), chunk):
            V[i:i + chunk] = field_fn(pts[i:i + chunk])
        self.V = V.reshape(n, n, 2)
        self.n_samples = len(pts)

    def lookup(self, z):
        ij = np.rint((z - self.lo) / self.h).astype(np.int64)
        ij = np.clip(ij, 0, self.n - 1)
        return self.V[ij[:, 0], ij[:, 1]]

    def iterate(self, z, clip_lo, clip_hi, steps=ORACLE_STEPS):
        z = z.copy()
        for _ in range(steps):
            z = np.clip(z + self.lookup(z), clip_lo, clip_hi)
        return z


def eval_carrier(carrier, frame, orc):
    ep = carrier.iterate(orc["probes"], frame.clip_lo, frame.clip_hi)
    d = np.linalg.norm(ep - orc["ep"], axis=1)
    side = orc["side"]
    epg = carrier.iterate(orc["G"], frame.clip_lo, frame.clip_hi)
    vn = np.linalg.norm(carrier.lookup(epg), axis=1)
    attr = find_attractors(epg, vn)
    tgt = orc["attractors"]
    recovered = int(sum(np.min(np.linalg.norm(tgt[i] - attr, axis=1)) < BASIN_EPS
                        for i in range(len(tgt)))) if len(attr) else 0
    return dict(basin_joint=float((d < BASIN_EPS).mean()),
                basin_A=float((d[~side] < BASIN_EPS).mean()),
                basin_B=float((d[side] < BASIN_EPS).mean()),
                n_attractors=int(len(attr)),
                sites_recovered=recovered, sites_total=int(len(tgt)))


# ---------- main ----------

def main():
    t0 = time.time()
    rng = np.random.default_rng(7)
    xh, labels = make_world8()
    xA, xB = xh[labels < 4], xh[labels >= 4]

    print("training teachers A, B (denoising), B_mem (memorization) ...")
    A = train_teacher(xA, noise=0.1, seed=0)
    B = train_teacher(xB, noise=0.1, seed=0)
    B_mem = train_teacher(xB, noise=0.0, seed=0)

    frame = UnionFrame(A, B, xh, labels)
    orc = frame.oracle(rng)
    print(f"union oracle: {len(orc['attractors'])} attractors, ceilings "
          f"joint {orc['ceil_joint']:.2f} A {orc['ceil_A']:.2f} B {orc['ceil_B']:.2f}")

    # --- union: read carrier vs their SGD student (0.82 mean, 2.048M evals) ---
    car = LatticeCarrier(frame.field_np, frame.lo, frame.hi)
    m = eval_carrier(car, frame, orc)
    print(f"read carrier ({car.n_samples} samples): joint {m['basin_joint']:.2f} "
          f"A {m['basin_A']:.2f} B {m['basin_B']:.2f} "
          f"sites {m['sites_recovered']}/{m['sites_total']}")
    out = dict(oracle=dict(ceil_joint=orc["ceil_joint"], ceil_A=orc["ceil_A"],
                           ceil_B=orc["ceil_B"], n_attractors=int(len(orc["attractors"]))),
               carrier_samples=car.n_samples, union=m)

    # --- lineage: resample lookup field on a jittered lattice, 5 generations ---
    lineage = [m]
    cur = car
    for g in range(2, GENS + 1):
        jit = rng.random(2).astype(np.float32)  # origin shift in [0, h)
        cur = LatticeCarrier(cur.lookup, frame.lo, frame.hi, origin_jitter=jit)
        e = eval_carrier(cur, frame, orc)
        lineage.append(e)
        print(f"  lineage gen {g}: joint {e['basin_joint']:.2f} "
              f"A {e['basin_A']:.2f} B {e['basin_B']:.2f}")
    out["lineage"] = lineage

    # --- onion filter: does disordered knowledge survive the read channel? ---
    frame_c = UnionFrame(A, B_mem, xh, labels)
    orc_c = frame_c.oracle(rng)
    print(f"contaminated oracle: ceilings joint {orc_c['ceil_joint']:.2f} "
          f"A {orc_c['ceil_A']:.2f} Bmem {orc_c['ceil_B']:.2f}")
    car_c = LatticeCarrier(frame_c.field_np, frame_c.lo, frame_c.hi)
    lineage_c = [eval_carrier(car_c, frame_c, orc_c)]
    cur = car_c
    for g in range(2, GENS + 1):
        jit = rng.random(2).astype(np.float32)
        cur = LatticeCarrier(cur.lookup, frame_c.lo, frame_c.hi, origin_jitter=jit)
        lineage_c.append(eval_carrier(cur, frame_c, orc_c))
    for g, e in enumerate(lineage_c, 1):
        print(f"  contaminated gen {g}: joint {e['basin_joint']:.2f} "
              f"A {e['basin_A']:.2f} Bmem {e['basin_B']:.2f}")
    out["lineage_contaminated"] = lineage_c

    OUT.write_text(json.dumps(out, indent=2))
    print(f"\nwrote {OUT} ({time.time() - t0:.0f}s)")
    print("\ntheir gi.json for comparison: union student 0.82 (ceiling 0.98), "
          "lineage 0.91->0.83, contaminated B-side 0.24->0.19")


if __name__ == "__main__":
    main()

spike_2d_results.json

{
  "oracle": {
    "ceil_joint": 0.9725,
    "ceil_A": 0.9692307692307692,
    "ceil_B": 0.975609756097561,
    "n_attractors": 8
  },
  "carrier_samples": 262144,
  "union": {
    "basin_joint": 1.0,
    "basin_A": 1.0,
    "basin_B": 1.0,
    "n_attractors": 8,
    "sites_recovered": 8,
    "sites_total": 8
  },
  "lineage": [
    {
      "basin_joint": 1.0,
      "basin_A": 1.0,
      "basin_B": 1.0,
      "n_attractors": 8,
      "sites_recovered": 8,
      "sites_total": 8
    },
    {
      "basin_joint": 0.995,
      "basin_A": 1.0,
      "basin_B": 0.9902439024390244,
      "n_attractors": 8,
      "sites_recovered": 8,
      "sites_total": 8
    },
    {
      "basin_joint": 0.99,
      "basin_A": 1.0,
      "basin_B": 0.9804878048780488,
      "n_attractors": 8,
      "sites_recovered": 8,
      "sites_total": 8
    },
    {
      "basin_joint": 0.99,
      "basin_A": 1.0,
      "basin_B": 0.9804878048780488,
      "n_attractors": 8,
      "sites_recovered": 8,
      "sites_total": 8
    },
    {
      "basin_joint": 0.99,
      "basin_A": 1.0,
      "basin_B": 0.9804878048780488,
      "n_attractors": 8,
      "sites_recovered": 8,
      "sites_total": 8
    }
  ],
  "lineage_contaminated": [
    {
      "basin_joint": 0.8425,
      "basin_A": 1.0,
      "basin_B": 0.7028301886792453,
      "n_attractors": 19,
      "sites_recovered": 18,
      "sites_total": 19
    },
    {
      "basin_joint": 0.8425,
      "basin_A": 1.0,
      "basin_B": 0.7028301886792453,
      "n_attractors": 20,
      "sites_recovered": 18,
      "sites_total": 19
    },
    {
      "basin_joint": 0.8375,
      "basin_A": 1.0,
      "basin_B": 0.6933962264150944,
      "n_attractors": 19,
      "sites_recovered": 18,
      "sites_total": 19
    },
    {
      "basin_joint": 0.8325,
      "basin_A": 0.9946808510638298,
      "basin_B": 0.6886792452830188,
      "n_attractors": 20,
      "sites_recovered": 18,
      "sites_total": 19
    },
    {
      "basin_joint": 0.835,
      "basin_A": 0.9946808510638298,
      "basin_B": 0.6933962264150944,
      "n_attractors": 19,
      "sites_recovered": 18,
      "sites_total": 19
    }
  ]
}

spike_alpha.py

"""alpha(noise) sweep: is the memorization->generalization transition a
fractal->smooth basin-boundary transition?

F-73 observed (direction only): ordered teachers alpha 0.73-1.04, sub-ordered
0.56-0.70. Here the dedicated test at d=8 (transition known to sit between
noise 0.1 and 0.2 from the frontier grid), noise in {0, .05, .1, .15, .2, .3},
3 teacher seeds per level, same world.

Pre-registered criteria (BEFORE running):
  (a) mean alpha at noise=0.0 (memorization) < 0.7;
  (b) mean alpha at noise=0.3 (deep ordered) > 0.85;
  (c) the noise level where mean alpha crosses 0.85 lies within one grid step
      of where mean conv crosses 0.9 (co-located transitions);
  (d) informative either way: if alpha transitions at LOWER noise than conv,
      alpha is the more sensitive order parameter; if alpha stays flat while
      conv snaps, the fractal-boundary story is DEAD.

Run (one seed per job): uv run --no-project --with mlx --with numpy python spike_alpha.py <seed>
Appends to spike_alpha_results.jsonl
"""

import json
import sys
import time

import mlx.core as mx
import numpy as np

from spike_dscaling import (HERE, ITER, N_PROBES, EPS_FRACS, field,
                            find_attractors, iterate, make_world, train_teacher)

OUT = HERE / "spike_alpha_results.jsonl"
D = 8
NOISES = (0.0, 0.05, 0.1, 0.15, 0.2, 0.3)


def probe(noise, seed):
    t0 = time.time()
    rng = np.random.default_rng(D)          # same world for all configs
    xh = make_world(D, rng)
    teacher = train_teacher(xh, D, seed=seed, noise=noise)
    rng = np.random.default_rng(10_000 + seed)  # fresh probe rng per seed

    z_tr = np.array(teacher.enc(mx.array(xh)))
    mu, sd = z_tr.mean(0), z_tr.std(0)
    scale = float(np.linalg.norm(sd))
    lo = (mu - 8 * sd).astype(np.float32)
    hi = (mu + 8 * sd).astype(np.float32)

    starts = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), 256)] + 0.5 * sd * rng.standard_normal((256, D)),
        mu + 2.0 * sd * rng.standard_normal((256, D)),
    ]).astype(np.float32)
    ep_s = iterate(teacher, starts, ITER, lo, hi)
    vn = np.linalg.norm(np.array(field(teacher, mx.array(ep_s))), axis=1)
    attrs = find_attractors(ep_s, vn, vtol=5e-3 * scale, merge_r=0.1 * scale)
    conv = float((vn < 5e-3 * scale).mean())
    if len(attrs) > 1:
        dm = np.linalg.norm(attrs[:, None] - attrs[None, :], axis=-1)
        basin_eps = 0.25 * float(np.median(dm[dm > 0]))
    else:
        basin_eps = 0.25 * scale

    probes = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), N_PROBES // 2)]
        + 0.5 * sd * rng.standard_normal((N_PROBES // 2, D)),
        mu + 2.0 * sd * rng.standard_normal((N_PROBES // 2, D)),
    ]).astype(np.float32)
    ep_ref = iterate(teacher, probes, ITER, lo, hi)
    ep_pert = iterate(teacher, probes + (0.05 * scale)
                      * rng.standard_normal(probes.shape).astype(np.float32), ITER, lo, hi)
    ceiling = float((np.linalg.norm(ep_pert - ep_ref, axis=1) < basin_eps).mean())

    fr = []
    for f in EPS_FRACS:
        eps = f * basin_eps
        ep_e = iterate(teacher, probes + eps
                       * rng.standard_normal(probes.shape).astype(np.float32), ITER, lo, hi)
        fr.append(float((np.linalg.norm(ep_e - ep_ref, axis=1) > basin_eps).mean()))
    mask = [i for i, v in enumerate(fr) if v > 0]
    if len(mask) >= 3:
        xs = np.log([EPS_FRACS[i] for i in mask])
        ys = np.log([fr[i] for i in mask])
        alpha = float(np.polyfit(xs, ys, 1)[0])
    else:
        alpha = float("nan")

    rec = dict(d=D, noise=noise, seed=seed, n_attractors=int(len(attrs)),
               conv=conv, ceiling=ceiling, alpha=alpha, f_eps=fr,
               seconds=round(time.time() - t0))
    with OUT.open("a") as fh:
        fh.write(json.dumps(rec) + "\n")
    print(f"noise={noise} seed={seed}: attrs {len(attrs)} conv {conv:.2f} "
          f"ceil {ceiling:.2f} alpha {alpha:.2f} [{rec['seconds']}s]")


if __name__ == "__main__":
    seed = int(sys.argv[1])
    for noise in NOISES:
        probe(noise, seed)

spike_alpha_results.jsonl

{"d": 8, "noise": 0.0, "seed": 0, "n_attractors": 0, "conv": 0.0, "ceiling": 0.6640625, "alpha": 0.5697017566380679, "f_eps": [0.078125, 0.109375, 0.17578125, 0.2421875, 0.37109375, 0.55859375], "seconds": 16}
{"d": 8, "noise": 0.0, "seed": 1, "n_attractors": 3, "conv": 0.03515625, "ceiling": 0.71484375, "alpha": 0.5280631012561426, "f_eps": [0.1328125, 0.14453125, 0.2890625, 0.41015625, 0.60546875, 0.6796875], "seconds": 14}
{"d": 8, "noise": 0.0, "seed": 2, "n_attractors": 1, "conv": 0.04296875, "ceiling": 0.73828125, "alpha": 0.6844503847547813, "f_eps": [0.05078125, 0.05859375, 0.13671875, 0.20703125, 0.33203125, 0.45703125], "seconds": 16}
{"d": 8, "noise": 0.05, "seed": 1, "n_attractors": 0, "conv": 0.0, "ceiling": 0.79296875, "alpha": 0.6345902892578807, "f_eps": [0.03515625, 0.06640625, 0.08984375, 0.14453125, 0.21875, 0.33984375], "seconds": 15}
{"d": 8, "noise": 0.05, "seed": 0, "n_attractors": 5, "conv": 0.025390625, "ceiling": 0.79296875, "alpha": 0.6375153522736858, "f_eps": [0.07421875, 0.14453125, 0.25390625, 0.390625, 0.53125, 0.6875], "seconds": 17}
{"d": 8, "noise": 0.05, "seed": 2, "n_attractors": 3, "conv": 0.12890625, "ceiling": 0.7265625, "alpha": 0.5447994650729416, "f_eps": [0.09765625, 0.17578125, 0.26171875, 0.37890625, 0.4921875, 0.6875], "seconds": 18}
{"d": 8, "noise": 0.1, "seed": 1, "n_attractors": 5, "conv": 0.150390625, "ceiling": 0.7109375, "alpha": 0.6973353951228579, "f_eps": [0.0390625, 0.09375, 0.1328125, 0.19140625, 0.37109375, 0.46875], "seconds": 14}
{"d": 8, "noise": 0.1, "seed": 0, "n_attractors": 13, "conv": 0.44140625, "ceiling": 0.765625, "alpha": 0.6777192206306923, "f_eps": [0.04296875, 0.1015625, 0.15234375, 0.203125, 0.34765625, 0.51953125], "seconds": 17}
{"d": 8, "noise": 0.1, "seed": 2, "n_attractors": 11, "conv": 0.560546875, "ceiling": 0.7734375, "alpha": 0.6973173148330541, "f_eps": [0.01953125, 0.0390625, 0.04296875, 0.07421875, 0.12109375, 0.26171875], "seconds": 16}
{"d": 8, "noise": 0.15, "seed": 1, "n_attractors": 4, "conv": 1.0, "ceiling": 0.9609375, "alpha": 1.1816321722514036, "f_eps": [0.00390625, 0.00390625, 0.015625, 0.03515625, 0.06640625, 0.1875], "seconds": 15}
{"d": 8, "noise": 0.15, "seed": 0, "n_attractors": 4, "conv": 0.998046875, "ceiling": 0.9453125, "alpha": 0.9408143304947397, "f_eps": [0.0078125, 0.01171875, 0.03515625, 0.07421875, 0.078125, 0.20703125], "seconds": 19}
{"d": 8, "noise": 0.15, "seed": 2, "n_attractors": 7, "conv": 1.0, "ceiling": 0.9609375, "alpha": 0.9867104779116247, "f_eps": [0.0078125, 0.01171875, 0.0234375, 0.07421875, 0.10546875, 0.19921875], "seconds": 18}
{"d": 8, "noise": 0.2, "seed": 1, "n_attractors": 6, "conv": 1.0, "ceiling": 0.96875, "alpha": 1.051078857802584, "f_eps": [0.00390625, 0.01171875, 0.01171875, 0.0234375, 0.09375, 0.16015625], "seconds": 16}
{"d": 8, "noise": 0.2, "seed": 0, "n_attractors": 4, "conv": 1.0, "ceiling": 0.96875, "alpha": 1.2673750739438054, "f_eps": [0.0, 0.00390625, 0.015625, 0.0390625, 0.09375, 0.12890625], "seconds": 18}
{"d": 8, "noise": 0.2, "seed": 2, "n_attractors": 4, "conv": 1.0, "ceiling": 0.9765625, "alpha": 1.0499832163298541, "f_eps": [0.00390625, 0.0078125, 0.01171875, 0.0234375, 0.0703125, 0.1484375], "seconds": 18}
{"d": 8, "noise": 0.3, "seed": 1, "n_attractors": 7, "conv": 1.0, "ceiling": 0.9765625, "alpha": 0.8525921427926573, "f_eps": [0.0078125, 0.0078125, 0.015625, 0.02734375, 0.06640625, 0.12109375], "seconds": 16}
{"d": 8, "noise": 0.3, "seed": 0, "n_attractors": 8, "conv": 1.0, "ceiling": 0.96484375, "alpha": 1.0865992291658317, "f_eps": [0.00390625, 0.00390625, 0.01953125, 0.0234375, 0.05859375, 0.14453125], "seconds": 19}
{"d": 8, "noise": 0.3, "seed": 2, "n_attractors": 6, "conv": 1.0, "ceiling": 0.93359375, "alpha": 0.9140233316505503, "f_eps": [0.0, 0.01171875, 0.0390625, 0.03125, 0.09375, 0.1796875], "seconds": 18}

spike_curve.py

"""Closing runs: (1) GP carrier with the honest sigma (0.5 x median — selection
by NMSE alone picked the oversmoothed kernel, our own mean-residual trap);
(2) the basin(N) coverage curve for the NN carrier: 65k / 262k / 4M
(1M = 0.32 known from spike_mnist.py). Rising without saturation => the 16D
wall is budget; plateau below ceiling => intrinsic.

Run: uv run --no-project --with mlx --with numpy python spike_curve.py
"""

import json
import time

import mlx.core as mx
import numpy as np

from spike_mnist import (HERE, ITER_STEPS, build_table, knn_field,
                         train_or_load_teacher)
from spike_wiener import basin_metrics, field_metrics, gp_iterate, gp_predict, rbf_solve

OUT = HERE / "spike_curve_results.json"
M_GP = 16384
SIGMA_MULT = 0.5
NS = (65536, 262144, 4_000_000)


def nn_run(teacher, meta, Z, V):
    Zsq = mx.sum(Z * Z, axis=1)
    mx.eval(Zsq)
    z = mx.array(meta["probes"])
    lo, hi = mx.array(meta["clip_lo"]), mx.array(meta["clip_hi"])
    for _ in range(ITER_STEPS):
        z = z + knn_field(Z, V, Zsq, z, 1)
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    return basin_metrics(teacher, meta, np.array(z))


def main():
    t0 = time.time()
    teacher, meta = train_or_load_teacher()
    print(f"teacher: ceiling {float(meta['self_agree']):.2f}, "
          f"basin_eps {float(meta['basin_eps']):.3f}")
    results = []

    # --- (1) GP at honest sigma ---
    Z, V = build_table(teacher, meta)  # 1M, cached seed
    Zn, Vn = np.array(Z), np.array(V)
    rng = np.random.default_rng(42)
    perm = rng.permutation(len(Zn))
    fin_idx = perm[4096:4096 + M_GP]   # same split as spike_wiener
    ZM, VM = Zn[fin_idx], Vn[fin_idx]
    d_med = np.median(np.linalg.norm(
        ZM[rng.integers(0, M_GP, 2000)] - ZM[rng.integers(0, M_GP, 2000)], axis=1))
    sigma = SIGMA_MULT * d_med
    alpha = rbf_solve(ZM, VM, sigma)
    ZM_m = mx.array(ZM); alpha_m = mx.array(alpha.astype(np.float32))
    Zsq_m = mx.sum(ZM_m * ZM_m, axis=1)
    mx.eval(ZM_m, alpha_m, Zsq_m)
    Vh = np.array(gp_predict(ZM_m, Zsq_m, alpha_m, sigma, mx.array(meta["zeval"])))
    nmse, cos = field_metrics(Vh, meta["Vt_eval"])
    ep = gp_iterate(ZM_m, Zsq_m, alpha_m, sigma, meta["probes"],
                    meta["clip_lo"], meta["clip_hi"])
    basin, dec, med = basin_metrics(teacher, meta, ep)
    print(f"GP M={M_GP} sigma=0.5x: basin {basin:.2f}  dec {dec:.2f}  "
          f"median-err {med:.2f}  nmse {nmse:.3f}  cos {cos:.2f}")
    results.append(dict(carrier=f"gp_M{M_GP}_s0.5", basin_agreement=basin,
                        dec_agree=dec, endpoint_median=med,
                        field_nmse=nmse, field_cos=cos))

    # --- (2) basin(N) curve for NN k=1 ---
    for n in NS:
        t1 = time.time()
        if n <= len(Zn):
            Zt, Vt_ = mx.array(Zn[perm[:n]]), mx.array(Vn[perm[:n]])
        else:
            Zt, Vt_ = build_table(teacher, meta, n=n, seed=2)
        basin, dec, med = nn_run(teacher, meta, Zt, Vt_)
        print(f"NN k=1 N={n}: basin {basin:.2f}  dec {dec:.2f}  "
              f"median-err {med:.2f}  [{time.time()-t1:.0f}s]")
        results.append(dict(carrier=f"nn1_N{n}", basin_agreement=basin,
                            dec_agree=dec, endpoint_median=med))
        del Zt, Vt_

    OUT.write_text(json.dumps(dict(ceiling=float(meta["self_agree"]),
                                   known=dict(nn1_N1000000=0.32),
                                   results=results), indent=2))
    print(f"\nknown: NN-1M 0.32 | ceiling {float(meta['self_agree']):.2f} "
          f"[{time.time()-t0:.0f}s total]")


if __name__ == "__main__":
    main()

spike_curve_results.json

{
  "ceiling": 0.703125,
  "known": {
    "nn1_N1000000": 0.32
  },
  "results": [
    {
      "carrier": "gp_M16384_s0.5",
      "basin_agreement": 0.484375,
      "dec_agree": 0.47265625,
      "endpoint_median": 3.0590169429779053,
      "field_nmse": 0.12240374088287354,
      "field_cos": 0.8978049159049988
    },
    {
      "carrier": "nn1_N65536",
      "basin_agreement": 0.0234375,
      "dec_agree": 0.015625,
      "endpoint_median": 8.348264694213867
    },
    {
      "carrier": "nn1_N262144",
      "basin_agreement": 0.18359375,
      "dec_agree": 0.1171875,
      "endpoint_median": 4.842211723327637
    },
    {
      "carrier": "nn1_N4000000",
      "basin_agreement": 0.484375,
      "dec_agree": 0.44140625,
      "endpoint_median": 2.630967140197754
    }
  ]
}

spike_dscaling.py

"""d-scaling of the read/GP carriers + basin-boundary fractality.

Questions (F-72 open item "d-scaling"):
  1. Does the correlation-structure gain (GP-16k vs NN at matched/larger N)
     shrink with latent dimension d? d in {2,4,8,16,24,32}.
  2. Is anything special around d=24 (the last Viazovska/universal-optimality
     dimension)? Here only as an empirical smoothness check of gain(d).
  3. Fractality: uncertainty exponent alpha of the TEACHER's basin boundaries
     (Grebogi/McDonald/Ott/Yorke): fraction of probes whose endpoint changes
     under eps-perturbation scales as f(eps) ~ eps^alpha; boundary box
     dimension D_b = d - alpha. Smooth boundary => alpha ~ 1; fractal => <1.
     Coverage cost of basin transfer should track alpha/d, not d alone.

World per d: 8 Gaussian blobs, centers r=1.7 * random unit vectors in R^d,
std 0.32 (separation ratio ~constant across d), lifted tanh-nonlinearly to
R^64. Denoising AE teacher 64-128-128-d, noise 0.1, 4000 steps. Ceiling,
conv_frac, basin_eps computed per d (basin_eps = 0.25 * median
inter-attractor distance). All basins vs the teacher's own endpoints.

Run: uv run --no-project --with mlx --with numpy python spike_dscaling.py 2 4 8 16
     uv run --no-project --with mlx --with numpy python spike_dscaling.py 24 32
Appends per-dim results to spike_dscaling_results.jsonl
"""

import json
import sys
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np

HERE = Path(__file__).parent
OUT = HERE / "spike_dscaling_results.jsonl"

D_IN = 64
N_BLOBS = 8
R_CENTERS = 1.7
STD = 0.32
N_PER = 500
NOISE = 0.1
STEPS, BS, LR = 4000, 256, 1e-3
ITER = 800
N_PROBES = 256
NS = (16384, 262144, 1_000_000)
M_GP = 16384
SIGMA_MULT = 0.5
JITTER = 1e-8
EPS_FRACS = (1 / 64, 1 / 32, 1 / 16, 1 / 8, 1 / 4, 1 / 2)  # x basin_eps


class MLP(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.layers = [nn.Linear(a, b) for a, b in zip(dims[:-1], dims[1:])]

    def __call__(self, x):
        n = len(self.layers)
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < n - 1:
                x = nn.gelu(x)
        return x


class AE(nn.Module):
    def __init__(self, d_lat):
        super().__init__()
        self.enc = MLP([D_IN, 128, 128, d_lat])
        self.dec = MLP([d_lat, 128, 128, D_IN])

    def __call__(self, x):
        return self.dec(self.enc(x))


def field(model, z):
    return model.enc(model.dec(z)) - z


def iterate(model, z0, steps, lo, hi):
    z = mx.array(z0)
    lo, hi = mx.array(lo), mx.array(hi)
    for i in range(steps):
        z = model.enc(model.dec(z))
        z = mx.maximum(mx.minimum(z, hi), lo)
        if i % 50 == 0:
            mx.eval(z)
    mx.eval(z)
    return np.array(z)


def find_attractors(endpoints, vnorms, vtol, merge_r):
    pts = endpoints[vnorms < vtol]
    centers = []
    for p in pts:
        if all(np.linalg.norm(p - c) > merge_r for c in centers):
            centers.append(p)
    return np.array(centers) if centers else np.zeros((0, endpoints.shape[1]))


def make_world(d, rng):
    u = rng.normal(size=(N_BLOBS, d))
    centers = R_CENTERS * u / np.linalg.norm(u, axis=1, keepdims=True)
    labels = np.repeat(np.arange(N_BLOBS), N_PER)
    x = centers[labels] + STD * rng.normal(size=(len(labels), d))
    W1 = rng.normal(size=(d, D_IN)) / np.sqrt(d)
    b1 = 0.1 * rng.normal(size=(D_IN,))
    W2 = rng.normal(size=(D_IN, D_IN)) / np.sqrt(D_IN)
    return (np.tanh(x @ W1 + b1) @ W2).astype(np.float32)


def train_teacher(xh, d, seed=0, noise=NOISE):
    mx.random.seed(seed)
    model = AE(d)
    mx.eval(model.parameters())
    opt = optim.Adam(learning_rate=LR)
    X = mx.array(xh)
    rng = np.random.default_rng(seed)

    def loss_fn(m, xb, xn):
        return mx.mean((m(xn) - xb) ** 2)

    lg = nn.value_and_grad(model, loss_fn)
    for _ in range(STEPS):
        idx = mx.array(rng.integers(0, len(xh), BS))
        xb = X[idx]
        xn = xb + noise * mx.array(rng.standard_normal((BS, D_IN)).astype(np.float32))
        _, g = lg(model, xb, xn)
        opt.update(model, g)
        mx.eval(model.parameters(), opt.state)
    return model


def build_table(teacher, mu, sd, d, n, seed=1):
    rng = np.random.default_rng(1000 + seed)
    zs, total = [], 0
    bs = 8192
    while total < n:
        half = bs // 2
        broad = (mu + 2.0 * sd * rng.standard_normal((half, d))).astype(np.float32)
        core = mx.array((mu + 1.5 * sd * rng.standard_normal((bs - half, d))).astype(np.float32))
        for _ in range(int(rng.integers(2, 25))):
            core = teacher.enc(teacher.dec(core))
        mx.eval(core)
        zs.append(np.concatenate([broad, np.array(core)]))
        total += bs
    Z = np.concatenate(zs)[:n]
    V = np.empty_like(Z)
    for i in range(0, n, 8192):
        V[i:i + 8192] = np.array(field(teacher, mx.array(Z[i:i + 8192])))
    return Z, V


def knn_iterate(Zt, Vt, z0, lo, hi, steps=ITER):
    Zsq = mx.sum(Zt * Zt, axis=1)
    z = mx.array(z0)
    lo, hi = mx.array(lo), mx.array(hi)
    for _ in range(steps):
        d2 = Zsq[None, :] - 2 * (z @ Zt.T)
        z = z + Vt[mx.argmin(d2, axis=1)]
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    return np.array(z)


def gp_carrier(ZM, VM, sigma):
    Z = ZM.astype(np.float64)
    sq = (Z ** 2).sum(1)
    d2 = sq[:, None] - 2 * Z @ Z.T + sq[None, :]
    K = np.exp(-np.maximum(d2, 0) / (2 * sigma ** 2))
    K[np.diag_indices_from(K)] += JITTER * K.shape[0]
    alpha = np.linalg.solve(K, VM.astype(np.float64))
    return mx.array(ZM), mx.array(alpha.astype(np.float32))


def gp_iterate(ZM_m, alpha_m, sigma, z0, lo, hi, steps=ITER):
    Zsq = mx.sum(ZM_m * ZM_m, axis=1)
    z = mx.array(z0)
    lo, hi = mx.array(lo), mx.array(hi)
    for _ in range(steps):
        d2 = Zsq[None, :] - 2 * (z @ ZM_m.T) + mx.sum(z * z, axis=1)[:, None]
        z = z + mx.exp(-mx.maximum(d2, 0.0) / (2 * sigma ** 2)) @ alpha_m
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    return np.array(z)


def run_dim(d, noise=NOISE):
    t0 = time.time()
    rng = np.random.default_rng(d)
    xh = make_world(d, rng)
    teacher = train_teacher(xh, d, noise=noise)

    z_tr = np.array(teacher.enc(mx.array(xh)))
    mu, sd = z_tr.mean(0), z_tr.std(0)
    scale = float(np.linalg.norm(sd))
    lo = (mu - 8 * sd).astype(np.float32)
    hi = (mu + 8 * sd).astype(np.float32)

    starts = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), 256)] + 0.5 * sd * rng.standard_normal((256, d)),
        mu + 2.0 * sd * rng.standard_normal((256, d)),
    ]).astype(np.float32)
    ep_s = iterate(teacher, starts, ITER, lo, hi)
    vn = np.linalg.norm(np.array(field(teacher, mx.array(ep_s))), axis=1)
    attrs = find_attractors(ep_s, vn, vtol=5e-3 * scale, merge_r=0.1 * scale)
    conv = float((vn < 5e-3 * scale).mean())
    if len(attrs) > 1:
        dm = np.linalg.norm(attrs[:, None] - attrs[None, :], axis=-1)
        basin_eps = 0.25 * float(np.median(dm[dm > 0]))
    else:
        basin_eps = 0.25 * scale

    probes = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), N_PROBES // 2)]
        + 0.5 * sd * rng.standard_normal((N_PROBES // 2, d)),
        mu + 2.0 * sd * rng.standard_normal((N_PROBES // 2, d)),
    ]).astype(np.float32)
    ep_ref = iterate(teacher, probes, ITER, lo, hi)
    ep_pert = iterate(teacher, probes + (0.05 * scale)
                      * rng.standard_normal(probes.shape).astype(np.float32), ITER, lo, hi)
    ceiling = float((np.linalg.norm(ep_pert - ep_ref, axis=1) < basin_eps).mean())

    # --- fractality: uncertainty exponent of teacher basin boundaries ---
    fr = []
    for f in EPS_FRACS:
        eps = f * basin_eps
        ep_e = iterate(teacher, probes + eps
                       * rng.standard_normal(probes.shape).astype(np.float32), ITER, lo, hi)
        fr.append(float((np.linalg.norm(ep_e - ep_ref, axis=1) > basin_eps).mean()))
    mask = [i for i, v in enumerate(fr) if v > 0]
    if len(mask) >= 3:
        xs = np.log([EPS_FRACS[i] for i in mask])
        ys = np.log([fr[i] for i in mask])
        alpha = float(np.polyfit(xs, ys, 1)[0])
    else:
        alpha = float("nan")

    # --- carriers ---
    Z, V = build_table(teacher, mu, sd, d, max(NS))
    res_nn = {}
    for n in NS:
        ep = knn_iterate(mx.array(Z[:n]), mx.array(V[:n]), probes, lo, hi)
        res_nn[n] = float((np.linalg.norm(ep - ep_ref, axis=1) < basin_eps).mean())

    perm = np.random.default_rng(42).permutation(len(Z))[:M_GP]
    ZM, VM = Z[perm], V[perm]
    d_med = float(np.median(np.linalg.norm(
        ZM[rng.integers(0, M_GP, 2000)] - ZM[rng.integers(0, M_GP, 2000)], axis=1)))
    sigma = SIGMA_MULT * d_med
    ZM_m, alpha_m = gp_carrier(ZM, VM, sigma)
    ep = gp_iterate(ZM_m, alpha_m, sigma, probes, lo, hi)
    basin_gp = float((np.linalg.norm(ep - ep_ref, axis=1) < basin_eps).mean())

    rec = dict(d=d, noise=noise, n_attractors=int(len(attrs)), conv=conv, ceiling=ceiling,
               basin_eps=basin_eps, scale=scale, alpha=alpha, f_eps=fr,
               nn={str(k): v for k, v in res_nn.items()}, gp_16k=basin_gp,
               seconds=round(time.time() - t0))
    with OUT.open("a") as fh:
        fh.write(json.dumps(rec) + "\n")
    print(f"d={d}: attrs {len(attrs)} conv {conv:.2f} ceil {ceiling:.2f} "
          f"alpha {alpha:.2f} | NN {res_nn[NS[0]]:.2f}/{res_nn[NS[1]]:.2f}/"
          f"{res_nn[NS[2]]:.2f} | GP-16k {basin_gp:.2f} [{rec['seconds']}s]")


if __name__ == "__main__":
    for arg in sys.argv[1:]:
        if ":" in arg:
            ds, ns = arg.split(":")
            run_dim(int(ds), noise=float(ns))
        else:
            run_dim(int(arg))

spike_dscaling_results.jsonl

{"d": 2, "n_attractors": 2, "conv": 0.9609375, "ceiling": 0.9765625, "basin_eps": 0.9323017597198486, "scale": 1.5630770921707153, "alpha": 0.8738169732842052, "f_eps": [0.0078125, 0.00390625, 0.00390625, 0.03125, 0.0390625, 0.08984375], "nn": {"16384": 0.953125, "262144": 0.953125, "1000000": 0.953125}, "gp_16k": 0.9765625, "seconds": 144}
{"d": 24, "n_attractors": 2, "conv": 0.05859375, "ceiling": 0.75, "basin_eps": 1.2125375270843506, "scale": 1.6877076625823975, "alpha": 0.6965928751616873, "f_eps": [0.05859375, 0.140625, 0.2109375, 0.3828125, 0.4765625, 0.734375], "nn": {"16384": 0.0, "262144": 0.0, "1000000": 0.0}, "gp_16k": 0.0234375, "seconds": 147}
{"d": 4, "n_attractors": 3, "conv": 0.98828125, "ceiling": 0.9609375, "basin_eps": 0.6162179112434387, "scale": 2.109988212585449, "alpha": 0.5631492364423594, "f_eps": [0.0, 0.01953125, 0.0234375, 0.0234375, 0.04296875, 0.1015625], "nn": {"16384": 0.93359375, "262144": 0.96484375, "1000000": 0.984375}, "gp_16k": 0.5625, "seconds": 157}
{"d": 32, "n_attractors": 7, "conv": 0.39453125, "ceiling": 0.7890625, "basin_eps": 0.15200293064117432, "scale": 1.5878543853759766, "alpha": 0.7009823095079734, "f_eps": [0.01953125, 0.03515625, 0.0625, 0.08984375, 0.15625, 0.22265625], "nn": {"16384": 0.0, "262144": 0.0, "1000000": 0.0}, "gp_16k": 0.00390625, "seconds": 157}
{"d": 8, "n_attractors": 15, "conv": 0.4375, "ceiling": 0.79296875, "basin_eps": 0.9196683764457703, "scale": 2.4838883876800537, "alpha": 0.5861776112612977, "f_eps": [0.05859375, 0.109375, 0.13671875, 0.203125, 0.3203125, 0.48828125], "nn": {"16384": 0.015625, "262144": 0.37890625, "1000000": 0.4453125}, "gp_16k": 0.55078125, "seconds": 97}
{"d": 16, "n_attractors": 8, "conv": 0.494140625, "ceiling": 0.84765625, "basin_eps": 1.6302363872528076, "scale": 2.1428210735321045, "alpha": 0.6907876805490315, "f_eps": [0.0625, 0.0546875, 0.140625, 0.2265625, 0.37890625, 0.5078125], "nn": {"16384": 0.0, "262144": 0.00390625, "1000000": 0.26953125}, "gp_16k": 0.58203125, "seconds": 136}
{"d": 8, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.97265625, "basin_eps": 0.9710250496864319, "scale": 2.6066699028015137, "alpha": 0.727158124277616, "f_eps": [0.01171875, 0.015625, 0.01171875, 0.03125, 0.05859375, 0.1484375], "nn": {"16384": 0.80078125, "262144": 0.9609375, "1000000": 0.96875}, "gp_16k": 0.98828125, "seconds": 160}
{"d": 24, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.91015625, "basin_eps": 1.0080623626708984, "scale": 2.7132787704467773, "alpha": 0.7588185869465107, "f_eps": [0.01953125, 0.015625, 0.05859375, 0.06640625, 0.1015625, 0.24609375], "nn": {"16384": 0.00390625, "262144": 0.13671875, "1000000": 0.1796875}, "gp_16k": 0.90625, "seconds": 167}
{"d": 16, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.95703125, "basin_eps": 1.0177536010742188, "scale": 2.6999692916870117, "alpha": 1.0379641440300558, "f_eps": [0.00390625, 0.01171875, 0.0234375, 0.03125, 0.07421875, 0.1875], "nn": {"16384": 0.1640625, "262144": 0.32421875, "1000000": 0.6171875}, "gp_16k": 0.95703125, "seconds": 163}
{"d": 32, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.91796875, "basin_eps": 0.9295673370361328, "scale": 2.519338369369507, "alpha": 0.9928596952474479, "f_eps": [0.01171875, 0.01171875, 0.0390625, 0.08203125, 0.140625, 0.28125], "nn": {"16384": 0.0, "262144": 0.03125, "1000000": 0.04296875}, "gp_16k": 0.85546875, "seconds": 170}

spike_frontier.py

"""Orderability frontier: does the contractive (ordered) phase shrink with
latent dimension — and is there a d beyond which no denoising level orders
the teacher?

The d-scaling spike showed teachers at fixed noise 0.1 fall out of the
ordered regime for d >= 8 (conv 0.44-0.06), making carrier comparison
ill-posed there. Here: teacher-only grid over (d, noise), reporting
attractor count / conv_frac / ceiling. Ordered = conv >= 0.9.

Run: uv run --no-project --with mlx --with numpy python spike_frontier.py 8 16
     uv run --no-project --with mlx --with numpy python spike_frontier.py 24 32
Appends to spike_frontier_results.jsonl
"""

import json
import sys
import time

import mlx.core as mx
import numpy as np

from spike_dscaling import (HERE, ITER, N_PROBES, AE, field, find_attractors,
                            iterate, make_world, train_teacher)

OUT = HERE / "spike_frontier_results.jsonl"
NOISES = (0.2, 0.35, 0.5)


def probe_teacher(d, noise):
    t0 = time.time()
    rng = np.random.default_rng(d)
    xh = make_world(d, rng)
    teacher = train_teacher(xh, d, noise=noise)

    z_tr = np.array(teacher.enc(mx.array(xh)))
    mu, sd = z_tr.mean(0), z_tr.std(0)
    scale = float(np.linalg.norm(sd))
    lo = (mu - 8 * sd).astype(np.float32)
    hi = (mu + 8 * sd).astype(np.float32)

    starts = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), 256)] + 0.5 * sd * rng.standard_normal((256, d)),
        mu + 2.0 * sd * rng.standard_normal((256, d)),
    ]).astype(np.float32)
    ep_s = iterate(teacher, starts, ITER, lo, hi)
    vn = np.linalg.norm(np.array(field(teacher, mx.array(ep_s))), axis=1)
    attrs = find_attractors(ep_s, vn, vtol=5e-3 * scale, merge_r=0.1 * scale)
    conv = float((vn < 5e-3 * scale).mean())
    if len(attrs) > 1:
        dm = np.linalg.norm(attrs[:, None] - attrs[None, :], axis=-1)
        basin_eps = 0.25 * float(np.median(dm[dm > 0]))
    else:
        basin_eps = 0.25 * scale

    probes = np.concatenate([
        z_tr[rng.integers(0, len(z_tr), N_PROBES // 2)]
        + 0.5 * sd * rng.standard_normal((N_PROBES // 2, d)),
        mu + 2.0 * sd * rng.standard_normal((N_PROBES // 2, d)),
    ]).astype(np.float32)
    ep_ref = iterate(teacher, probes, ITER, lo, hi)
    ep_pert = iterate(teacher, probes + (0.05 * scale)
                      * rng.standard_normal(probes.shape).astype(np.float32), ITER, lo, hi)
    ceiling = float((np.linalg.norm(ep_pert - ep_ref, axis=1) < basin_eps).mean())

    rec = dict(d=d, noise=noise, n_attractors=int(len(attrs)), conv=conv,
               ceiling=ceiling, seconds=round(time.time() - t0))
    with OUT.open("a") as fh:
        fh.write(json.dumps(rec) + "\n")
    print(f"d={d} noise={noise}: attrs {len(attrs)} conv {conv:.2f} "
          f"ceil {ceiling:.2f} [{rec['seconds']}s]")


if __name__ == "__main__":
    for arg in sys.argv[1:]:
        for noise in NOISES:
            probe_teacher(int(arg), noise)

spike_frontier_results.jsonl

{"d": 8, "noise": 0.2, "n_attractors": 4, "conv": 1.0, "ceiling": 0.95703125, "seconds": 27}
{"d": 24, "noise": 0.2, "n_attractors": 5, "conv": 1.0, "ceiling": 0.93359375, "seconds": 27}
{"d": 8, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.97265625, "seconds": 14}
{"d": 24, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.91015625, "seconds": 13}
{"d": 8, "noise": 0.5, "n_attractors": 8, "conv": 1.0, "ceiling": 0.95703125, "seconds": 13}
{"d": 24, "noise": 0.5, "n_attractors": 8, "conv": 1.0, "ceiling": 0.95703125, "seconds": 13}
{"d": 16, "noise": 0.2, "n_attractors": 7, "conv": 1.0, "ceiling": 0.94140625, "seconds": 15}
{"d": 32, "noise": 0.2, "n_attractors": 7, "conv": 1.0, "ceiling": 0.90234375, "seconds": 17}
{"d": 16, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.95703125, "seconds": 15}
{"d": 32, "noise": 0.35, "n_attractors": 8, "conv": 1.0, "ceiling": 0.91796875, "seconds": 13}
{"d": 16, "noise": 0.5, "n_attractors": 8, "conv": 1.0, "ceiling": 0.93359375, "seconds": 14}
{"d": 32, "noise": 0.5, "n_attractors": 8, "conv": 1.0, "ceiling": 0.921875, "seconds": 15}

spike_mnist.py

"""Read-carrier at 16D: does the MNIST precision wall apply to a non-fit carrier?

Their result (knowledge-as-dynamics, runs/mnist_summary.json): at latent 16,
ALL methods fail at basin level — field arm 0.09, best (outdistill) 0.19,
teacher self-agreement ceiling 0.71. Diagnosis: SGD regression from sparse
probes plateaus at NMSE 0.049; residual compounds over ~800 iterations.

Carrier here: a table of (z, V_teacher(z)) pairs sampled from THEIR OWN
sampling distribution (teacher latent moments + deep teacher trajectories —
same calibration-level leak they document), read back by k-NN lookup.
1M samples vs their student's 3.07M field evaluations (12000 steps x 256).
No gradients. If this also fails, the wall is dimension-intrinsic (coverage);
if it holds, the wall is a property of the fit carrier.

Teacher/eval protocol adapted line-for-line from the reviewed source of
github.com/Veso-AI-Open-Source/knowledge-as-dynamics (MIT),
src/m5_fielddistill/mnist.py, so numbers are comparable.

Run: uv run --no-project --with mlx --with numpy python spike_mnist.py
"""

import gzip
import json
import time
import urllib.request
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np

HERE = Path(__file__).parent
DATA = HERE / "data"
MIRROR = "https://ossci-datasets.s3.amazonaws.com/mnist/"
FILES = {"train_x": "train-images-idx3-ubyte.gz", "train_y": "train-labels-idx1-ubyte.gz",
         "test_x": "t10k-images-idx3-ubyte.gz", "test_y": "t10k-labels-idx1-ubyte.gz"}

D_IN, D_LAT = 784, 16
ENC_DIMS = [784, 256, 128, 16]
DEC_DIMS = [16, 128, 256, 784]
N_TRAIN, N_TEST = 8000, 2000
STEPS, BS, LR = 3500, 256, 1e-3
TEACHER_NOISE = 0.3          # their best (0.2-0.5 sweep)
ITER_STEPS = 800
N_PROBES = 256
N_TABLE = 1_000_000

TEACHER_W = HERE / "mnist_teacher.safetensors"
META = HERE / "mnist_meta.npz"
OUT = HERE / "spike_mnist_results.json"


class MLP(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.layers = [nn.Linear(a, b) for a, b in zip(dims[:-1], dims[1:])]

    def __call__(self, x):
        n = len(self.layers)
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < n - 1:
                x = nn.gelu(x)
        return x


class AE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = MLP(ENC_DIMS)
        self.dec = MLP(DEC_DIMS)

    def __call__(self, x):
        return self.dec(self.enc(x))


def field(model, z):
    return model.enc(model.dec(z)) - z


def iterate(model, z, steps, clip_lo, clip_hi):
    lo, hi = mx.array(clip_lo), mx.array(clip_hi)
    for i in range(steps):
        z = model.enc(model.dec(z))
        z = mx.maximum(mx.minimum(z, hi), lo)
        if i % 20 == 0:
            mx.eval(z)
    mx.eval(z)
    return z


def find_attractors(endpoints, vnorms, vtol, merge_r):
    pts = endpoints[vnorms < vtol]
    centers = []
    for p in pts:
        if all(np.linalg.norm(p - c) > merge_r for c in centers):
            centers.append(p)
    return np.array(centers) if centers else np.zeros((0, endpoints.shape[1]))


def load_mnist():
    DATA.mkdir(parents=True, exist_ok=True)
    arrs = {}
    for key, fname in FILES.items():
        p = DATA / fname
        if not p.exists():
            print(f"downloading {fname} ...")
            urllib.request.urlretrieve(MIRROR + fname, p)
        raw = gzip.open(p, "rb").read()
        if "x" in key:
            arrs[key] = (np.frombuffer(raw, np.uint8, offset=16)
                         .reshape(-1, 784).astype(np.float32) / 255.0)
        else:
            arrs[key] = np.frombuffer(raw, np.uint8, offset=8).astype(np.int64)
    return (arrs["train_x"][:N_TRAIN], arrs["test_x"][:N_TEST])


def train_or_load_teacher():
    xtr, xte = load_mnist()
    mx.random.seed(0)
    model = AE()
    mx.eval(model.parameters())
    if TEACHER_W.exists() and META.exists():
        model.load_weights(str(TEACHER_W))
        mx.eval(model.parameters())
        return model, dict(np.load(META))

    t0 = time.time()
    opt = optim.Adam(learning_rate=LR)
    X = mx.array(xtr)
    rng = np.random.default_rng(0)

    def loss_fn(m, xb, xn):
        return mx.mean((m(xn) - xb) ** 2)

    lg = nn.value_and_grad(model, loss_fn)
    for i in range(STEPS):
        idx = mx.array(rng.integers(0, N_TRAIN, BS))
        xb = X[idx]
        xn = xb + TEACHER_NOISE * mx.array(rng.standard_normal((BS, D_IN)).astype(np.float32))
        _, g = lg(model, xb, xn)
        opt.update(model, g)
        mx.eval(model.parameters(), opt.state)
        if (i + 1) % 1000 == 0:
            print(f"  teacher step {i+1}/{STEPS} [{time.time()-t0:.0f}s]")

    z_tr = np.array(model.enc(X))
    z_te = np.array(model.enc(mx.array(xte)))
    mu, sd = z_tr.mean(0), z_tr.std(0)
    scale = float(np.linalg.norm(sd))
    clip_lo = (mu - 8 * sd).astype(np.float32)
    clip_hi = (mu + 8 * sd).astype(np.float32)

    rng = np.random.default_rng(7)
    starts = np.concatenate([
        z_te[rng.integers(0, N_TEST, 256)] + 0.5 * sd * rng.standard_normal((256, D_LAT)),
        mu + 2.0 * sd * rng.standard_normal((256, D_LAT)),
    ]).astype(np.float32)
    ep = np.array(iterate(model, mx.array(starts), ITER_STEPS, clip_lo, clip_hi))
    vn = np.linalg.norm(np.array(field(model, mx.array(ep))), axis=1)
    attractors = find_attractors(ep, vn, vtol=1e-2, merge_r=0.1 * scale)
    if len(attractors) > 1:
        dm = np.linalg.norm(attractors[:, None] - attractors[None, :], axis=-1)
        basin_eps = 0.25 * float(np.median(dm[dm > 0]))
    else:
        basin_eps = 0.25 * scale

    probes = np.concatenate([
        z_te[rng.integers(0, N_TEST, N_PROBES // 2)]
        + 0.5 * sd * rng.standard_normal((N_PROBES // 2, D_LAT)),
        mu + 2.0 * sd * rng.standard_normal((N_PROBES // 2, D_LAT)),
    ]).astype(np.float32)
    ep_probes = np.array(iterate(model, mx.array(probes), ITER_STEPS, clip_lo, clip_hi))
    ep_pert = np.array(iterate(model, mx.array(
        probes + (0.05 * scale) * rng.standard_normal(probes.shape).astype(np.float32)),
        ITER_STEPS, clip_lo, clip_hi))
    self_agree = float((np.linalg.norm(ep_pert - ep_probes, axis=1) < basin_eps).mean())
    conv_frac = float((vn < 1e-2).mean())

    zeval = np.concatenate([
        z_te[rng.integers(0, N_TEST, 512)] + 0.2 * sd * rng.standard_normal((512, D_LAT)),
        z_te[rng.integers(0, N_TEST, 512)] + 1.0 * sd * rng.standard_normal((512, D_LAT)),
        mu + 2.0 * sd * rng.standard_normal((512, D_LAT)),
    ]).astype(np.float32)
    Vt_eval = np.array(field(model, mx.array(zeval)))

    model.save_weights(str(TEACHER_W))
    np.savez(META, mu=mu, sd=sd, scale=scale, clip_lo=clip_lo, clip_hi=clip_hi,
             attractors=attractors, basin_eps=basin_eps, probes=probes,
             ep_probes=ep_probes, self_agree=self_agree, conv_frac=conv_frac,
             zeval=zeval, Vt_eval=Vt_eval)
    print(f"teacher: {len(attractors)} attractors, conv {conv_frac:.2f}, "
          f"ceiling {self_agree:.2f}, basin_eps {basin_eps:.3f}, scale {scale:.2f}")
    return model, dict(np.load(META))


# ---------- the read carrier: k-NN table over sampled field ----------

def build_table(teacher, meta, n=N_TABLE, seed=1):
    """Same sampling distribution as their field arm's sample_z: half broad
    prior from teacher latent moments, half deep teacher-trajectory points."""
    mu, sd = meta["mu"], meta["sd"]
    rng = np.random.default_rng(1000 + seed)
    zs = []
    bs = 4096
    t0 = time.time()
    while sum(len(z) for z in zs) < n:
        half = bs // 2
        broad = (mu + 2.0 * sd * rng.standard_normal((half, D_LAT))).astype(np.float32)
        core = mx.array((mu + 1.5 * sd * rng.standard_normal((bs - half, D_LAT)))
                        .astype(np.float32))
        depth = int(rng.integers(2, 25))
        for _ in range(depth):
            core = teacher.enc(teacher.dec(core))
        mx.eval(core)
        zs.append(np.concatenate([broad, np.array(core)]))
    Z = np.concatenate(zs)[:n]
    V = np.empty_like(Z)
    for i in range(0, n, 8192):
        V[i:i + 8192] = np.array(field(teacher, mx.array(Z[i:i + 8192])))
    print(f"table: {n} (z, V) pairs [{time.time()-t0:.0f}s]")
    return mx.array(Z), mx.array(V)


def knn_field(Z, V, Zsq, zq, k):
    """k-NN readout of the table at query points zq (mlx, matmul distances)."""
    d2 = Zsq[None, :] - 2 * (zq @ Z.T)          # + |zq|^2, constant per row
    if k == 1:
        idx = mx.argmin(d2, axis=1)
        return V[idx]
    idx = mx.argpartition(d2, kth=k - 1, axis=1)[:, :k]
    dk = mx.take_along_axis(d2, idx, axis=1)
    dk = dk - mx.min(dk, axis=1, keepdims=True)
    w = mx.exp(-dk / (mx.mean(dk, axis=1, keepdims=True) + 1e-9))
    w = w / mx.sum(w, axis=1, keepdims=True)
    return mx.sum(V[idx] * w[:, :, None], axis=1)


def carrier_iterate(Z, V, Zsq, z0, clip_lo, clip_hi, k, steps=ITER_STEPS):
    z = mx.array(z0)
    lo, hi = mx.array(clip_lo), mx.array(clip_hi)
    for i in range(steps):
        z = z + knn_field(Z, V, Zsq, z, k)
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    return np.array(z)


def eval_carrier(teacher, meta, Z, V, k):
    Zsq = mx.sum(Z * Z, axis=1)
    mx.eval(Zsq)
    t0 = time.time()
    ep = carrier_iterate(Z, V, Zsq, meta["probes"], meta["clip_lo"], meta["clip_hi"], k)
    dists = np.linalg.norm(ep - meta["ep_probes"], axis=1)
    basin = float((dists < float(meta["basin_eps"])).mean())
    img_s = np.array(teacher.dec(mx.array(ep.astype(np.float32))))
    img_t = np.array(teacher.dec(mx.array(meta["ep_probes"])))
    dec_agree = float((((img_s - img_t) ** 2).mean(axis=1) < 0.01).mean())

    # field fidelity off the table's own points (their zeval mix)
    Vh = np.array(knn_field(Z, V, Zsq, mx.array(meta["zeval"]), k))
    Vt = meta["Vt_eval"]
    nmse = float(((Vh - Vt) ** 2).mean() / (Vt ** 2).mean())
    mag = np.linalg.norm(Vt, axis=1)
    mask = mag > np.percentile(mag, 10)
    cos = float((np.sum(Vh[mask] * Vt[mask], axis=1)
                 / (np.linalg.norm(Vh[mask], axis=1) * mag[mask] + 1e-12)).mean())
    return dict(k=k, basin_agreement=basin, dec_agree=dec_agree,
                endpoint_median=float(np.median(dists)),
                field_nmse=nmse, field_cos=cos,
                eval_seconds=round(time.time() - t0, 1))


def main():
    teacher, meta = train_or_load_teacher()
    print(f"teacher: {len(meta['attractors'])} attractors, conv {float(meta['conv_frac']):.2f}, "
          f"ceiling {float(meta['self_agree']):.2f}, basin_eps {float(meta['basin_eps']):.3f}")
    Z, V = build_table(teacher, meta)
    results = []
    for k in (1, 8):
        r = eval_carrier(teacher, meta, Z, V, k)
        results.append(r)
        print(f"read carrier k={k}: basin {r['basin_agreement']:.2f}  "
              f"dec {r['dec_agree']:.2f}  median-endpoint-err {r['endpoint_median']:.2f}  "
              f"nmse {r['field_nmse']:.3f}  cos {r['field_cos']:.2f}  "
              f"[{r['eval_seconds']}s]")
    OUT.write_text(json.dumps(dict(
        ceiling=float(meta["self_agree"]), basin_eps=float(meta["basin_eps"]),
        n_table=N_TABLE, results=results), indent=2))
    print(f"\ntheir mnist_summary.json for comparison: field 0.09, fieldplus 0.07, "
          f"outdistill 0.19, ceiling 0.71 (nmse: field 0.049)")


if __name__ == "__main__":
    main()

spike_mnist_results.json

{
  "ceiling": 0.703125,
  "basin_eps": 2.450249195098877,
  "n_table": 1000000,
  "results": [
    {
      "k": 1,
      "basin_agreement": 0.31640625,
      "dec_agree": 0.2421875,
      "endpoint_median": 3.8904805183410645,
      "field_nmse": 0.37175649404525757,
      "field_cos": 0.5767523050308228,
      "eval_seconds": 42.9
    },
    {
      "k": 8,
      "basin_agreement": 0.3125,
      "dec_agree": 0.2421875,
      "endpoint_median": 3.4760305881500244,
      "field_nmse": 0.22140631079673767,
      "field_cos": 0.6998631358146667,
      "eval_seconds": 426.6
    }
  ]
}

spike_wiener.py

"""Wiener/GP carrier at 16D: cosmology's answer to the sparse-probe field problem.

Third carrier type on the MNIST tier. Same teacher, same 1M-sample table as
spike_mnist.py, but the field is reconstructed by kernel ridge regression
(RBF kernel = GP posterior mean = Wiener filter for a stationary field):
interpolation that respects the field's correlation structure, the method
cosmologists use to reconstruct velocity fields from sparse galaxy probes
(POTENT / Hoffman-Ribak constrained realizations).

Controls:
  - NN carrier on a table subsampled to the SAME M points as the GP's
    inducing set (isolates correlation-structure gain from sample-count gain).
  - Length scale selected on held-out table points only (no extra teacher access).

Reference numbers: NN-1M 0.32, their field arm 0.09, outdistill 0.19, ceiling 0.70.

Run: uv run --no-project --with mlx --with numpy python spike_wiener.py
"""

import json
import time

import mlx.core as mx
import numpy as np

from spike_mnist import (HERE, ITER_STEPS, build_table, knn_field,
                         train_or_load_teacher)

OUT = HERE / "spike_wiener_results.json"
M_SELECT = 4096       # inducing size for length-scale selection (cheap solves)
M_FINAL = 16384       # inducing size for the real run
N_VAL = 4096          # held-out table points for selection
JITTER = 1e-8
SIGMA_MULTS = (0.25, 0.5, 1.0, 2.0)


def rbf_solve(ZM, VM, sigma):
    """alpha = (K + jitter*tr/M * I)^-1 VM, K_ij = exp(-|zi-zj|^2 / 2sigma^2). f64."""
    Z = ZM.astype(np.float64)
    sq = (Z ** 2).sum(1)
    d2 = sq[:, None] - 2 * Z @ Z.T + sq[None, :]
    K = np.exp(-np.maximum(d2, 0) / (2 * sigma ** 2))
    K[np.diag_indices_from(K)] += JITTER * K.shape[0]
    return np.linalg.solve(K, VM.astype(np.float64))


def gp_predict(ZM_m, Zsq_m, alpha_m, sigma, zq):
    d2 = Zsq_m[None, :] - 2 * (zq @ ZM_m.T) + mx.sum(zq * zq, axis=1)[:, None]
    Kq = mx.exp(-mx.maximum(d2, 0.0) / (2 * sigma ** 2))
    return Kq @ alpha_m


def gp_iterate(ZM_m, Zsq_m, alpha_m, sigma, z0, clip_lo, clip_hi, steps=ITER_STEPS):
    z = mx.array(z0)
    lo, hi = mx.array(clip_lo), mx.array(clip_hi)
    for _ in range(steps):
        z = z + gp_predict(ZM_m, Zsq_m, alpha_m, sigma, z)
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    return np.array(z)


def field_metrics(Vh, Vt):
    nmse = float(((Vh - Vt) ** 2).mean() / (Vt ** 2).mean())
    mag = np.linalg.norm(Vt, axis=1)
    mask = mag > np.percentile(mag, 10)
    cos = float((np.sum(Vh[mask] * Vt[mask], axis=1)
                 / (np.linalg.norm(Vh[mask], axis=1) * mag[mask] + 1e-12)).mean())
    return nmse, cos


def basin_metrics(teacher, meta, ep):
    dists = np.linalg.norm(ep - meta["ep_probes"], axis=1)
    basin = float((dists < float(meta["basin_eps"])).mean())
    img_s = np.array(teacher.dec(mx.array(ep.astype(np.float32))))
    img_t = np.array(teacher.dec(mx.array(meta["ep_probes"])))
    dec_agree = float((((img_s - img_t) ** 2).mean(axis=1) < 0.01).mean())
    return basin, dec_agree, float(np.median(dists))


def main():
    t0 = time.time()
    teacher, meta = train_or_load_teacher()
    print(f"teacher: ceiling {float(meta['self_agree']):.2f}, "
          f"basin_eps {float(meta['basin_eps']):.3f}")
    Z, V = build_table(teacher, meta)          # same seed -> same 1M table
    Zn, Vn = np.array(Z), np.array(V)
    rng = np.random.default_rng(42)
    perm = rng.permutation(len(Zn))
    val_idx = perm[:N_VAL]
    Zval, Vval = Zn[val_idx], Vn[val_idx]
    pool = perm[N_VAL:]

    # --- length-scale selection at M_SELECT on held-out table points ---
    sel_idx = pool[:M_SELECT]
    ZM, VM = Zn[sel_idx], Vn[sel_idx]
    d_med = np.median(np.linalg.norm(
        ZM[rng.integers(0, M_SELECT, 2000)] - ZM[rng.integers(0, M_SELECT, 2000)], axis=1))
    print(f"median pairwise distance (table subsample): {d_med:.2f}")
    best = None
    for mult in SIGMA_MULTS:
        sigma = mult * d_med
        alpha = rbf_solve(ZM, VM, sigma)
        ZM_m = mx.array(ZM); alpha_m = mx.array(alpha.astype(np.float32))
        Zsq_m = mx.sum(ZM_m * ZM_m, axis=1)
        Vh = np.array(gp_predict(ZM_m, Zsq_m, alpha_m, sigma, mx.array(Zval)))
        nmse, cos = field_metrics(Vh, Vval)
        print(f"  sigma = {mult:.2f} x median: held-out nmse {nmse:.4f} cos {cos:.2f}")
        if best is None or nmse < best[1]:
            best = (mult, nmse)
    mult = best[0]
    print(f"selected sigma = {mult} x median")

    # --- final GP carrier at M_FINAL ---
    fin_idx = pool[:M_FINAL]
    ZM, VM = Zn[fin_idx], Vn[fin_idx]
    sigma = mult * d_med
    t1 = time.time()
    alpha = rbf_solve(ZM, VM, sigma)
    print(f"solved {M_FINAL} x {M_FINAL} system [{time.time()-t1:.0f}s]")
    ZM_m = mx.array(ZM); alpha_m = mx.array(alpha.astype(np.float32))
    Zsq_m = mx.sum(ZM_m * ZM_m, axis=1)
    mx.eval(ZM_m, alpha_m, Zsq_m)

    Vh = np.array(gp_predict(ZM_m, Zsq_m, alpha_m, sigma, mx.array(meta["zeval"])))
    nmse, cos = field_metrics(Vh, meta["Vt_eval"])
    ep = gp_iterate(ZM_m, Zsq_m, alpha_m, sigma, meta["probes"],
                    meta["clip_lo"], meta["clip_hi"])
    basin, dec, med = basin_metrics(teacher, meta, ep)
    gp_res = dict(carrier=f"gp_M{M_FINAL}", sigma_mult=mult, basin_agreement=basin,
                  dec_agree=dec, endpoint_median=med, field_nmse=nmse, field_cos=cos)
    print(f"GP carrier (M={M_FINAL}): basin {basin:.2f}  dec {dec:.2f}  "
          f"median-err {med:.2f}  nmse {nmse:.3f}  cos {cos:.2f}")

    # --- control: NN carrier on the SAME M points (budget-matched) ---
    Zs, Vs = mx.array(ZM), mx.array(VM)
    Zsq_s = mx.sum(Zs * Zs, axis=1)
    mx.eval(Zs, Vs, Zsq_s)
    z = mx.array(meta["probes"])
    lo, hi = mx.array(meta["clip_lo"]), mx.array(meta["clip_hi"])
    for _ in range(ITER_STEPS):
        z = z + knn_field(Zs, Vs, Zsq_s, z, 1)
        z = mx.maximum(mx.minimum(z, hi), lo)
        mx.eval(z)
    basin_n, dec_n, med_n = basin_metrics(teacher, meta, np.array(z))
    Vh_n = np.array(knn_field(Zs, Vs, Zsq_s, mx.array(meta["zeval"]), 1))
    nmse_n, cos_n = field_metrics(Vh_n, meta["Vt_eval"])
    nn_res = dict(carrier=f"nn1_M{M_FINAL}", basin_agreement=basin_n, dec_agree=dec_n,
                  endpoint_median=med_n, field_nmse=nmse_n, field_cos=cos_n)
    print(f"NN carrier (same M={M_FINAL}): basin {basin_n:.2f}  dec {dec_n:.2f}  "
          f"median-err {med_n:.2f}  nmse {nmse_n:.3f}  cos {cos_n:.2f}")

    OUT.write_text(json.dumps(dict(
        ceiling=float(meta["self_agree"]), results=[gp_res, nn_res]), indent=2))
    print(f"\nreference: NN-1M 0.32 | their field 0.09 | outdistill 0.19 | "
          f"ceiling {float(meta['self_agree']):.2f}  [{time.time()-t0:.0f}s total]")


if __name__ == "__main__":
    main()

spike_wiener_results.json

{
  "ceiling": 0.703125,
  "results": [
    {
      "carrier": "gp_M16384",
      "sigma_mult": 2.0,
      "basin_agreement": 0.0,
      "dec_agree": 0.0,
      "endpoint_median": 7.072707653045654,
      "field_nmse": 0.10303101688623428,
      "field_cos": 0.809187114238739
    },
    {
      "carrier": "nn1_M16384",
      "basin_agreement": 0.0,
      "dec_agree": 0.00390625,
      "endpoint_median": 11.727006912231445,
      "field_nmse": 0.5944164395332336,
      "field_cos": 0.3631992042064667
    }
  ]
}