1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 17
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 10    if device == 'cuda': torch.zeros(1, device='cuda')
 11except Exception:
 12    device = 'cpu'
 13
 14torch.set_num_threads(4)
 15
 16def orbit_tangent(x):
 17    return torch.stack((-x[:, 1], x[:, 0]), dim=1)
 18
 19def projector(R, eps=1e-6):
 20    # R is [batch,n,r]; G=I. Handles rank-deficient R through damping.
 21    b, n, r = R.shape
 22    eye = torch.eye(r, device=R.device, dtype=R.dtype).expand(b, r, r)
 23    return R @ torch.linalg.solve(R.transpose(1,2) @ R + eps*eye,
 24                                   R.transpose(1,2))
 25
 26def verify_math():
 27    x = torch.tensor([[1.2, -0.7]], dtype=torch.float64)
 28    r = orbit_tangent(x).unsqueeze(-1)
 29    P = projector(r, 1e-10)[0]
 30    rv = r[0,:,0:1]
 31    radial = x[0,:,None]
 32    z = torch.zeros_like(P)
 33    origin = torch.zeros(1,2,1,dtype=torch.float64)
 34    Po = projector(origin, 1e-6)[0]
 35    return {
 36        'projector_idempotence': float(torch.linalg.norm(P@P-P)),
 37        'orbit_projection_error': float(torch.linalg.norm(P@rv-rv)),
 38        'normal_projection_norm': float(torch.linalg.norm(P@radial)),
 39        'rank_regular': int(torch.linalg.matrix_rank(r[0])),
 40        'rank_singular_at_origin': int(torch.linalg.matrix_rank(origin[0])),
 41        'singular_projector_norm': float(torch.linalg.norm(Po))
 42    }
 43
 44def rotate(x, theta):
 45    c, s = torch.cos(theta), torch.sin(theta)
 46    return torch.stack((c*x[:,0]-s*x[:,1], s*x[:,0]+c*x[:,1]), 1)
 47
 48def make_data(n, seed, radius_noise=.10):
 49    gen = torch.Generator().manual_seed(seed)
 50    y = torch.randint(0,2,(n,),generator=gen)
 51    radius = (1.0 + (2*y.float()-1)*.22 + radius_noise*torch.randn(n,generator=gen))
 52    theta = 2*math.pi*torch.rand(n,generator=gen)
 53    x = torch.stack((radius*torch.cos(theta), radius*torch.sin(theta)),1)
 54    return x, y
 55
 56class Encoder(nn.Module):
 57    def __init__(self):
 58        super().__init__()
 59        self.body=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,16),nn.Tanh())
 60        self.head=nn.Linear(16,2)
 61    def forward(self,x): return self.head(self.body(x))
 62
 63def orbit_penalty(model, x):
 64    # Exact JVP of f in the supplied orbit direction, without materializing J.
 65    x = x.detach().requires_grad_(True)
 66    v = orbit_tangent(x)
 67    h = model(x)
 68    # Use a per-output VJP loop to obtain exact Jv without materializing J.
 69    # to obtain exact Jv for the 2-dimensional representation.
 70    cols=[]
 71    for k in range(h.shape[1]):
 72        g=torch.autograd.grad(h[:,k].sum(), x, create_graph=True, retain_graph=True)[0]
 73        cols.append((g*v).sum(1))
 74    jv=torch.stack(cols,1)
 75    return (jv*jv).sum(1).mean()
 76
 77def train(reg, seed=17, steps=350):
 78    torch.manual_seed(seed)
 79    model=Encoder().to(device)
 80    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 81    x,y=make_data(768,seed)
 82    x,y=x.to(device),y.to(device)
 83    for step in range(steps):
 84        ix=torch.randint(0,len(x),(96,),device=device)
 85        xb,yb=x[ix],y[ix]
 86        logits=model(xb)
 87        loss=nn.functional.cross_entropy(logits,yb)
 88        if reg: loss=loss + .15*orbit_penalty(model,xb)
 89        opt.zero_grad(); loss.backward(); opt.step()
 90    return model
 91
 92@torch.no_grad()
 93def evaluate(model, seed=101):
 94    x,y=make_data(3000,seed)
 95    x,y=x.to(device),y.to(device)
 96    clean=(model(x).argmax(1)==y).float().mean().item()
 97    # unseen rotations, including angles outside a narrow training orientation
 98    th=(torch.rand(len(x),device=device)*2-1)*math.pi
 99    xr=rotate(x,th)
100    robust=(model(xr).argmax(1)==y).float().mean().item()
101    drift=(model(xr)-model(x)).norm(dim=1).mean().item()
102    return clean,robust,drift
103
104def main():
105    check=verify_math()
106    base=train(False); idea=train(True)
107    b=evaluate(base); q=evaluate(idea)
108    out={'device':device,'math_check':check,
109         'baseline':{'clean_accuracy':b[0],'rotated_accuracy':b[1],'representation_drift':b[2]},
110         'quotient_regularizer':{'clean_accuracy':q[0],'rotated_accuracy':q[1],'representation_drift':q[2]}}
111    with open('results.json','w') as f: json.dump(out,f,indent=2)
112    print(json.dumps(out,indent=2))
113
114if __name__=='__main__': main()