1import json, math, random, time
  2import numpy as np
  3import torch
  4from sklearn.neighbors import kneighbors_graph
  5
  6SEED=487
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10except Exception:
 11    device=torch.device('cpu')
 12
 13# Circle data makes the intrinsic dimension and graph energy explicit.
 14N=720
 15ang=np.linspace(0,2*np.pi,N,endpoint=False)
 16rng=np.random.default_rng(SEED)
 17X=np.c_[np.cos(ang),np.sin(ang)].astype('float32')
 18# mildly nonuniform labels, but smooth along the circle
 19Y=(np.sin(3*ang)+.25*np.cos(ang)>0).astype('int64')
 20perm=rng.permutation(N); X=X[perm]; Y=Y[perm]; ang=ang[perm]
 21# kNN graph, undirected edge list, unit weights
 22G=kneighbors_graph(X, n_neighbors=8, mode='connectivity', include_self=False).toarray()
 23G=np.maximum(G,G.T); edges=np.argwhere(np.triu(G,1)>0)
 24
 25def energy(v, ee=edges):
 26    return float(np.mean((v[ee[:,0]]-v[ee[:,1]])**2)) if len(ee) else 0.
 27
 28def graph_est(v, idx, p):
 29    mask=np.zeros(N,dtype=bool); mask[idx]=True
 30    keep=mask[edges[:,0]] & mask[edges[:,1]]
 31    # Horvitz-Thompson estimate for both endpoints being sampled
 32    return energy(v[...], edges[keep])*len(edges[keep])/max(len(edges),1)/(p*p) if keep.any() else 0.
 33
 34def q_score(v, idx, probe):
 35    # Formula Q(S), with graph energy estimated by sampled edges.
 36    refn=float(np.mean(v[probe]**2)); sn=float(np.mean(v[idx]**2))
 37    refb=energy(v[probe], edges[np.isin(edges,probe).all(axis=1)]) if False else energy(v)
 38    # use full graph energy as the fixed probe reference; candidate is HT graph estimate
 39    sb=graph_est(v,idx,len(idx)/N)
 40    return abs(sn-refn)/(refn+1e-8)+0.7*abs(sb-refb)/(refb+1e-8)
 41
 42# Stage-1 claimed phenomenon: low-frequency functions have smaller graph energy and
 43# geometry-aware acceptance reduces norm/energy distortion.
 44math_rows=[]
 45probe=rng.choice(N,128,replace=False)
 46for freq in [1,2,4,8,16]:
 47    v=np.sin(freq*ang)
 48    qs=[]; errs=[]; accepted=[]
 49    for t in range(250):
 50        idx=rng.choice(N,64,replace=False)
 51        qs.append(q_score(v,idx,probe)); errs.append(abs(np.mean(v[idx]**2)-np.mean(v**2)))
 52        # accept among three attempts, exactly the audit mechanism
 53        cand=[rng.choice(N,64,replace=False) for _ in range(3)]
 54        best=min(cand,key=lambda z:q_score(v,z,probe)); accepted.append(abs(np.mean(v[best]**2)-np.mean(v**2)))
 55    math_rows.append({'freq':freq,'energy':energy(v),'random_abs_norm_error':float(np.mean(errs)),
 56                      'audited_abs_norm_error':float(np.mean(accepted)),
 57                      'q_error_corr':float(np.corrcoef(qs,errs)[0,1])})
 58
 59class MLP(torch.nn.Module):
 60    def __init__(self):
 61        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(2,32),torch.nn.Tanh(),torch.nn.Linear(32,2))
 62    def forward(self,x): return self.net(x)
 63
 64def flatgrad(model):
 65    return torch.cat([p.grad.detach().flatten() for p in model.parameters() if p.grad is not None])
 66
 67def run(audited, steps=90):
 68    torch.manual_seed(SEED+int(audited)); model=MLP().to(device)
 69    xt=torch.tensor(X,device=device); yt=torch.tensor(Y,device=device)
 70    opt=torch.optim.SGD(model.parameters(),lr=.12)
 71    losses=[]; cosines=[]; biases=[]; rejects=[]; audit_time=0.; train_time=time.time()
 72    for step in range(steps):
 73        # Full-data audit values are an inexpensive stand-in for a fixed probe audit.
 74        with torch.no_grad():
 75            logits=model(xt); lv=torch.nn.functional.cross_entropy(logits,yt,reduction='none').cpu().numpy()
 76            logv=logits[:,0].cpu().numpy()
 77        # full gradient reference (small toy only)
 78        model.zero_grad(); full=torch.nn.functional.cross_entropy(model(xt),yt); full.backward(); gf=flatgrad(model).cpu()
 79        t0=time.time(); attempts=0
 80        def score(z): return q_score(lv,z,probe)+0.35*q_score(logv,z,probe)
 81        candidates=[]
 82        for j in range(3 if audited else 1):
 83            z=rng.choice(N,64,replace=False); candidates.append((score(z),z))
 84        best,z=min(candidates,key=lambda x:x[0]); attempts=len(candidates); audit_time+=time.time()-t0
 85        # baseline uses one random draw; idea uses best of three (two retries maximum)
 86        rejects.append(attempts-1)
 87        batch_loss=float(np.mean(lv[z])); biases.append(abs(batch_loss-float(np.mean(lv))))
 88        xb=xt[z]; yb=yt[z]
 89        opt.zero_grad(); lb=torch.nn.functional.cross_entropy(model(xb),yb); lb.backward(); gb=flatgrad(model).cpu()
 90        cosines.append(float(torch.dot(gf,gb)/(gf.norm()*gb.norm()+1e-12)))
 91        opt.step(); losses.append(float(lb.detach().cpu()))
 92    return {'loss':float(np.mean(losses[-20:])), 'gradient_cosine':float(np.mean(cosines)),
 93            'loss_bias':float(np.mean(biases)), 'retries':float(np.mean(rejects)),
 94            'audit_overhead':float(audit_time/max(time.time()-train_time,1e-9)),
 95            'val_acc':float((model(xt).argmax(1)==yt).float().mean().cpu())}
 96
 97baseline=run(False); idea=run(True)
 98out={'device':str(device),'math_check':math_rows,'baseline':baseline,'idea':idea,
 99     'notes':'Three candidate batches for audited mode; same 64 examples per optimizer update, attempted-example overhead is 3x in audited mode.'}
100with open('results.json','w') as f: json.dump(out,f,indent=2)
101print(json.dumps(out,indent=2))