1import os, sys, json, math, time, random
 2import numpy as np
 3import torch
 4from torch import nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report, count_params
 7
 8EPOCHS=8; BATCH=128; NTRAIN=800; NTEST=400
 9LRS=[1e-3,3e-3,6e-3]
10
11class SU11Scan(nn.Module):
12    """Shared potential/readout architecture; exact or Euler transition is the only difference."""
13    def __init__(self, exact=True, pairs=16, h=.12, xi_scale=.07):
14        super().__init__(); self.exact=exact; self.pairs=pairs; self.h=h; self.xi_scale=xi_scale
15        self.fmap=nn.Linear(3,2*pairs); self.head=nn.Linear(4*pairs,1)
16    def forward(self,x,return_state=False):
17        B=x.shape[0]; seq=x.view(B,-1,3); dev=x.device
18        a=torch.ones(B,self.pairs,device=dev,dtype=torch.complex64); b=torch.zeros_like(a)
19        xi=torch.arange(self.pairs,device=dev,dtype=torch.float32)*self.xi_scale
20        for k in range(seq.shape[1]):
21            raw=self.fmap(seq[:,k]); f=torch.complex(raw[:,:self.pairs],raw[:,self.pairs:])
22            f=.7*torch.tanh(f.real)+1j*.7*torch.tanh(f.imag)
23            r=torch.abs(f).float(); phase=torch.exp((2j*math.pi*k*xi).to(torch.complex64))[None,:]
24            if self.exact:
25                q=self.h*r; c=torch.cosh(q).to(torch.complex64)
26                s=torch.where(r>1e-7,torch.sinh(q)/r,torch.full_like(r,self.h)).to(torch.complex64)
27            else:
28                c=torch.ones_like(r).to(torch.complex64); s=torch.full_like(r,self.h).to(torch.complex64)
29            aa=c*a+s*torch.conj(f)*phase*b; bb=c*b+s*f*torch.conj(phase)*a; a,b=aa,bb
30        feat=torch.cat([a.real,a.imag,b.real,b.imag],1); out=self.head(feat)
31        return (out,a,b) if return_state else out
32
33def seed_all(s):
34    random.seed(s); np.random.seed(s); torch.manual_seed(s)
35def ds(seed): return get_dataset('dynamics',seed,n_train=NTRAIN,n_test=NTEST)
36def train_one(kind,lr,seed,capture=False):
37    seed_all(seed); d=ds(seed); m=SU11Scan(exact=(kind=='idea'))
38    net,metric,_=train_model(m,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
39    if net is None: raise RuntimeError('training failed')
40    out={'metric':float(metric)}
41    if capture:
42        with torch.no_grad():
43            x=d['xte'].to(next(net.parameters()).device); _,a,b=net(x,True)
44            inv=(a.abs()**2-b.abs()**2-1).abs()
45            out.update(invariant_mean_error=float(inv.mean().cpu()),invariant_max_error=float(inv.max().cpu()),mean_abs_a=float(a.abs().mean().cpu()),mean_abs_b=float(b.abs().mean().cpu()))
46    return out
47def metric_fn(kind,lr): return lambda s: train_one(kind,lr,s)['metric']
48def main():
49    t=time.time(); grid=[{'lr':v} for v in LRS]
50    base=sweep_baseline(lambda c:metric_fn('baseline',c['lr']),grid)
51    idea_sweep=[]
52    for c in grid:
53        r=evaluate(metric_fn('idea',c['lr']),seeds=(0,1,2,3)); idea_sweep.append({'cfg':c,'mean':r['mean']})
54    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']; idea=evaluate(metric_fn('idea',best['lr']))
55    bs=train_one('baseline',base['best_cfg']['lr'],0,True); ins=train_one('idea',best['lr'],0,True)
56    sig={'prediction':'Exact exponential preserves |a|^2-|b|^2=1 while Euler drift accumulates.','predicted_exact_invariant_error':0.0,'observed_exact_mean_error':ins['invariant_mean_error'],'observed_exact_max_error':ins['invariant_max_error'],'observed_euler_mean_error':bs['invariant_mean_error'],'observed_euler_max_error':bs['invariant_max_error'],'confirmed':bool(ins['invariant_max_error']<1e-5 and bs['invariant_max_error']>ins['invariant_max_error']*10)}
57    rep=make_report('dynamics','su11_scan_matched',base,idea,{'mechanism_signature':sig,'idea_sweep':idea_sweep,'idea_best_cfg':best,'baseline_definition':'Euler discretization of the same learned SU(1,1) generator; exact exponential is the sole intervention.','protocol_note':'Identical trained systems, data, parameters, optimizer, epochs, and learning-rate grid; lower MSE is better.'})
58    rep['timing_seconds']=time.time()-t; rep['parameter_counts']={'baseline':count_params(SU11Scan(False)),'idea':count_params(SU11Scan(True))}
59    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
60    print(json.dumps(rep,indent=2))
61if __name__=='__main__': main()