1import json
  2import sys
  3from pathlib import Path
  4import numpy as np
  5import torch
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
  9from bench.protocol import DEFAULT_SEEDS
 10
 11TRACK = 'tabular'
 12MODEL = 'mlp_tiny'
 13EPOCHS = 20
 14BATCH = 128
 15# Shared union: every lr and scale tested by the idea is also tested by baseline.
 16GRID = [
 17    {'lr': 1e-3, 'scale': 0.7},
 18    {'lr': 3e-3, 'scale': 1.0},
 19    {'lr': 1e-2, 'scale': 1.3},
 20]
 21
 22
 23def coulomb_bank(n, beta=2.0, eta=0.02, steps=180, eps=1e-5, seed=0):
 24    rng = np.random.default_rng(seed)
 25    a = np.linspace(0, 2*np.pi, n, endpoint=False) + rng.normal(0, .03, n)
 26    z = np.c_[np.cos(a), np.sin(a)] * np.sqrt(n / 2.0)
 27    for _ in range(steps):
 28        d = z[:, None, :] - z[None, :, :]
 29        r2 = np.sum(d*d, axis=-1) + eps
 30        np.fill_diagonal(r2, np.inf)
 31        rep = np.sum(d / r2[:, :, None], axis=1) / n
 32        z += eta * (-z + rep) + np.sqrt(2 * eta / beta) * rng.normal(size=z.shape)
 33    # Normalize to unit per-coordinate RMS, matching a conventional small init.
 34    z /= max(np.sqrt(np.mean(z*z)), 1e-8)
 35    return z
 36
 37
 38def set_first_layer(net, mode, scale, seed):
 39    # A 2-D Coulomb bank is lifted to the 10-D input rows by a fixed random
 40    # orthonormal projection. The only system difference is this initialization.
 41    layer = net[0]
 42    with torch.no_grad():
 43        if mode == 'coulomb':
 44            z = coulomb_bank(layer.out_features, seed=seed)
 45            rng = np.random.default_rng(seed + 991)
 46            q, _ = np.linalg.qr(rng.normal(size=(layer.in_features, 2)))
 47            w = z @ q.T
 48            layer.weight.copy_(torch.tensor(w * (scale * 0.12), dtype=layer.weight.dtype))
 49            layer.bias.zero_()
 50        else:
 51            torch.manual_seed(seed + 991)
 52            torch.nn.init.normal_(layer.weight, mean=0.0, std=scale * 0.12)
 53            torch.nn.init.zeros_(layer.bias)
 54
 55
 56def train_one(seed, cfg, mode, return_net=False):
 57    np.random.seed(seed)
 58    torch.manual_seed(seed)
 59    ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 60    net = make_model(MODEL, ds['input_shape'], ds['out_dim'])
 61    set_first_layer(net, mode, cfg['scale'], seed)
 62    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
 63                                    log=lambda *_: None)
 64    if return_net:
 65        return metric, net
 66    return metric
 67
 68
 69def baseline_factory(cfg):
 70    return lambda seed: train_one(seed, cfg, 'gaussian')
 71
 72
 73def stats_from_net(net):
 74    w = net[0].weight.detach().cpu().numpy()
 75    d = np.sqrt(((w[:, None] - w[None, :])**2).sum(-1))
 76    np.fill_diagonal(d, np.inf)
 77    return float(np.median(d.min(axis=1))), float(np.mean(d.min(axis=1))), float(np.std(w, axis=0).mean())
 78
 79
 80def main():
 81    # Baseline is tuned on the protocol's four-seed sweep, then re-evaluated on 8.
 82    base = sweep_baseline(baseline_factory, GRID)
 83    best = base['best_cfg']
 84    idea_grid = [best, {'lr': 1e-3, 'scale': 0.7}, {'lr': 1e-2, 'scale': 1.3}]
 85    # Deduplicate while preserving the required three-config idea budget.
 86    uniq = []
 87    for c in idea_grid:
 88        if c not in uniq: uniq.append(c)
 89    idea_cfg = min(uniq, key=lambda c: np.mean([train_one(s, c, 'coulomb') for s in (0,1,2,3)]))
 90    idea_vals = [train_one(s, idea_cfg, 'coulomb') for s in DEFAULT_SEEDS]
 91    # Paired trained-model signature: compare final first-layer nearest-neighbor spacing.
 92    bnn, inn = [], []
 93    for s in DEFAULT_SEEDS:
 94        _, bn = train_one(s, best, 'gaussian', True)
 95        _, cn = train_one(s, idea_cfg, 'coulomb', True)
 96        bnn.append(stats_from_net(bn))
 97        inn.append(stats_from_net(cn))
 98    bmean = [float(np.mean(x)) for x in zip(*bnn)]
 99    imean = [float(np.mean(x)) for x in zip(*inn)]
100    # The stage-1 prediction is higher non-collapse spacing after training;
101    # confirmation requires a positive, practically nontrivial observed ratio.
102    ratio = imean[0] / max(bmean[0], 1e-12)
103    sig = {'quantity': 'trained first-layer median nearest-neighbor distance',
104           'prediction': 'Coulomb repulsion preserves larger inter-row spacing as N=64 bank grows',
105           'baseline_observed': bmean[0], 'idea_observed': imean[0],
106           'observed_ratio_idea_over_baseline': ratio,
107           'confirmed': bool(ratio > 1.05),
108           'n_models': 8}
109    report = make_report(TRACK, MODEL, base, {'cfg': idea_cfg, 'mean': float(np.mean(idea_vals)),
110                                             'std': float(np.std(idea_vals)), 'per_seed': [float(x) for x in idea_vals],
111                                             'n': len(idea_vals)},
112                         {'mechanism_signature': sig,
113                          'track_choice': 'tabular: initialization is explicitly covered by the benchmark mapping; same mlp_tiny and training loop used.',
114                          'idea_sweep': [{'cfg': c} for c in uniq],
115                          'trained_geometry': {'baseline_median_nn': bmean[0], 'idea_median_nn': imean[0],
116                                              'baseline_mean_nn': bmean[1], 'idea_mean_nn': imean[1]}})
117    Path('bench_report.json').write_text(json.dumps(report, indent=2))
118    print(json.dumps(report, indent=2))
119
120if __name__ == '__main__':
121    main()