1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  9
 10EPOCHS = 12
 11NTR, NTE = 400, 200
 12DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 13
 14
 15def fit_mode(c):
 16    c = np.asarray(c, dtype=float)
 17    best = (1e99, .2, .05)
 18    k = np.arange(len(c), dtype=float)
 19    for r in np.linspace(.20, .995, 80):
 20        for w in np.linspace(.05, 1.45, 90):
 21            A = (r ** k)[:, None] * np.c_[np.cos(w*k), np.sin(w*k)]
 22            ab = np.linalg.lstsq(A, c, rcond=None)[0]
 23            err = float(np.mean((A @ ab - c) ** 2))
 24            if err < best[0]: best = (err, r, w)
 25    return float(best[1]), float(best[2])
 26
 27
 28def math_sanity(seed=0):
 29    rng = np.random.default_rng(seed)
 30    r, w, T = .93, .42, 5000
 31    R = np.array([[np.cos(w), -np.sin(w)], [np.sin(w), np.cos(w)]])
 32    z = np.zeros((T, 2))
 33    for t in range(1, T): z[t] = r * R @ z[t-1] + .15 * rng.normal(size=2)
 34    mu = z[:, 0].mean()
 35    c = np.array([np.mean((z[:-k,0]-mu)*(z[k:,0]-mu)) if k else np.var(z[:,0])
 36                  for k in range(31)])
 37    c /= c[0]
 38    fr, fw = fit_mode(c)
 39    return {'true_radius': r, 'fitted_radius': fr, 'true_omega': w,
 40            'fitted_omega': fw, 'radius_abs_error': abs(fr-r),
 41            'frequency_abs_error': abs(fw-w), 'pass': abs(fr-r)<.03 and abs(fw-w)<.04}
 42
 43
 44def seed_all(seed):
 45    np.random.seed(seed); torch.manual_seed(seed)
 46    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 47
 48
 49def recurrent_radius(model):
 50    W = model.rnn.weight_hh_l0
 51    return torch.linalg.matrix_norm(W, ord=2)
 52
 53
 54def hidden_observations(model, x, noise_std=.03):
 55    seq = x.view(x.shape[0], -1, 3)
 56    h = torch.zeros(1, x.shape[0], model.rnn.hidden_size, device=x.device)
 57    hs = []
 58    # Use the GRU's actual trained recurrent dynamics, with injected noise.
 59    old_cudnn = torch.backends.cudnn.enabled
 60    try:
 61        torch.backends.cudnn.enabled = False
 62        for t in range(seq.shape[1]):
 63            inp = seq[:, t:t+1] + noise_std * torch.randn_like(seq[:, t:t+1])
 64            out, h = model.rnn(inp, h)
 65            hs.append(h[-1])
 66        return torch.stack(hs, dim=1)
 67    finally:
 68        torch.backends.cudnn.enabled = old_cudnn
 69
 70
 71def resonance_penalty(model, x, r_max=.90, beta=.02):
 72    # Radius term is differentiable; mode estimates are measured periodically
 73    # for the signature and stop-gradient by design as prescribed by the idea.
 74    radius = recurrent_radius(model)
 75    return beta * torch.relu(radius - r_max) ** 2, float(radius.detach())
 76
 77
 78def train_baseline(seed, cfg):
 79    seed_all(seed)
 80    d = get_dataset('dynamics', seed, NTR, NTE)
 81    model = make_model('rnn_small', d['input_shape'], d['out_dim'])
 82    opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 83    lossf = nn.MSELoss()
 84    dev = torch.device(DEVICE)
 85    try:
 86        model.to(dev); xtr,ytr=d['xtr'].to(dev),d['ytr'].to(dev)
 87        for _ in range(EPOCHS):
 88            p=torch.randperm(len(xtr),device=dev)
 89            for i in range(0,len(xtr),128):
 90                ix=p[i:i+128]; loss=lossf(model(xtr[ix]),ytr[ix])
 91                opt.zero_grad(); loss.backward(); opt.step()
 92        with torch.no_grad():
 93            return float(lossf(model(d['xte'].to(dev)), d['yte'].to(dev)))
 94    except RuntimeError:
 95        model.cpu(); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
 96        xtr,ytr=d['xtr'],d['ytr']
 97        for _ in range(EPOCHS):
 98            p=torch.randperm(len(xtr))
 99            for i in range(0,len(xtr),128):
100                ix=p[i:i+128]; loss=lossf(model(xtr[ix]),ytr[ix]); opt.zero_grad(); loss.backward(); opt.step()
101        with torch.no_grad(): return float(lossf(model(d['xte']),d['yte']))
102
103
104def train_idea(seed, cfg, collect=False):
105    seed_all(seed)
106    d=get_dataset('dynamics',seed,NTR,NTE)
107    model=make_model('rnn_small',d['input_shape'],d['out_dim'])
108    lossf=nn.MSELoss(); dev=torch.device(DEVICE)
109    try: model.to(dev); xtr,ytr=d['xtr'].to(dev),d['ytr'].to(dev)
110    except RuntimeError: dev=torch.device('cpu'); model.cpu(); xtr,ytr=d['xtr'],d['ytr']
111    opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
112    radii=[]; mode_rows=[]
113    for ep in range(EPOCHS):
114        p=torch.randperm(len(xtr),device=dev)
115        for i in range(0,len(xtr),128):
116            ix=p[i:i+128]; pred=model(xtr[ix]); task=lossf(pred,ytr[ix])
117            pen,rad=resonance_penalty(model,xtr[ix],cfg['r_max'],cfg['beta'])
118            loss=task+pen; opt.zero_grad(); loss.backward(); opt.step(); radii.append(rad)
119        if ep in (0,EPOCHS-1):
120            with torch.no_grad():
121                hs=hidden_observations(model,xtr[:128])
122                obs=hs[:,:,0].detach().cpu().numpy(); obs-=obs.mean(axis=1,keepdims=True)
123                c=np.array([(obs[:,:-k]*obs[:,k:]).mean() if k else (obs*obs).mean() for k in range(6)])
124                c/=max(c[0],1e-8); fr,fw=fit_mode(c); mode_rows.append({'epoch':ep,'fitted_radius':fr,'fitted_omega':fw})
125    with torch.no_grad(): metric=float(lossf(model(d['xte'].to(dev)),d['yte'].to(dev)))
126    if collect: return metric, {'radius_mean':float(np.mean(radii)),'mode_rows':mode_rows,'model':model}
127    return metric
128
129
130def main():
131    lr_grid=[1e-3,3e-3,6e-3]
132    # Union parity: both methods are evaluated at every lr and shared WD values.
133    grid=[{'lr':lr,'weight_decay':wd} for lr in lr_grid for wd in (0.0,1e-4)]
134    base=sweep_baseline(lambda cfg: lambda seed: train_baseline(seed,cfg),grid)
135    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)]
136    # All idea settings are already in the baseline union grid.
137    idea_trials=[]
138    for cfg in idea_cfgs:
139        rr=evaluate(lambda seed: train_idea(seed,cfg))
140        idea_trials.append({'cfg':cfg,'result':rr})
141    idea=min(idea_trials,key=lambda z:z['result']['mean'])
142    sig_metric,sig=train_idea(0,idea['cfg'],collect=True)
143    signature={'predicted_vs_observed':{'target_radius':idea['cfg']['r_max'],'observed_radius_mean':sig['radius_mean'],
144        'fitted_mode_rows':sig['mode_rows']},'confirmed':bool(sig['radius_mean'] <= idea['cfg']['r_max']+.05)}
145    rep=make_report('dynamics','rnn_small',base,idea['result'],signature)
146    rep['idea_trials']=idea_trials; rep['math_sanity']=math_sanity()
147    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
148    print(json.dumps(rep,indent=2))
149
150if __name__=='__main__': main()