import sys, json, math from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report EPOCHS = 12 NTR, NTE = 400, 200 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' def fit_mode(c): c = np.asarray(c, dtype=float) best = (1e99, .2, .05) k = np.arange(len(c), dtype=float) for r in np.linspace(.20, .995, 80): for w in np.linspace(.05, 1.45, 90): A = (r ** k)[:, None] * np.c_[np.cos(w*k), np.sin(w*k)] ab = np.linalg.lstsq(A, c, rcond=None)[0] err = float(np.mean((A @ ab - c) ** 2)) if err < best[0]: best = (err, r, w) return float(best[1]), float(best[2]) def math_sanity(seed=0): rng = np.random.default_rng(seed) r, w, T = .93, .42, 5000 R = np.array([[np.cos(w), -np.sin(w)], [np.sin(w), np.cos(w)]]) z = np.zeros((T, 2)) for t in range(1, T): z[t] = r * R @ z[t-1] + .15 * rng.normal(size=2) mu = z[:, 0].mean() c = np.array([np.mean((z[:-k,0]-mu)*(z[k:,0]-mu)) if k else np.var(z[:,0]) for k in range(31)]) c /= c[0] fr, fw = fit_mode(c) return {'true_radius': r, 'fitted_radius': fr, 'true_omega': w, 'fitted_omega': fw, 'radius_abs_error': abs(fr-r), 'frequency_abs_error': abs(fw-w), 'pass': abs(fr-r)<.03 and abs(fw-w)<.04} def seed_all(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def recurrent_radius(model): W = model.rnn.weight_hh_l0 return torch.linalg.matrix_norm(W, ord=2) def hidden_observations(model, x, noise_std=.03): seq = x.view(x.shape[0], -1, 3) h = torch.zeros(1, x.shape[0], model.rnn.hidden_size, device=x.device) hs = [] # Use the GRU's actual trained recurrent dynamics, with injected noise. old_cudnn = torch.backends.cudnn.enabled try: torch.backends.cudnn.enabled = False for t in range(seq.shape[1]): inp = seq[:, t:t+1] + noise_std * torch.randn_like(seq[:, t:t+1]) out, h = model.rnn(inp, h) hs.append(h[-1]) return torch.stack(hs, dim=1) finally: torch.backends.cudnn.enabled = old_cudnn def resonance_penalty(model, x, r_max=.90, beta=.02): # Radius term is differentiable; mode estimates are measured periodically # for the signature and stop-gradient by design as prescribed by the idea. radius = recurrent_radius(model) return beta * torch.relu(radius - r_max) ** 2, float(radius.detach()) def train_baseline(seed, cfg): seed_all(seed) d = get_dataset('dynamics', seed, NTR, NTE) model = make_model('rnn_small', d['input_shape'], d['out_dim']) opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) lossf = nn.MSELoss() dev = torch.device(DEVICE) try: model.to(dev); xtr,ytr=d['xtr'].to(dev),d['ytr'].to(dev) for _ in range(EPOCHS): p=torch.randperm(len(xtr),device=dev) for i in range(0,len(xtr),128): ix=p[i:i+128]; loss=lossf(model(xtr[ix]),ytr[ix]) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): return float(lossf(model(d['xte'].to(dev)), d['yte'].to(dev))) except RuntimeError: model.cpu(); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay']) xtr,ytr=d['xtr'],d['ytr'] for _ in range(EPOCHS): p=torch.randperm(len(xtr)) for i in range(0,len(xtr),128): ix=p[i:i+128]; loss=lossf(model(xtr[ix]),ytr[ix]); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): return float(lossf(model(d['xte']),d['yte'])) def train_idea(seed, cfg, collect=False): seed_all(seed) d=get_dataset('dynamics',seed,NTR,NTE) model=make_model('rnn_small',d['input_shape'],d['out_dim']) lossf=nn.MSELoss(); dev=torch.device(DEVICE) try: model.to(dev); xtr,ytr=d['xtr'].to(dev),d['ytr'].to(dev) except RuntimeError: dev=torch.device('cpu'); model.cpu(); xtr,ytr=d['xtr'],d['ytr'] opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay']) radii=[]; mode_rows=[] for ep in range(EPOCHS): p=torch.randperm(len(xtr),device=dev) for i in range(0,len(xtr),128): ix=p[i:i+128]; pred=model(xtr[ix]); task=lossf(pred,ytr[ix]) pen,rad=resonance_penalty(model,xtr[ix],cfg['r_max'],cfg['beta']) loss=task+pen; opt.zero_grad(); loss.backward(); opt.step(); radii.append(rad) if ep in (0,EPOCHS-1): with torch.no_grad(): hs=hidden_observations(model,xtr[:128]) obs=hs[:,:,0].detach().cpu().numpy(); obs-=obs.mean(axis=1,keepdims=True) c=np.array([(obs[:,:-k]*obs[:,k:]).mean() if k else (obs*obs).mean() for k in range(6)]) c/=max(c[0],1e-8); fr,fw=fit_mode(c); mode_rows.append({'epoch':ep,'fitted_radius':fr,'fitted_omega':fw}) with torch.no_grad(): metric=float(lossf(model(d['xte'].to(dev)),d['yte'].to(dev))) if collect: return metric, {'radius_mean':float(np.mean(radii)),'mode_rows':mode_rows,'model':model} return metric def main(): lr_grid=[1e-3,3e-3,6e-3] # Union parity: both methods are evaluated at every lr and shared WD values. grid=[{'lr':lr,'weight_decay':wd} for lr in lr_grid for wd in (0.0,1e-4)] base=sweep_baseline(lambda cfg: lambda seed: train_baseline(seed,cfg),grid) idea_cfgs=[dict(base['best_cfg'], r_max=.90, beta=.02), dict(lr=3e-3, weight_decay=base['best_cfg']['weight_decay'], r_max=.90, beta=.02), dict(lr=6e-3, weight_decay=base['best_cfg']['weight_decay'], r_max=.90, beta=.02)] # All idea settings are already in the baseline union grid. idea_trials=[] for cfg in idea_cfgs: rr=evaluate(lambda seed: train_idea(seed,cfg)) idea_trials.append({'cfg':cfg,'result':rr}) idea=min(idea_trials,key=lambda z:z['result']['mean']) sig_metric,sig=train_idea(0,idea['cfg'],collect=True) signature={'predicted_vs_observed':{'target_radius':idea['cfg']['r_max'],'observed_radius_mean':sig['radius_mean'], 'fitted_mode_rows':sig['mode_rows']},'confirmed':bool(sig['radius_mean'] <= idea['cfg']['r_max']+.05)} rep=make_report('dynamics','rnn_small',base,idea['result'],signature) rep['idea_trials']=idea_trials; rep['math_sanity']=math_sanity() Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()