1import sys, json, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 8
 9SEEDS=tuple(range(8)); EPOCHS=10; BATCH=128
10LR_GRID=[1e-3,3e-3,1e-2]
11SIGMA_GRID=[0.01,0.03,0.1]
12
13class MatchedRNN(nn.Module):
14    """Matched recurrent baseline/idea; noise is added after every GRUCell step."""
15    def __init__(self, out_dim, sigma=0.0):
16        super().__init__(); self.sigma=float(sigma)
17        self.inp=nn.Linear(3,64); self.cell=nn.GRUCell(64,64); self.head=nn.Linear(64,out_dim)
18    def forward(self,x,return_hidden=False):
19        seq=x.reshape(x.shape[0],-1,3)
20        h=torch.tanh(self.inp(seq[:,0])); states=[]
21        for t in range(seq.shape[1]):
22            z=self.inp(seq[:,t])
23            h=self.cell(z,h)
24            if self.training and self.sigma>0: h=h+self.sigma*torch.randn_like(h)
25            states.append(h)
26        hs=torch.stack(states,1)
27        out=self.head(h)
28        return (out,hs) if return_hidden else out
29
30def seed_all(s):
31    random.seed(s); np.random.seed(s); torch.manual_seed(s)
32    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
33
34def train_one(seed,cfg,with_noise=False,signature=False):
35    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
36    sigma=cfg.get('sigma',0.0) if with_noise else 0.0
37    net=MatchedRNN(ds['out_dim'],sigma)
38    trained,metric,_=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH,log=lambda *a,**k:None)
39    if trained is None: return float('inf') if not signature else {'metric':float('inf')}
40    if not signature: return float(metric)
41    trained.eval(); dev=next(trained.parameters()).device; x=ds['xte'].to(dev)
42    with torch.no_grad(): _,h=trained(x,return_hidden=True)
43    return {'metric':float(metric),'hidden_variance':float(h.var().item()),
44            'temporal_energy':float(((h[:,1:]-h[:,:-1])**2).mean().item()),'sigma':float(sigma)}
45
46def main():
47    grid=[{'lr':lr,'epochs':EPOCHS,'sigma':0.0} for lr in LR_GRID]
48    base=sweep_baseline(lambda c:lambda s:train_one(s,c,False),grid,seeds=(0,1,2,3))
49    best_lr=base['best_cfg']['lr']
50    idea_cfgs=[{'lr':best_lr,'epochs':EPOCHS,'sigma':s} for s in SIGMA_GRID]
51    runs=[]
52    for c in idea_cfgs:
53        runs.append({'cfg':c,'result':evaluate(lambda s,c=c:train_one(s,c,True),seeds=SEEDS)})
54    best_cfg,best=min(((q['cfg'],q['result']) for q in runs),key=lambda z:z[1]['mean'])
55    sig0=train_one(0,best_cfg,True,True)
56    sig_base=train_one(0,{'lr':best_cfg['lr'],'epochs':EPOCHS,'sigma':0.0},False,True)
57    sig_rows=[]
58    for q in runs:
59        s=q['cfg']['sigma']; z=train_one(0,q['cfg'],True,True)
60        sig_rows.append({'sigma':s,'sigma2':s*s,'observed_temporal_energy':z['temporal_energy']})
61    x=np.array([r['sigma2'] for r in sig_rows]); y=np.array([r['observed_temporal_energy'] for r in sig_rows])
62    slope=float(np.dot(x,y)/np.dot(x,x)); r2=float(1-np.sum((y-slope*x)**2)/max(np.sum((y-y.mean())**2),1e-12))
63    signature={'prediction':'trained recurrent hidden temporal energy is nonzero and increases with sigma^2',
64      'trained_model_observed':sig_rows,'baseline_sigma0':sig_base,'best_sigma':sig0,
65      'through_origin_slope':slope,'through_origin_r2':r2,
66      'confirmed':bool(r2>0.8 and y[-1]>y[0]),
67      'note':'measured from trained benchmark model hidden trajectories'}
68    report=make_report('dynamics','rnn_small',base,best,extra={'mechanism_signature':signature,
69      'idea_sweep':runs,'track_match':'stability/control and recurrent state evolution -> dynamics'})
70    report['custom_track']=None
71    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
72if __name__=='__main__': main()