1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from scipy.linalg import expm
  6from torch import nn
  7
  8SEED=2704
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10
 11def cross_matrix(n):
 12    x,y,z=n
 13    return np.array([[0,-z,y],[z,0,-x],[-y,x,0]],float)
 14
 15def rot(axis, angle):
 16    axis=np.asarray(axis,float); axis=axis/np.linalg.norm(axis)
 17    K=cross_matrix(axis)
 18    return np.eye(3)*math.cos(angle)+(1-math.cos(angle))*np.outer(axis,axis)+math.sin(angle)*K
 19
 20def spin_seq(events, axes, angles, s0=np.array([0.31,-0.52,0.796])):
 21    s=np.asarray(s0,float); s=s/np.linalg.norm(s)
 22    for e in events: s=rot(axes[e],angles[e])@s
 23    return s
 24
 25def mechanism_checks():
 26    # Prediction 1: rotations preserve norm (up to floating point).
 27    axes=[np.array([1.,0,0]),np.array([0,1.,0])]
 28    rng=np.random.default_rng(SEED)
 29    maxerr=0
 30    expm_err=0
 31    for _ in range(100):
 32        a=rng.normal(size=3); a/=np.linalg.norm(a); q=rng.uniform(-math.pi,math.pi)
 33        expm_err=max(expm_err, np.max(np.abs(rot(a,q)-expm(q*cross_matrix(a)))))
 34    for _ in range(1000):
 35        s=rng.normal(size=3); s/=np.linalg.norm(s)
 36        for e in rng.integers(0,2,100): s=rot(axes[e],rng.uniform(-math.pi,math.pi))@s
 37        maxerr=max(maxerr,abs(np.linalg.norm(s)-1))
 38    # Prediction 2: AB-vs-BA difference is zero for parallel axes; for small equal
 39    # angles orthogonal axes it scales quadratically, with commutator coefficient.
 40    orth=[]; par=[]
 41    s0=np.array([.31,-.52,.796]); s0/=np.linalg.norm(s0)
 42    for th in np.logspace(-3,-0.2,10):
 43        orth.append(np.linalg.norm(spin_seq([0,1],axes,[th,th],s0)-spin_seq([1,0],axes,[th,th],s0)))
 44        pa=[np.array([1.,0,0]),np.array([1.,0,0])]
 45        par.append(np.linalg.norm(spin_seq([0,1],pa,[th,th],s0)-spin_seq([1,0],pa,[th,th],s0)))
 46    slope=np.polyfit(np.log(np.logspace(-3,-0.2,10)[:5]),np.log(np.array(orth)[:5]),1)[0]
 47    # Prediction 3: sensitivity grows with nonparallelity, using fixed small angles.
 48    th=.08; deltas=[]; sins=[]
 49    for phi in np.linspace(0,.5*math.pi,9):
 50        a=[np.array([1.,0,0]),np.array([math.cos(phi),math.sin(phi),0.])]
 51        deltas.append(np.linalg.norm(spin_seq([0,1],a,[th,th],s0)-spin_seq([1,0],a,[th,th],s0)))
 52        sins.append(abs(math.sin(phi)))
 53    corr=np.corrcoef(deltas,sins)[0,1]
 54    # Quantitative predictions: norm error=0; exponent is 2; parallel effect=0;
 55    # fixed-angle order sensitivity is proportional to sin(axis angle).
 56    return {'predictions': {
 57        'rotation_exponential_max_error_predicted_0': expm_err,
 58        'sphere_norm_error_predicted_0': maxerr,
 59        'small_angle_order_effect_exponent_predicted_2': slope,
 60        'parallel_axis_order_effect_predicted_0': max(par),
 61        'axis_misalignment_dependence_corr_predicted_1': corr
 62    }, 'orthogonal_delta_at_theta_.08':deltas[-1]}
 63
 64class CountModel(nn.Module):
 65    def __init__(self):
 66        super().__init__(); self.fc=nn.Linear(2,2)
 67    def forward(self,x): return self.fc(torch.stack([(x==0).float().sum(1),(x==1).float().sum(1)],1))
 68
 69class GRUModel(nn.Module):
 70    def __init__(self, spin=False):
 71        super().__init__(); self.spin=spin; self.gru=nn.GRU(2,16,batch_first=True)
 72        if spin:
 73            self.axis=nn.Parameter(torch.tensor([[1.,0,0],[0,1.,0]]))
 74            self.theta=nn.Parameter(torch.tensor([1.,1.]))
 75            self.fc=nn.Linear(19,2)
 76        else: self.fc=nn.Linear(16,2)
 77    def forward(self,x):
 78        h,_=self.gru(torch.nn.functional.one_hot(x,2).float()); z=h[:,-1]
 79        if not self.spin: return self.fc(z)
 80        axes=self.axis/(self.axis.norm(dim=1,keepdim=True)+1e-8)
 81        s=torch.zeros(x.size(0),3,device=x.device); s[:,2]=1
 82        for t in range(x.size(1)):
 83            n=axes[x[:,t]]; th=torch.tanh(self.theta[x[:,t]])*math.pi
 84            s=s*torch.cos(th[:,None])+torch.cross(n,s,dim=1)*torch.sin(th[:,None])+n*(n*s).sum(1,keepdim=True)*(1-torch.cos(th[:,None]))
 85        return self.fc(torch.cat([z,s],1))
 86
 87def dataset(n, L=64):
 88    # Equal counts and identical suffix; label is ordering of first two events.
 89    x=np.tile(np.array([0,1]*((L-2)//2)),(n,1)); x=np.concatenate([np.zeros((n,1),int),x],1)[:,:L]
 90    # correct dimensions and balanced suffix
 91    suffix=np.array([0,1]*((L-2)//2),int)
 92    X=[]; y=[]
 93    for i in range(n):
 94        lab=i%2; prefix=[0,1] if lab==0 else [1,0]
 95        X.append(prefix+suffix.tolist()); y.append(lab)
 96    return torch.tensor(np.array(X),dtype=torch.long),torch.tensor(y)
 97
 98def train(model, Xtr,ytr,Xte,yte, epochs=100):
 99    opt=torch.optim.Adam(model.parameters(),lr=.02); lossfn=nn.CrossEntropyLoss()
100    for _ in range(epochs):
101        opt.zero_grad(); lossfn(model(Xtr),ytr).backward(); opt.step()
102    with torch.no_grad():
103        pred=model(Xte).argmax(1); acc=(pred==yte).float().mean().item(); final=lossfn(model(Xte),yte).item()
104    return acc,final
105
106def mini_experiment():
107    X,y=dataset(512); Xtr,ytr=X[:384],y[:384]; Xte,yte=X[384:],y[384:]
108    out={}
109    for name,model in [('count',CountModel()),('gru',GRUModel()),('gru_spin',GRUModel(True))]:
110        out[name]=train(model,Xtr,ytr,Xte,yte)
111    return out
112
113def main():
114    checks=mechanism_checks(); exp=mini_experiment()
115    result={'seed':SEED,'mechanism_checks':checks,'classification':exp}
116    Path('results.json').write_text(json.dumps(result,indent=2))
117    print(json.dumps(result,indent=2))
118if __name__=='__main__': main()