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, make_model, train_model, evaluate, sweep_baseline, make_report
 8
 9SEED=449; TAU=.20; ALPHA=(.5,.5); HIDDEN=64
10
11class SwitchedSSM(nn.Module):
12    def __init__(self, input_shape, out_dim, comm_lambda=0.0):
13        super().__init__(); self.d=HIDDEN; self.lam=float(comm_lambda)
14        eye=torch.eye(self.d)
15        self.A=nn.Parameter(-.7*eye[None].repeat(2,1,1)+.03*torch.randn(2,self.d,self.d))
16        self.U=nn.Parameter(.08*torch.randn(2,3,self.d)); self.head=nn.Linear(self.d,out_dim)
17    def forward(self,x):
18        z=x.reshape(x.shape[0],-1,3); h=torch.zeros(x.shape[0],self.d,device=x.device,dtype=x.dtype)
19        B=[ALPHA[i]*TAU*self.A[i] for i in range(2)]
20        E=[torch.matrix_exp(b) for b in B]
21        for t in range(z.shape[1]):
22            for i in range(2): h=h@E[i].T+(ALPHA[i]*z[:,t,:])@self.U[i]
23        return self.head(h)
24    def comm_penalty(self):
25        B0=ALPHA[0]*TAU*self.A[0]; B1=ALPHA[1]*TAU*self.A[1]
26        return ((B0@B1-B1@B0)**2).sum()
27
28def train_idea(model, ds, epochs=30, lr=.003, batch=128, weight_decay=0.):
29    # The intervention is a new loss, hence a local loop; all other settings match train_model.
30    try: device='cuda' if torch.cuda.is_available() else 'cpu'; model=model.to(device)
31    except Exception: device='cpu'; model=model.to(device)
32    x,y=ds['xtr'].to(device),ds['ytr'].to(device); opt=torch.optim.Adam(model.parameters(),lr=lr,weight_decay=weight_decay)
33    lossf=nn.CrossEntropyLoss() if ds['task']=='classification' else nn.MSELoss(); hist=[]
34    for _ in range(epochs):
35        model.train(); order=torch.randperm(len(x),device=device); total=0.
36        for ix in order.split(batch):
37            pred=model(x[ix]); loss=lossf(pred,y[ix])+model.lam*model.comm_penalty()
38            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.); opt.step(); total+=float(loss.detach())
39        hist.append(total/max(1,(len(x)+batch-1)//batch))
40    model.eval()
41    with torch.no_grad():
42        pred=model(ds['xte'].to(device)); metric=float(lossf(pred,ds['yte'].to(device)).cpu())
43    return model,metric,hist
44
45def idea_factory(cfg, seed, ds):
46    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
47    return train_idea(SwitchedSSM(ds['input_shape'],ds['out_dim'],cfg['comm_lambda']),ds,epochs=cfg['epochs'],lr=cfg['lr'])
48
49def baseline_factory(cfg, seed, ds):
50    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
51    # Matched baseline: exactly the same switched architecture, but no proposed penalty.
52    net=SwitchedSSM(ds['input_shape'],ds['out_dim'],comm_lambda=0.0)
53    return train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=0.)
54
55def eval_side(factory, cfg, seeds):
56    vals=[]; models=[]
57    for s in seeds:
58        ds=get_dataset('dynamics',s,n_train=400,n_test=200)
59        out=factory(cfg,s,ds); models.append(out[0]); vals.append(float(out[1]))
60    return {'config':cfg,'per_seed':vals,'mean':float(np.mean(vals)),'models':models}
61
62def model_input_sensitivity(models, seeds):
63    vals=[]
64    for model,s in zip(models,seeds):
65        ds=get_dataset('dynamics',s,n_train=8,n_test=8); dev=next(model.parameters()).device
66        x=ds['xte'][:4].to(dev).clone().requires_grad_(True)
67        model.eval(); out=model(x).sum(); g=torch.autograd.grad(out,x)[0]
68        vals.append(float(g.detach().norm()/max(1,g.numel())**0.5))
69    return float(np.mean(vals))
70
71def diagnostics(models,seeds):
72    comm=[]; gain=[]
73    for model,s in zip(models,seeds):
74        A=model.A.detach().cpu(); B=[.5*TAU*A[i] for i in range(2)]
75        phi=torch.matrix_exp(B[1])@torch.matrix_exp(B[0]); c=torch.linalg.matrix_norm(B[0]@B[1]-B[1]@B[0]).item()**2
76        gain.append(float(torch.linalg.svdvals(phi).max())); comm.append(c)
77    return {'observed_commutator_mean':float(np.mean(comm)),'observed_cycle_gain_mean':float(np.mean(gain)),'n_models':len(models)}
78
79def main():
80    seeds=list(range(8)); grid=[{'lr':lr,'epochs':20} for lr in (.001,.003,.009)]
81    # Equal union: baseline evaluates every learning rate considered by the idea.
82    base_sweep=[]
83    for cfg in grid: base_sweep.append(eval_side(baseline_factory,cfg,seeds[:4]))
84    best=min(base_sweep,key=lambda r:r['mean']); best_cfg=best['config']
85    base_full=eval_side(baseline_factory,best_cfg,seeds)
86    base_block={'best_cfg':best_cfg,'sweep':[{'config':r['config'],'mean':r['mean'],'per_seed':r['per_seed']} for r in base_sweep],'full':{'per_seed':base_full['per_seed'],'mean':base_full['mean']}}
87    idea_grid=[dict(best_cfg,comm_lambda=lam) for lam in (.0,.01,.05)]
88    idea_runs=[eval_side(lambda c,s,d: idea_factory(c,s,d),cfg,seeds) for cfg in idea_grid]
89    idea=min(idea_runs,key=lambda r:r['mean'])
90    report=make_report('dynamics','rnn_small',base_block,{'config':idea['config'],'per_seed':idea['per_seed'],'mean':idea['mean']},extra={'mechanism_signature':{'prediction':'commutator regularization lowers trained ordered-cycle noncommutativity and cycle gain','observed':diagnostics(idea['models'],seeds),'baseline_observed':{'cycle_gain_mean':diagnostics(base_full['models'],seeds)['observed_cycle_gain_mean'],'observed_commutator_mean':diagnostics(base_full['models'],seeds)['observed_commutator_mean'],'input_sensitivity_mean':model_input_sensitivity(base_full['models'],seeds),'note':'Matched switched SSM baseline with lambda=0'},'idea_input_sensitivity_mean':model_input_sensitivity(idea['models'],seeds),'confirmed': diagnostics(idea['models'],seeds)['observed_commutator_mean'] < diagnostics(base_full['models'],seeds)['observed_commutator_mean']},'track_justification':'Dynamics is structurally matched to switched latent stability and Lyapunov contraction.'})
91    report['idea_sweep']=[{'config':r['config'],'mean':r['mean'],'per_seed':r['per_seed']} for r in idea_runs]
92    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
93if __name__=='__main__': main()