1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2871
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def sigmoid(x):
 10    x = np.clip(x, -60.0, 60.0)
 11    return 1.0 / (1.0 + np.exp(-x))
 12
 13
 14def toy_checks():
 15    # Prediction 1: q(E)=1/2 exactly at E=E_star and transition is about 4*tau.
 16    E_star, tau = 2.0, 0.25
 17    qstar = float(sigmoid((E_star-E_star)/tau))
 18    e10 = E_star - tau*math.log(9.0)
 19    e90 = E_star + tau*math.log(9.0)
 20    transition_width = e90-e10
 21
 22    # Prediction 2: in the active, unforced regime dE/dt=-2*c*E,
 23    # hence log(E) slope=-2c, independently of initial energy.
 24    decay = []
 25    for c in [0.2, 0.5, 1.0, 1.5]:
 26        dt=1e-3; n=10000; v=math.sqrt(2*8.0); logs=[]; ts=[]
 27        for k in range(n):
 28            E=.5*v*v
 29            if k % 10 == 0:
 30                logs.append(math.log(E)); ts.append(k*dt)
 31            q=sigmoid((E-E_star)/tau)
 32            v += -dt*q*c*v
 33        # fit only well-active portion, E > 3 E_star
 34        mask=np.array([math.exp(x)>3*E_star for x in logs])
 35        slope=np.polyfit(np.array(ts)[mask],np.array(logs)[mask],1)[0]
 36        decay.append({'c':c,'predicted_log_slope':-2*c,'observed_log_slope':float(slope),
 37                      'relative_error':float(abs(slope+2*c)/(2*c))})
 38
 39    # Prediction 3: bounded adversarial forcing g=-G sign(v), active equilibrium
 40    # has E*=G^2/(2c^2), from G*sqrt(2E)=2cE.
 41    forced=[]; G=1.0
 42    for c in [0.25,0.5,1.0,2.0]:
 43        dt=2e-4; n=250000; v=0.05
 44        vals=[]
 45        # Choose a low threshold so every predicted equilibrium is well inside q≈1.
 46        forced_E_star, forced_tau = 0.01, 0.001
 47        for k in range(n):
 48            E=.5*v*v; q=sigmoid((E-forced_E_star)/forced_tau)
 49            g=-G*(1.0 if v >= 0 else -1.0)
 50            v += dt*(-g-q*c*v)
 51            if k >= n//2: vals.append(.5*v*v)
 52        observed=float(np.mean(vals[-50000:]))
 53        predicted=G*G/(2*c*c)
 54        forced.append({'c':c,'predicted_plateau_E':predicted,
 55                       'observed_plateau_E':observed,
 56                       'relative_error':abs(observed-predicted)/predicted})
 57    return {'gate_transition':{'E_star':E_star,'tau':tau,'predicted_q_at_Estar':.5,
 58                               'observed_q_at_Estar':qstar,
 59                               'predicted_width_10_to_90':4.3944491547*tau,
 60                               'observed_width_10_to_90':transition_width},
 61            'unforced_decay':decay,'forced_plateau':forced}
 62
 63
 64def train_benchmark():
 65    # sklearn digits is a tiny, reproducible MNIST-like classification task.
 66    import torch
 67    import torch.nn as nn
 68    from sklearn.datasets import load_digits
 69    from sklearn.model_selection import train_test_split
 70    torch.manual_seed(SEED); np.random.seed(SEED)
 71    try:
 72        device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 73        # Trigger a tiny allocation so CUDA initialization errors are caught.
 74        if device.type=='cuda': torch.empty(1,device=device)
 75    except Exception:
 76        device=torch.device('cpu')
 77    x,y=load_digits(return_X_y=True)
 78    x=x.astype('float32')/16.0
 79    xa,xb,ya,yb=train_test_split(x,y,test_size=.25,random_state=SEED,stratify=y)
 80    tx=torch.tensor(xa); ty=torch.tensor(ya,dtype=torch.long)
 81    vx=torch.tensor(xb); vy=torch.tensor(yb,dtype=torch.long)
 82    def make():
 83        return nn.Sequential(nn.Linear(64,96),nn.ReLU(),nn.Linear(96,64),nn.ReLU(),nn.Linear(64,10)).to(device)
 84    def run(mode, lr):
 85        torch.manual_seed(SEED)
 86        model=make(); lossfn=nn.CrossEntropyLoss()
 87        # Explicit velocity implements theta <- theta + v, matching the stated interface.
 88        vel=[torch.zeros_like(p,device=device) for p in model.parameters()]
 89        bs=128; losses=[]; energies=[]; gates=[]
 90        t0=time.time()
 91        for step in range(300):
 92            idx=((torch.arange(bs)+step*bs) % len(tx))
 93            out=model(tx[idx].to(device)); loss=lossfn(out,ty[idx].to(device))
 94            gs=torch.autograd.grad(loss,tuple(model.parameters()))
 95            E=.5*sum(float((v*v).sum().detach().cpu()) for v in vel)
 96            q=float(sigmoid((E-0.02)/0.01)) if mode=='gated' else 0.0
 97            clip_scale=1.0
 98            if mode=='clip':
 99                total_norm=torch.sqrt(sum((g*g).sum() for g in gs))
100                clip_scale=min(1.0, 1.0/(float(total_norm.detach().cpu())+1e-12))
101            with torch.no_grad():
102                for j,(p,g) in enumerate(zip(model.parameters(),gs)):
103                    gg=g*clip_scale
104                    vel[j].mul_(0.9*(1.0-lr*1.0*q)).add_(gg,alpha=-lr)
105                    p.add_(vel[j])
106            losses.append(float(loss.detach().cpu())); energies.append(E); gates.append(q)
107        with torch.no_grad():
108            acc=float((model(vx.to(device)).argmax(1).cpu()==vy).float().mean())
109        return {'final_train_loss':losses[-1],'eval_accuracy':acc,'max_energy':max(energies),
110                'mean_last50_energy':float(np.mean(energies[-50:])),'mean_gate_last50':float(np.mean(gates[-50:])),
111                'seconds':time.time()-t0,'device':str(device)}
112    results={}
113    for lr in [0.03,0.1,0.3]:
114        for mode in ['momentum','clip','gated']:
115            try: results[f'{mode}_lr{lr}']=run(mode,lr)
116            except Exception as e:
117                results[f'{mode}_lr{lr}']={'error':repr(e)}
118    return results
119
120if __name__=='__main__':
121    out={'seed':SEED,'toy':toy_checks()}
122    try: out['benchmark']=train_benchmark()
123    except Exception as e: out['benchmark_error']=repr(e)
124    Path('results.json').write_text(json.dumps(out,indent=2))
125    print(json.dumps(out,indent=2))