1import json, sys
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, train_model, sweep_baseline, make_report
 7
 8SEED=2112; EPOCHS=12; NTRAIN=1200; NTEST=400; BATCH=128; DT=0.05
 9
10class GenericGRU(nn.Module):
11    def __init__(self):
12        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.head=nn.Linear(64,1)
13    def forward(self,x):
14        _,h=self.rnn(x.view(x.shape[0],-1,3)); return self.head(h[-1])
15
16class ReducedGRU(nn.Module):
17    def __init__(self, penalty=0.0):
18        super().__init__(); self.penalty=penalty
19        self.rnn=nn.GRU(3,64,batch_first=True)
20        self.cy_raw=nn.Linear(64,1); self.cv_raw=nn.Linear(64,1)
21        self.cu=nn.Linear(64,1)
22        self.res=nn.Sequential(nn.Linear(64,32),nn.Tanh(),nn.Linear(32,1))
23    def forward(self,x):
24        seq=x.view(x.shape[0],-1,3); _,h=self.rnn(seq); z=h[-1]; last=seq[:,-1]
25        cy=torch.nn.functional.softplus(self.cy_raw(z))
26        # signed coefficient induced by positive alpha,beta,r: r-alpha-beta
27        cv=self.cv_raw(z)-torch.nn.functional.softplus(self.cv_raw(z)).detach()*0.0
28        acc=cy*last[:,0:1]+cv*last[:,1:2]+self.cu(z)*last[:,2:3]+0.1*self.res(z)
29        return last[:,0:1]+DT*last[:,1:2]+0.5*DT*DT*acc
30
31def seed_all(s):
32    np.random.seed(s); torch.manual_seed(s)
33    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
34
35def train_idea(cfg, seed, return_model=False):
36    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST)
37    model=ReducedGRU(cfg.get('penalty',0.0))
38    # train_model is the canonical loop; this is an end-to-end structural readout.
39    model,metric,hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH)
40    if return_model: return metric,model,ds
41    return metric
42
43def train_base(cfg, seed, return_model=False):
44    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST); model=GenericGRU()
45    model,metric,hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH)
46    if return_model: return metric,model,ds
47    return metric
48
49def main():
50    # Union parity: every idea lr is evaluated by the baseline sweep too.
51    grid=[{'lr':x,'epochs':EPOCHS} for x in (0.001,0.003,0.006)]
52    base=sweep_baseline(lambda c: lambda s: train_base(c,s),grid,seeds=tuple(range(4)))
53    idea_cfgs=[{'lr':x,'epochs':EPOCHS,'penalty':p} for x,p in ((0.001,0.0),(0.003,0.0),(0.006,0.0))]
54    idea_runs=[]
55    for cfg in idea_cfgs:
56        vals=[train_idea(cfg,s) for s in range(8)]
57        idea_runs.append({'cfg':cfg,'mean':float(np.mean(vals)),'per_seed':vals})
58    best=min(idea_runs,key=lambda r:r['mean']); idea={'best_cfg':best['cfg'],'sweep':idea_runs,'mean':best['mean'],'std':float(np.std(best['per_seed'])),'per_seed':best['per_seed'],'n':8}
59    # Signature from trained models: compare structural predicted acceleration to finite differences.
60    sig=[]
61    for s in range(8):
62        im,model,ds=train_idea(best['cfg'],s,True); model.eval(); x=ds['xte'];
63        model = model.cpu()
64        with torch.no_grad(): pred=model(x.cpu()).numpy().ravel()
65        a=x.numpy().reshape(len(x),-1,3)[:,-1,0]; v=x.numpy().reshape(len(x),-1,3)[:,-1,1]
66        observed=(ds['yte'].numpy().ravel()-a-DT*v)/(0.5*DT*DT)
67        sig.append([float(np.corrcoef(pred,ds['yte'].numpy().ravel())[0,1]),float(np.sqrt(np.mean((pred-ds['yte'].numpy().ravel())**2)))])
68    report=make_report('dynamics','rnn_small',base,idea,{'prediction':'second-order structured dynamics improves extrapolative neural prediction','trained_model_signature':{'output_target_corr_mean':float(np.mean([x[0] for x in sig])),'output_rmse_mean':float(np.mean([x[1] for x in sig])),'confirmed':False,'note':'This built-in task has no aggregate two-compartment ground truth, so the exact cy=beta*r invariant cannot be quantitatively retested.'},'structural_match':'dynamics stability/control'})
69    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
70    print(json.dumps(report,indent=2))
71if __name__=='__main__': main()