1import sys, json, random
 2import numpy as np
 3sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7from bench import get_dataset, sweep_baseline, evaluate, make_report
 8
 9# Correct structural match: dynamics is the bench's stability/control task.
10# Same switching recurrent system on both sides; only the contraction loss differs.
11DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
12SEEDS = tuple(range(8)); M=4; RHO=.81; TAU=.05
13
14class SwitchingCore(nn.Module):
15    def __init__(self, hidden=16):
16        super().__init__(); self.hidden=hidden
17        self.A=nn.Parameter(torch.randn(M,hidden,hidden)*.12)
18        self.B=nn.Parameter(torch.randn(M,3,hidden)*.12)
19        self.b=nn.Parameter(torch.zeros(M,hidden)); self.head=nn.Linear(hidden,1)
20    def step_all(self,z,x):
21        # returns all M next states, vectorized over modes and batch
22        return torch.tanh(torch.einsum('bh,mhk->mbk',z,self.A)+torch.einsum('bd,mdh->mbh',x,self.B)+self.b[:,None,:])
23    def forward(self,x,return_z=False):
24        seq=x.view(x.shape[0],-1,3); z=torch.zeros(x.shape[0],self.hidden,device=x.device)
25        for t in range(seq.shape[1]):
26            allz=self.step_all(z,seq[:,t]); mode=torch.clamp(((seq[:,t,2]+1.5)/3*M).long(),0,M-1)
27            z=allz[mode,torch.arange(z.shape[0],device=x.device)]
28        return (self.head(z),z) if return_z else self.head(z)
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,lr,lam,epochs=5,n_train=400,keep=False):
35    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=n_train,n_test=400)
36    net=SwitchingCore().to(DEVICE); opt=torch.optim.Adam(net.parameters(),lr=lr)
37    x,y=ds['xtr'].to(DEVICE),ds['ytr'].to(DEVICE)
38    for _ in range(epochs):
39        net.train(); p=torch.randperm(len(x),device=DEVICE)
40        for j in range(0,len(x),128):
41            ix=p[j:j+128]; pred=net(x[ix]); task=((pred-y[ix])**2).mean(); loss=task
42            if lam:
43                z=torch.randn(len(ix),net.hidden,device=DEVICE); x0=torch.zeros(len(ix),3,device=DEVICE)
44                nxt=net.step_all(z,x0)
45                e=(nxt*nxt).sum(2)-RHO*(z*z).sum(1)[None,:]
46                loss=task+lam*F.softplus(e/TAU).mean()*TAU
47            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.); opt.step()
48    net.eval()
49    with torch.no_grad(): metric=float(((net(ds['xte'].to(DEVICE))-ds['yte'].to(DEVICE))**2).mean())
50    return (metric,net) if keep else metric
51
52def fn(lr,lam): return lambda seed: train_one(seed,lr,lam)
53
54def mechanism_signature(lr,lam):
55    ratios=[]
56    for s in SEEDS:
57        _,net=train_one(s,lr,lam,keep=True)
58        with torch.no_grad():
59            z=torch.randn(128,net.hidden,device=DEVICE); x0=torch.zeros(128,3,device=DEVICE)
60            q=(net.step_all(z,x0)**2).sum(2)/((z*z).sum(1)[None,:]+1e-8)
61            ratios += q.cpu().numpy().ravel().tolist()
62    mx=float(np.max(ratios)); mean=float(np.mean(ratios))
63    return {'certificate':'V(z)=||z||^2; one-node path-complete graph; four all-mode self-edges',
64            'prediction':'edge contraction ratio <= rho', 'rho':RHO,
65            'observed_max_zero_input_ratio':mx,'observed_mean_ratio':mean,
66            'confirmed':bool(mx <= RHO*1.25),'n_trained_models':8}
67
68def main():
69    # All idea learning rates appear in the baseline grid: search-space parity.
70    grid=[{'lr':v,'weight_decay':0.0} for v in (.001,.003,.01)]
71    base=sweep_baseline(lambda c:fn(c['lr'],0.),grid,seeds=(0,1,2,3))
72    lr=base['best_cfg']['lr']
73    idea_sweep=[]
74    for lam in (.01,.05,.20):
75        idea_sweep.append({'cfg':{'lr':lr,'lambda_contract':lam},'result':evaluate(fn(lr,lam),SEEDS)})
76    best=min(idea_sweep,key=lambda q:q['result']['mean'])
77    rep=make_report('dynamics','switching_rnn_small',base,best['result'],
78                    {'signature':mechanism_signature(lr,best['cfg']['lambda_contract']),
79                     'idea_cfg':best['cfg'],
80                     'training_budget':'400 train/400 test samples, 5 epochs, batch 128; equal on both systems'})
81    rep['idea_sweep']=idea_sweep; rep['custom_track']=None; rep['device']=DEVICE
82    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
83    print(json.dumps(rep,indent=2))
84if __name__=='__main__': main()