1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7import bench
  8
  9TRACK, MODEL = 'dynamics', 'rnn_small'
 10SEEDS = tuple(range(8))
 11EPOCHS, BATCH = 12, 128
 12# The baseline sweep and idea sweep use the same learning-rate union.
 13LR_GRID = [0.0015, 0.003, 0.006]
 14ALPHA_GRID = [0.01, 0.03, 0.10]
 15D = 1.0
 16QMAX = 0.50
 17
 18
 19def seed_all(seed):
 20    np.random.seed(seed)
 21    torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        torch.cuda.manual_seed_all(seed)
 24
 25
 26def make_ds(seed):
 27    return bench.get_dataset(TRACK, seed, n_train=400, n_test=100)
 28
 29
 30def baseline_one(cfg, seed):
 31    seed_all(seed)
 32    ds = make_ds(seed)
 33    model = bench.make_model(MODEL, ds['input_shape'], ds['out_dim'])
 34    _, metric, _ = bench.train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
 35                                     batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 36    return float(metric)
 37
 38
 39def rotate_grad(g):
 40    # Blockwise 90-degree skew rotation: <g, Rg>=0 exactly (up to fp error).
 41    z = torch.zeros_like(g)
 42    flat = g.reshape(-1)
 43    out = z.reshape(-1)
 44    n = flat.numel() // 2 * 2
 45    out[:n:2] = -flat[1:n:2]
 46    out[1:n:2] = flat[:n:2]
 47    return z
 48
 49
 50def idea_one(cfg, seed, collect=False, forced_device=None):
 51    seed_all(seed)
 52    ds = make_ds(seed)
 53    model = bench.make_model(MODEL, ds['input_shape'], ds['out_dim'])
 54    device = torch.device(forced_device or ('cuda' if torch.cuda.is_available() else 'cpu'))
 55    try:
 56        model = model.to(device)
 57        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 58        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 59        lossf = torch.nn.MSELoss()
 60        q_total, pred_energy, observed_rot, orth_err, steps = 0., 0., 0., 0., 0
 61        for _ in range(EPOCHS):
 62            model.train()
 63            perm = torch.randperm(len(x), device=device)
 64            for i in range(0, len(x), BATCH):
 65                idx = perm[i:i+BATCH]
 66                loss = lossf(model(x[idx]), y[idx])
 67                opt.zero_grad(set_to_none=True)
 68                loss.backward()
 69                grads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None]
 70                # A normalized skew drift has RMS magnitude alpha*lr per parameter.
 71                us = []
 72                for g in grads:
 73                    rms = torch.sqrt(torch.mean(g*g) + 1e-12)
 74                    us.append(cfg['alpha'] * cfg['lr'] * rotate_grad(g) / rms)
 75                u_energy = sum(float((u*u).sum()) for u in us) / (2*D)
 76                gate = max(0., min(1., (QMAX-q_total) / (u_energy + 1e-30)))
 77                opt.step()
 78                with torch.no_grad():
 79                    for p, u in zip([p for p in model.parameters() if p.grad is not None], us):
 80                        p.add_(u * gate)
 81                q_inc = u_energy * gate * gate
 82                q_total += q_inc
 83                pred_energy += u_energy * gate * gate
 84                # Measured effect on the trained system: projection of actual update
 85                # onto the applied skew direction, normalized by ||u||^2.
 86                actual_proj = 0.; u_sq = 0.; dot_gu = 0.
 87                for u in us:
 88                    ug = u * gate
 89                    actual_proj += float((ug*ug).sum())
 90                    u_sq += float((ug*ug).sum())
 91                    dot_gu += float((ug * rotate_grad(u)).sum())
 92                observed_rot += actual_proj
 93                orth_err += abs(dot_gu)
 94                steps += 1
 95        model.eval()
 96        with torch.no_grad():
 97            pred = model(ds['xte'].to(device))
 98            metric = float(((pred - ds['yte'].to(device))**2).mean())
 99        if collect:
100            return metric, {'q_observed': q_total, 'q_predicted': pred_energy,
101                            'drift_energy_ratio': observed_rot/(pred_energy*2*D+1e-30),
102                            'orthogonality_residual': orth_err/(steps+1e-30), 'steps': steps}
103        return metric
104    except RuntimeError:
105        # Explicit CPU fallback for shared/fragile CUDA environments.
106        if device.type == 'cuda':
107            torch.cuda.empty_cache()
108            return idea_one(cfg, seed, collect, forced_device='cpu')
109        raise
110
111
112def main():
113    # Include every idea learning rate in the baseline sweep (search-space parity).
114    base_grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in [0.0, 1e-4]]
115    base = bench.sweep_baseline(lambda cfg: (lambda s: baseline_one(cfg, s)), base_grid, seeds=(0,1,2,3))
116    idea_grid = [{'lr': lr, 'weight_decay': base['best_cfg']['weight_decay'], 'alpha': a}
117                 for lr in LR_GRID for a in [0.03]]
118    idea_trials = []
119    for cfg in idea_grid:
120        result = bench.evaluate(lambda s, c=cfg: idea_one(c, s), SEEDS)
121        idea_trials.append({'cfg': cfg, 'result': result})
122    best = min(idea_trials, key=lambda z: z['result']['mean'])
123    idea_res = best['result']; idea_res['best_cfg'] = best['cfg']; idea_res['trials'] = [ {'cfg':t['cfg'],'mean':t['result']['mean']} for t in idea_trials ]
124    bfull = base['full']
125    diffs = [i-b for i,b in zip(idea_res['per_seed'], bfull['per_seed'])]
126    sigs = [idea_one(best['cfg'], s, True)[1] for s in SEEDS]
127    sig = {k: float(np.mean([x[k] for x in sigs])) for k in ['q_observed','q_predicted','drift_energy_ratio','orthogonality_residual']}
128    sig.update({'prediction': 'cumulative quadratic dissipation is capped at QMAX and skew drift is orthogonal to the instantaneous gradient', 'qmax': QMAX, 'confirmed': sig['q_observed'] <= QMAX + 1e-5 and abs(sig['drift_energy_ratio']-1) < .05})
129    report = bench.make_report(TRACK, MODEL, base, idea_res, {'mechanism_signature': sig, 'paired_deltas': diffs, 'permutation_p': bench.permutation_pvalue(diffs)})
130    Path('bench_report.json').write_text(json.dumps(report, indent=2))
131    print(json.dumps(report, indent=2))
132
133if __name__ == '__main__': main()