1import sys, json, random, numpy as np, torch
 2import torch.nn as nn
 3sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 4from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 5
 6SEEDS = tuple(range(8))
 7# Union used by both sides; baseline sweep includes all idea learning rates.
 8GRID = [{'lr': 1e-3, 'epochs': 12}, {'lr': 3e-3, 'epochs': 12}, {'lr': 1e-2, 'epochs': 12}]
 9
10class LiftRNN(nn.Module):
11    """Finite m=2 skew-dilation recurrent cell with Cayley propagation."""
12    def __init__(self, input_dim=3, out_dim=1, hidden=64, eta=.05):
13        super().__init__(); self.d=hidden; self.eta=eta
14        self.rawS=nn.Parameter(torch.randn(hidden,hidden)*.02)
15        self.rawA=nn.Parameter(torch.randn(hidden,hidden)*.02)
16        self.inp=nn.Linear(input_dim, 2*hidden)
17        self.head=nn.Linear(hidden,out_dim)
18    def forward(self,x):
19        x=x.view(x.shape[0],-1,3)
20        B=x.shape[0]; dev=x.device
21        S=self.rawS+self.rawS.T; A=self.rawA-self.rawA.T
22        z0=torch.zeros((2*self.d,),device=dev,dtype=x.dtype)
23        L=torch.zeros((2*self.d,2*self.d),device=dev,dtype=x.dtype)
24        L[:self.d,:self.d]=A; L[:self.d,self.d:]=S
25        L[self.d:,:self.d]=-S; L[self.d:,self.d:]=A
26        I=torch.eye(2*self.d,device=dev,dtype=x.dtype)
27        C=torch.linalg.solve(I-self.eta*L/2,I+self.eta*L/2)
28        z=torch.zeros((B,2*self.d),device=dev,dtype=x.dtype)
29        for k in range(x.shape[1]):
30            z=z@C.T + self.inp(x[:,k])
31            # bounded nonlinear readout is only the task head; recurrent lift remains linear/unitary
32        return self.head(z[:,:self.d])
33
34def baseline_fn(cfg, seed):
35    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
36    return make_base(3,1)
37def make_base(input_dim=3,out_dim=1):
38    class Base(nn.Module):
39        def __init__(self):
40            super().__init__(); self.rnn=nn.GRU(input_dim,64,batch_first=True); self.head=nn.Linear(64,out_dim)
41        def forward(self,x):
42            _,h=self.rnn(x.view(x.shape[0],-1,input_dim)); return self.head(h[-1])
43    return Base()
44def idea_fn(cfg, seed):
45    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
46    return LiftRNN(3,1,64,eta=.05)
47
48def run_one(fn, cfg, seed):
49    ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
50    net,metric,hist=train_model(fn(cfg,seed),ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128)
51    if net is None: raise RuntimeError('training failed')
52    return metric, net, ds
53
54def eval_cfg(fn,cfg,seeds=SEEDS):
55    vals=[]
56    for s in seeds: vals.append(run_one(fn,cfg,s)[0])
57    return {'per_seed':vals,'mean':float(np.mean(vals)),'std':float(np.std(vals,ddof=1))}
58
59def signature():
60    # measured on a trained benchmark model: compare Cayley norm preservation to Euler growth
61    metric,net,ds=run_one(idea_fn,GRID[1],0); net.eval(); d=net.d
62    with torch.no_grad():
63        S=net.rawS+net.rawS.T; A=net.rawA-net.rawA.T
64        L=torch.zeros((2*d,2*d)); L[:d,:d]=A; L[:d,d:]=S; L[d:,:d]=-S; L[d:,d:]=A
65        I=torch.eye(2*d); eta=net.eta
66        C=torch.linalg.solve(I-eta*L/2,I+eta*L/2); E=torch.randn(2*d)
67        zc=E.clone(); ze=E.clone()
68        for _ in range(20): zc=C@zc; ze=(I+eta*L)@ze
69        cay=float((zc.norm()/E.norm()-1).abs()); euler=float(ze.norm()/E.norm())
70    return {'prediction':'Cayley preserves lifted norm; explicit Euler grows it',
71            'observed_cayley_relative_norm_error':cay,
72            'observed_euler_20_step_norm_ratio':euler,
73            'trained_test_mse':float(metric),'confirmed':bool(cay<1e-5 and euler>1.0001)}
74
75def main():
76    # baseline sweep is explicitly over the same union as the idea settings
77    # Invoke the canonical harness sweep: make_fn(cfg) returns train_fn(seed).
78    from bench import sweep_baseline
79    def tuned_make(cfg):
80        def train_fn(seed):
81            return run_one(baseline_fn, cfg, int(seed))[0]
82        return train_fn
83    tuned=sweep_baseline(tuned_make, GRID, seeds=tuple(range(4)))
84    sweep=[]
85    for cfg in GRID:
86        r=eval_cfg(baseline_fn,cfg); sweep.append({'cfg':cfg,**r})
87    best=min(sweep,key=lambda r:r['mean'])
88    base={'best_cfg':best['cfg'],'sweep':sweep,'harness_tuning':tuned,'full':eval_cfg(baseline_fn,best['cfg'])}
89    idea=[]
90    for cfg in GRID:
91        r=eval_cfg(idea_fn,cfg); idea.append({'cfg':cfg,**r})
92    ibest=min(idea,key=lambda r:r['mean']); idea_res={k:ibest[k] for k in ('per_seed','mean','std')}
93    rep=make_report('dynamics','rnn_small',base,idea_res,{'idea_sweep':idea,'mechanism_signature':signature()})
94    rep['custom_track']=None
95    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
96    print(json.dumps(rep,indent=2))
97if __name__=='__main__': main()