1import json, math, time, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 475
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9try:
 10    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 11    if device == 'cuda':
 12        torch.set_default_device('cuda')
 13        torch.zeros(1, device='cuda')
 14except Exception:
 15    device = 'cpu'
 16
 17# Real CMV-style specialization. M covers (0,1),(2,3),...; L covers
 18# (1,2),(3,4),...,(n-1,0). Scatter avoids in-place autograd mutations.
 19def apply_pair_factor(x, starts, theta, period):
 20    n = x.shape[-1]
 21    starts = torch.as_tensor(starts, device=x.device, dtype=torch.long)
 22    ends = (starts + 1) % n
 23    t = theta[torch.arange(len(starts), device=x.device) % period]
 24    c, s = torch.cos(t), torch.sin(t)
 25    u, v = x[:, starts], x[:, ends]
 26    vals = torch.stack((c[None, :] * u - s[None, :] * v,
 27                        s[None, :] * u + c[None, :] * v), dim=-1)
 28    idx = torch.stack((starts, ends), dim=-1).reshape(-1)
 29    # Each coordinate occurs once for each disjoint factor.
 30    return torch.zeros_like(x).scatter(1, idx[None, :].expand(x.shape[0], -1),
 31                                        vals.reshape(x.shape[0], -1))
 32
 33def cmv_apply(x, theta, period):
 34    n = x.shape[-1]
 35    y = apply_pair_factor(x, list(range(0, n, 2)), theta, period)
 36    return apply_pair_factor(y, list(range(1, n, 2)), theta, period)
 37
 38def cmv_matrix(n, theta, period):
 39    return cmv_apply(torch.eye(n, device=theta.device, dtype=theta.dtype), theta, period)
 40
 41class CMVRNN(nn.Module):
 42    def __init__(self, n, inp, period):
 43        super().__init__(); self.n=n
 44        self.theta=nn.Parameter(torch.empty(period).uniform_(-math.pi, math.pi))
 45        self.inp=nn.Linear(inp,n); self.out=nn.Linear(n,1)
 46    def forward(self, x):
 47        h=torch.zeros(x.shape[0],self.n,device=x.device)
 48        for t in range(x.shape[1]):
 49            h=torch.tanh(cmv_apply(h,self.theta,self.theta.numel())+self.inp(x[:,t]))
 50        return self.out(h).squeeze(-1)
 51
 52class DenseRNN(nn.Module):
 53    def __init__(self,n,inp,kind='orthogonal'):
 54        super().__init__(); self.n=n
 55        q,_=torch.linalg.qr(torch.randn(n,n)); self.W=nn.Parameter(q if kind=='orthogonal' else 1.05*q)
 56        self.inp=nn.Linear(inp,n); self.out=nn.Linear(n,1)
 57    def forward(self,x):
 58        h=torch.zeros(x.shape[0],self.n,device=x.device)
 59        for t in range(x.shape[1]): h=torch.tanh(h@self.W.T+self.inp(x[:,t]))
 60        return self.out(h).squeeze(-1)
 61
 62def math_check():
 63    n=32; b=128; x=torch.randn(b,n)
 64    th=torch.linspace(-1.2,1.1,8,device=x.device)
 65    U=cmv_matrix(n,th,8)
 66    err=float((U.T@U-torch.eye(n,device=x.device)).abs().max())
 67    ratios=[]; v=x
 68    for _ in range(100):
 69        v=cmv_apply(v,th,8); ratios.append((v.norm(dim=1)/x.norm(dim=1)).detach().cpu().numpy())
 70    ratios=np.concatenate(ratios)
 71    q,_=torch.linalg.qr(torch.randn(n,n,device=x.device)); expansive=1.05*q
 72    y=x.clone(); dense_rat=[]
 73    for _ in range(100):
 74        y=y@expansive.T; dense_rat.append((y.norm(dim=1)/x.norm(dim=1)).cpu().numpy())
 75    return {'orthogonality_max_error':err,
 76            'cmv_norm_ratio_max_abs_error':float(np.max(np.abs(ratios-1))),
 77            'cmv_norm_ratio_mean':float(ratios.mean()),
 78            'expansive_dense_norm_ratio_at_100':float(np.mean(dense_rat[-1]))}
 79
 80def train(model, steps=120, seq=40, batch=64):
 81    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 82    lossfn=nn.BCEWithLogitsLoss(); vals=[]; start=time.time()
 83    for step in range(steps):
 84        x=torch.randn(batch,seq,1,device=device); y=(x.sum(1).squeeze(-1)>0).float()
 85        opt.zero_grad(set_to_none=True); loss=lossfn(model(x),y); loss.backward(); opt.step()
 86        if step%20==0 or step==steps-1:
 87            with torch.no_grad():
 88                xx=torch.randn(512,seq,1,device=device); yy=(xx.sum(1).squeeze(-1)>0).float()
 89                acc=((model(xx)>0)==(yy>0.5)).float().mean().item()
 90            vals.append((step,float(loss),acc))
 91    return {'final_loss':vals[-1][1],'final_acc':vals[-1][2],
 92            'trace':vals,'seconds':time.time()-start}
 93
 94def main():
 95    torch.set_num_threads(4)
 96    check=math_check(); n=32; inp=1
 97    models={'dense_orthogonal':DenseRNN(n,inp,'orthogonal'),
 98            'cmv_period_8':CMVRNN(n,inp,8), 'cmv_period_32':CMVRNN(n,inp,32)}
 99    results={}
100    for name,m in models.items():
101        results[name]={'parameters':sum(p.numel() for p in m.parameters()), **train(m)}
102    out={'seed':SEED,'device':device,'math_check':check,'models':results}
103    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
104if __name__=='__main__': main()