1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6from sklearn.datasets import make_moons
  7from sklearn.model_selection import train_test_split
  8
  9SEED = 636
 10random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 11try:
 12    if torch.cuda.is_available():
 13        device = torch.device('cuda')
 14    else: device = torch.device('cpu')
 15except Exception:
 16    device = torch.device('cpu')
 17
 18
 19def cosine(a, b):
 20    den = np.linalg.norm(a)*np.linalg.norm(b)
 21    return float(np.dot(a,b)/den) if den > 1e-12 else 0.0
 22
 23def scores(updates, projection):
 24    """Updates is a list of (n_chunks, parameter_dim) arrays, one tensor per layer.
 25    Projection puts layer/chunk directions in one common low-dimensional space."""
 26    local = []
 27    directions = []
 28    for U in updates:
 29        V = np.asarray([projection @ x for x in U])
 30        V = V / np.maximum(np.linalg.norm(V, axis=1, keepdims=True), 1e-12)
 31        local.extend(cosine(V[i], V[i+1]) for i in range(len(V)-1))
 32        directions.append(np.mean(V, axis=0))
 33    u = np.asarray(directions)
 34    u = u / np.maximum(np.linalg.norm(u, axis=1, keepdims=True), 1e-12)
 35    mean = u.mean(axis=0)
 36    G = ((u-mean).T @ (u-mean))/len(u)
 37    eig = np.linalg.eigvalsh(G)
 38    tr = float(np.trace(G))
 39    C = float(eig[-1]/tr) if tr > 1e-12 else 0.0
 40    return float(np.mean(local)), C
 41
 42def controller_eta(eta, Q, C, kappa=1.0, C0=.45, q0=.3):
 43    # Formula from the proposal, with the stated high-Q gate.
 44    if Q > q0 and C < C0:
 45        return eta * math.exp(-kappa*(1-Q)*max(C0-C, 0.0))
 46    return eta
 47
 48def math_check():
 49    # Verify the controller on the smallest clear construction: each block has
 50    # internally aligned neighboring updates (Q high), while four blocks point
 51    # into distinct orthogonal global sectors (C low).
 52    rng = np.random.default_rng(4)
 53    d = 4
 54    sectors = np.eye(d)
 55    competing = [np.tile(sectors[i], (4, 1)) + .001*rng.normal(size=(4, d))
 56                 for i in range(4)]
 57    Q, C = scores(competing, np.eye(d))
 58    eta0 = 1.; eta1 = controller_eta(eta0, Q, C, kappa=2., C0=.8)
 59    # A single global sector has negligible dispersion; the covariance ratio is
 60    # conventionally set to zero when its trace vanishes.
 61    aligned = [np.tile(np.array([1., 0, 0, 0]), (4, 1))
 62               + .001*rng.normal(size=(4, d)) for _ in range(4)]
 63    Qa, Ca = scores(aligned, np.eye(d))
 64    return {'Q_competing': Q, 'C_competing': C,
 65            'eta_ratio_competing': eta1/eta0,
 66            'Q_aligned': Qa, 'C_aligned': Ca,
 67            'eta_ratio_aligned': controller_eta(1, Qa, Ca, C0=.8),
 68            'controller_triggered_on_competing': bool(eta1 < eta0)}
 69
 70class Net(nn.Module):
 71    def __init__(self):
 72        super().__init__()
 73        self.layers = nn.ModuleList([nn.Linear(2,32), nn.Linear(32,32), nn.Linear(32,2)])
 74    def forward(self,x):
 75        x=torch.tanh(self.layers[0](x)); x=torch.tanh(self.layers[1](x)); return self.layers[2](x)
 76
 77def run(controlled, steps=900):
 78    X,y=make_moons(n_samples=1800, noise=.22, random_state=SEED)
 79    X=(X-X.mean(0))/X.std(0)
 80    xt,xv,yt,yv=train_test_split(X,y,test_size=.3,random_state=SEED,stratify=y)
 81    x=torch.tensor(xt,dtype=torch.float32,device=device); y=torch.tensor(yt,dtype=torch.long,device=device)
 82    xv=torch.tensor(xv,dtype=torch.float32,device=device); yv=torch.tensor(yv,dtype=torch.long,device=device)
 83    torch.manual_seed(SEED)
 84    model=Net().to(device)
 85    opt=torch.optim.SGD(model.parameters(),lr=.12,momentum=.9)
 86    params=list(model.parameters())
 87    # A shared random sketch makes block directions comparable despite different tensor sizes.
 88    rng=np.random.default_rng(99); d=18
 89    projections=[rng.normal(0,1/math.sqrt(d),(d,p.numel()//p.shape[0] if p.ndim>1 else p.numel())) for p in params]
 90    # For each layer, combine weight and bias rows into equal-sized row chunks by flattening tensor.
 91    history=[]; eta_trace=[]; q_trace=[]; c_trace=[]; loss_trace=[]
 92    batch=96
 93    for step in range(steps):
 94        ix=torch.randint(0,len(x),(batch,),device=device)
 95        before=[p.detach().clone() for p in params]
 96        opt.zero_grad(); loss=nn.functional.cross_entropy(model(x[ix]),y[ix]); loss.backward(); opt.step()
 97        if step % 30 == 29:
 98            updates=[]
 99            for li,(p,old) in enumerate(zip(params,before)):
100                delta=(p.detach()-old).cpu().numpy().reshape(-1)
101                # four neighboring temporal-like shards within each tensor; local consistency is adjacent shards.
102                chunks=np.array_split(delta,4)
103                # project with a fresh deterministic resize matrix for each chunk to common coordinates
104                U=[]
105                for ci,ch in enumerate(chunks):
106                    rr=np.random.default_rng(1000+li*10+ci)
107                    R=rr.normal(0,1/math.sqrt(d),(d,len(ch)))
108                    U.append(R@ch)
109                updates.append(np.asarray(U))
110            Q,C=scores(updates,np.eye(d))
111            eta=opt.param_groups[0]['lr']
112            # Baseline cosine schedule is applied every diagnostic interval; controller adds coarsening slowdown.
113            base=.12*(.015 + .985*0.5*(1+math.cos(math.pi*(step+1)/steps)))
114            neweta=base
115            if controlled and Q>.3 and C<.45:
116                neweta=controller_eta(base,Q,C,kappa=2.5,C0=.45,q0=.3)
117            for g in opt.param_groups:g['lr']=neweta
118            q_trace.append(Q); c_trace.append(C); eta_trace.append(neweta)
119            with torch.no_grad():
120                vl=nn.functional.cross_entropy(model(xv),yv).item(); acc=(model(xv).argmax(1)==yv).float().mean().item()
121            history.append({'step':step+1,'loss':vl,'acc':acc,'Q':Q,'C':C,'lr':neweta})
122    return {'final_acc':history[-1]['acc'],'final_loss':history[-1]['loss'],
123            'mean_Q':float(np.mean(q_trace)),'mean_C':float(np.mean(c_trace)),
124            'highQ_lowC_fraction':float(np.mean([(q>.3 and c<.45) for q,c in zip(q_trace,c_trace)])),
125            'min_lr':float(min(eta_trace)), 'records':history}
126
127def main():
128    math_result=math_check()
129    # Separate process-level failure fallback is handled by rerunning on CPU if CUDA errors.
130    try:
131        base=run(False); idea=run(True)
132    except Exception as e:
133        global device
134        device=torch.device('cpu')
135        base=run(False); idea=run(True)
136        math_result['cuda_error']=repr(e)
137    out={'device':str(device),'math_check':math_result,'baseline':base,'idea':idea}
138    Path('results.json').write_text(json.dumps(out,indent=2))
139    print(json.dumps({'device':str(device),'math_check':math_result,
140                      'baseline':{k:v for k,v in base.items() if k!='records'},
141                      'idea':{k:v for k,v in idea.items() if k!='records'}},indent=2))
142if __name__=='__main__': main()