1import sys, json, math, random
  2import numpy as np
  3sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  4import torch
  5from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  6
  7SEEDS = tuple(range(8))
  8SWEEP_SEEDS = tuple(range(4))
  9LR_GRID = [0.003, 0.006, 0.009]
 10EPOCHS = 18
 11BATCH = 128
 12CLIP = 1.0
 13
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 18
 19
 20def run(seed, lr, scheduled):
 21    seed_all(seed)
 22    requested = 'cuda' if torch.cuda.is_available() else 'cpu'
 23    try:
 24        return _run_device(seed, lr, scheduled, requested)
 25    except Exception:
 26        if requested == 'cuda':
 27            try: torch.cuda.empty_cache()
 28            except Exception: pass
 29            return _run_device(seed, lr, scheduled, 'cpu')
 30        raise
 31
 32
 33def _run_device(seed, lr, scheduled, device_name):
 34    device = torch.device(device_name)
 35    d = get_dataset('tabular', seed, n_train=400, n_test=200)
 36    xtr, ytr = d['xtr'].to(device), d['ytr'].to(device)
 37    xte, yte = d['xte'].to(device), d['yte'].to(device)
 38    model = make_model('mlp_tiny', d['input_shape'], d['out_dim']).to(device)
 39    opt = torch.optim.SGD(model.parameters(), lr=lr)
 40    fhat = None; qhist = []; gaphist = []; train_hist = []
 41    local_streak = 0; switched = False; switch_step = None; clip_steps = 0
 42    step = 0
 43    n = len(xtr)
 44    for epoch in range(EPOCHS):
 45        # deterministic per-seed permutation, equivalent data budget in both arms
 46        gen = torch.Generator(device='cpu').manual_seed(seed * 1000 + epoch)
 47        perm = torch.randperm(n, generator=gen, device='cpu')
 48        for start in range(0, n, BATCH):
 49            idx = perm[start:start+BATCH].to(device)
 50            xb, yb = xtr[idx], ytr[idx]
 51            opt.zero_grad(set_to_none=True)
 52            pred = model(xb)
 53            loss = torch.nn.functional.mse_loss(pred, yb)
 54            loss.backward()
 55            g2 = 0.0
 56            for p in model.parameters():
 57                if p.grad is not None: g2 += float((p.grad.detach() ** 2).sum().cpu())
 58            g = math.sqrt(max(g2, 0.0)); lv = float(loss.detach().cpu())
 59            if fhat is None: fhat = lv
 60            else: fhat = min(fhat, 0.98 * fhat + 0.02 * lv)
 61            gap = max(lv - fhat, 1e-8)
 62            q = g / math.sqrt(gap)
 63            qhist.append(q); gaphist.append(gap); train_hist.append(lv)
 64            if scheduled and len(qhist) >= 20:
 65                qwin = np.asarray(qhist[-100:]); gapwin = np.asarray(gaphist[-100:])
 66                floor = max(float(np.percentile(qwin, 10)) / 2.0, 1e-5)
 67                candidate = gap <= float(np.percentile(gapwin, 35)) and q >= floor
 68                local_streak = local_streak + 1 if candidate else 0
 69                if local_streak >= 5 and not switched:
 70                    for group in opt.param_groups: group['lr'] *= 1.5
 71                    switched = True; switch_step = step
 72            pre = torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP)
 73            clip_steps += int(float(pre) > CLIP)
 74            opt.step(); step += 1
 75    with torch.no_grad():
 76        test_loss = float(torch.nn.functional.mse_loss(model(xte), yte).cpu())
 77    # Re-test the claimed local prediction on the trained model's observed trajectory.
 78    qarr, garr = np.asarray(qhist), np.asarray(gaphist)
 79    tail = min(40, len(qarr))
 80    observed_slope = float(np.polyfit(np.arange(tail), np.log(np.maximum(garr[-tail:], 1e-12)), 1)[0]) if tail >= 3 else float('nan')
 81    predicted_slope = float(-np.mean(qarr[-tail:] ** 2)) if tail else float('nan')
 82    return {'metric': test_loss, 'switched': switched, 'switch_step': switch_step,
 83            'lr_final': float(opt.param_groups[0]['lr']), 'clip_steps': clip_steps,
 84            'observed_log_gap_slope': observed_slope, 'predicted_minus_q2': predicted_slope,
 85            'q_tail_mean': float(np.mean(qarr[-tail:])) if tail else float('nan'),
 86            'device': device_name}
 87
 88
 89def main():
 90    base = sweep_baseline(lambda cfg: (lambda seed: run(seed, cfg['lr'], False)['metric']),
 91                          [{'lr': x} for x in LR_GRID], seeds=SWEEP_SEEDS)
 92    best_lr = float(base['best_cfg']['lr'])
 93    idea_cfgs = [{'lr': x} for x in LR_GRID]
 94    idea_sweep = []
 95    for cfg in idea_cfgs:
 96        r = evaluate(lambda seed, x=cfg['lr']: run(seed, x, True)['metric'], seeds=SWEEP_SEEDS)
 97        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
 98    best_idea_cfg = min(idea_sweep, key=lambda z: z['mean'])['cfg']
 99    # Full 8-seed results at the best idea configuration, with full per-run diagnostics.
100    idea_runs = [run(s, best_idea_cfg['lr'], True) for s in SEEDS]
101    idea_res = {'mean': float(np.mean([r['metric'] for r in idea_runs])),
102                'std': float(np.std([r['metric'] for r in idea_runs])),
103                'per_seed': [r['metric'] for r in idea_runs], 'n': 8,
104                'cfg': best_idea_cfg, 'runs': idea_runs, 'sweep': idea_sweep}
105    rep = make_report('tabular', 'mlp_tiny', base, idea_res)
106    # Signature uses only trained-model observed numbers, not the toy identity.
107    obs = [r['observed_log_gap_slope'] for r in idea_runs]
108    pred = [r['predicted_minus_q2'] for r in idea_runs]
109    rel = [abs(a-b)/max(abs(b),1e-12) for a,b in zip(obs,pred) if np.isfinite(a) and np.isfinite(b)]
110    rep['mechanism_signature'] = {
111        'quantity': 'tail log(training loss gap) slope versus -mean(q^2), measured during trained NN runs',
112        'observed_mean': float(np.mean(obs)), 'predicted_mean': float(np.mean(pred)),
113        'relative_error_mean': float(np.mean(rel)) if rel else None,
114        'tolerance': 0.20, 'confirmed': bool(rel and np.mean(rel) <= 0.20),
115        'switch_fraction': float(np.mean([r['switched'] for r in idea_runs]))}
116    rep['audit'] = {'baseline_lr_grid': LR_GRID, 'idea_lr_grid': LR_GRID,
117                    'baseline_selected_lr': best_lr, 'idea_selected_lr': best_idea_cfg['lr'],
118                    'epochs': EPOCHS, 'batch': BATCH, 'clip_norm': CLIP,
119                    'domain_rationale': 'tabular is the prescribed structural track for optimizer and learning-rate schedule ideas'}
120    print(json.dumps(rep, indent=2))
121
122if __name__ == '__main__': main()