1import sys, json, math, 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, train_model, evaluate, sweep_baseline, make_report
 8
 9SEEDS = tuple(range(8))
10SWEEP = (0, 1, 2, 3)
11LRS = [1e-3, 3e-3, 1e-2]
12EPOCHS = 20
13NTRAIN, NTEST = 800, 300
14
15class BaseTransformer(nn.Module):
16    def __init__(self, win=32, d=64, depth=2):
17        super().__init__(); self.win=win
18        self.inp=nn.Linear(1,d); self.pos=nn.Parameter(torch.zeros(1,win,d)); nn.init.normal_(self.pos,std=.02)
19        layer=nn.TransformerEncoderLayer(d,nhead=2,dim_feedforward=128,batch_first=True,dropout=0.0)
20        self.enc=nn.TransformerEncoder(layer,depth); self.head=nn.Linear(win*d,1)
21    def features(self,x):
22        h=self.inp(x.unsqueeze(-1))+self.pos
23        return self.enc(h)
24    def forward(self,x): return self.head(self.features(x).flatten(1))
25
26class SpinTransformer(BaseTransformer):
27    def __init__(self, win=32):
28        super().__init__(win)
29        self.axis=nn.Parameter(torch.tensor([[1.,0.,0.],[0.,1.,0.]]))
30        self.raw_theta=nn.Parameter(torch.tensor([0.7,0.7]))
31        self.head=nn.Linear(win*64+3,1)
32    def spin(self,x):
33        axes=self.axis/(self.axis.norm(dim=1,keepdim=True)+1e-8)
34        theta=math.pi*torch.tanh(self.raw_theta)
35        s=torch.zeros(x.shape[0],3,device=x.device,dtype=x.dtype); s[:,2]=1.
36        # Sign events provide two event classes while retaining continuous sequence input.
37        for t in range(x.shape[1]):
38            e=(x[:,t]>=0).long(); n=axes[e]; a=theta[e]
39            ca=torch.cos(a)[:,None]; sa=torch.sin(a)[:,None]
40            s=s*ca+torch.cross(n,s,dim=1)*sa+n*(n*s).sum(1,keepdim=True)*(1-ca)
41        return s
42    def forward(self,x): return self.head(torch.cat([self.features(x).flatten(1),self.spin(x)],1))
43
44def seed_all(seed):
45    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
46    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
47
48def train_one(seed, lr, idea):
49    seed_all(seed)
50    d=get_dataset('sequence', seed, n_train=NTRAIN, n_test=NTEST)
51    model=SpinTransformer() if idea else BaseTransformer()
52    _, metric, _=train_model(model,d,epochs=EPOCHS,lr=lr,batch=128)
53    # train_model returns the standard independent test MSE.
54    return float(metric)
55
56def behavior_signature():
57    seed_all(12345); d=get_dataset('sequence',12345,n_train=256,n_test=64)
58    m=SpinTransformer(); _,_,_=train_model(m,d,epochs=EPOCHS,lr=3e-3,batch=128)
59    m.eval(); device=next(m.parameters()).device; x=d['xte'].to(device)
60    with torch.no_grad():
61        s=m.spin(x); norms=s.norm(dim=1).cpu().numpy()
62        # The trained model's observed order sensitivity: compare each sequence
63        # with its reversed event-class ordering, holding real magnitudes fixed.
64        xr=x.clone(); xr=xr.flip(1)
65        delta=(s-m.spin(xr)).norm(dim=1).mean().item()
66        comm=torch.linalg.norm(torch.bmm(torch.eye(3).expand(2,3,3),torch.zeros(2,3,3))).item()
67        axes=(m.axis/(m.axis.norm(dim=1,keepdim=True)+1e-8)).detach().cpu().numpy()
68        observed_comm=float(abs(np.linalg.det(np.stack([axes[0],axes[1],np.cross(axes[0],axes[1])]))))
69    return {'prediction':'ordered noncommuting updates preserve sphere norm and create order sensitivity',
70            'predicted_norm_error':0.0,'observed_max_norm_error':float(np.max(np.abs(norms-1))),
71            'predicted_order_delta_nonzero':True,'observed_mean_reversal_spin_delta':delta,
72            'observed_axis_cross_norm':float(np.linalg.norm(np.cross(axes[0],axes[1]))),
73            'confirmed': bool(np.max(np.abs(norms-1)) < 1e-5 and delta > 1e-5),
74            'note':'All values are measured from the trained benchmark spin model.'}
75
76def main():
77    # Baseline sweep over exactly the union of learning rates used by the idea.
78    base=sweep_baseline(lambda cfg: (lambda seed: train_one(seed,cfg['lr'],False)),
79                        [{'lr':lr} for lr in LRS], seeds=SWEEP)
80    idea_sweep=[]
81    for lr in LRS:
82        r=evaluate(lambda seed, lr=lr: train_one(seed,lr,True), seeds=SWEEP)
83        idea_sweep.append({'cfg':{'lr':lr},'mean':r['mean']})
84    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
85    idea=evaluate(lambda seed: train_one(seed,best['lr'],True), seeds=SEEDS)
86    report=make_report('sequence','transformer_tiny',base,idea,{
87        'idea_sweep':idea_sweep,'selected_idea_cfg':best,
88        'behavior':behavior_signature()})
89    Path('bench_report.json').write_text(json.dumps(report,indent=2))
90    print(json.dumps(report,indent=2))
91if __name__=='__main__': main()