1import json, math, os, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6from sklearn.datasets import load_digits
  7from sklearn.model_selection import train_test_split
  8from sklearn.preprocessing import StandardScaler
  9
 10SEED = 645
 11OUT = Path('results.json')
 12
 13def seed_all(seed=SEED):
 14    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 15    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 16
 17
 18def harmonic_check():
 19    # H(q,p)=1/2(q^2+p^2), exact leapfrog versus explicit Euler.
 20    h = 0.5; n = 2000
 21    q0, p0 = 1.0, 0.0
 22    ql, pl, qe, pe = q0, p0, q0, p0
 23    hl, he = [], []
 24    for _ in range(n):
 25        hl.append(0.5*(ql*ql+pl*pl)); he.append(0.5*(qe*qe+pe*pe))
 26        pl -= 0.5*h*ql
 27        ql += h*pl
 28        pl -= 0.5*h*ql
 29        # explicit Euler for dq/dt=p, dp/dt=-q
 30        qe, pe = qe + h*pe, pe - h*qe
 31    hl, he = np.asarray(hl), np.asarray(he)
 32    return {
 33        'step': h, 'steps': n,
 34        'leapfrog_energy_start': float(hl[0]),
 35        'leapfrog_energy_end': float(hl[-1]),
 36        'leapfrog_energy_abs_drift': float(abs(hl[-1]-hl[0])),
 37        'leapfrog_energy_range': float(hl.max()-hl.min()),
 38        'euler_energy_start': float(he[0]),
 39        'euler_energy_end': float(he[-1]),
 40        'euler_energy_growth_factor': float(he[-1]/he[0]),
 41        'euler_energy_range': float(he.max()-he.min()),
 42        'claim_observed': bool((hl.max()-hl.min()) < 0.01 and he[-1] > 100*he[0])
 43    }
 44
 45class MLP(nn.Module):
 46    def __init__(self):
 47        super().__init__()
 48        self.net = nn.Sequential(nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 10))
 49    def forward(self, x): return self.net(x)
 50
 51def make_data():
 52    x, y = load_digits(return_X_y=True)
 53    x = StandardScaler().fit_transform(x).astype('float32')
 54    xtr, xva, ytr, yva = train_test_split(x, y, test_size=.25, random_state=SEED, stratify=y)
 55    return (torch.tensor(xtr), torch.tensor(ytr, dtype=torch.long),
 56            torch.tensor(xva), torch.tensor(yva, dtype=torch.long))
 57
 58def evaluate(model, xv, yv):
 59    model.eval()
 60    with torch.no_grad():
 61        z = model(xv); loss = nn.functional.cross_entropy(z, yv).item()
 62        acc = (z.argmax(1) == yv).float().mean().item()
 63    model.train(); return loss, acc
 64
 65def batches(x, y, bs=128, seed=0):
 66    g = torch.Generator().manual_seed(seed)
 67    ix = torch.randperm(len(x), generator=g)
 68    for s in range(0, len(x), bs):
 69        j = ix[s:s+bs]; yield x[j], y[j]
 70
 71def flat_grads(model):
 72    return torch.cat([p.grad.detach().reshape(-1) for p in model.parameters()])
 73
 74def flat_params(model):
 75    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])
 76
 77def train(kind, xtr, ytr, xva, yva, grad_budget=240):
 78    # CPU is deliberately used for robust shared-machine reproducibility.
 79    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 80    model = MLP().to(device); xtr=xtr.to(device); ytr=ytr.to(device); xva=xva.to(device); yva=yva.to(device)
 81    if kind == 'sgd': opt = torch.optim.SGD(model.parameters(), lr=0.08, momentum=0.9)
 82    elif kind == 'adamw': opt = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.0)
 83    params = list(model.parameters())
 84    velocity = [torch.zeros_like(p) for p in params]
 85    mass = 1.0; h = 0.08
 86    losses=[]; grad_norms=[]; evals=0; step=0
 87    t0=time.time()
 88    while evals < grad_budget:
 89        # Fixed deterministic cycling through reshuffled minibatches.
 90        for xb, yb in batches(xtr, ytr, 128, SEED+step//100):
 91            if evals >= grad_budget: break
 92            if kind == 'leapfrog':
 93                model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb), yb); loss.backward()
 94                g1=flat_grads(model); grad_norms.append(float(g1.norm().cpu()))
 95                with torch.no_grad():
 96                    for p, v in zip(params, velocity): p.add_(h*v/mass); v.sub_(0.5*h*p.grad)
 97                # Correct kick-drift ordering requires first kick before drift; undo/restate compactly:
 98                # position was advanced using old v, so apply kick ordering from saved old state below is avoided
 99                # by using the equivalent implementation in the next branch's explicit state correction.
100                # This branch is replaced immediately by a clean two-kick update using old parameters.
101                raise RuntimeError('internal leapfrog branch should not execute')
102            else:
103                opt.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb), yb); loss.backward()
104                grad_norms.append(float(flat_grads(model).norm().detach().cpu())); opt.step(); evals += 1; step += 1
105            if step % 40 == 0: losses.append((step, evals, float(loss.detach().cpu())))
106    return {}
107
108def train_leapfrog(xtr,ytr,xva,yva,grad_budget=240):
109    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
110    model=MLP().to(device); xtr=xtr.to(device); ytr=ytr.to(device); xva=xva.to(device); yva=yva.to(device)
111    params=list(model.parameters()); vel=[torch.zeros_like(p) for p in params]; h=.08; evals=0; step=0; gn=[]; trace=[]; t0=time.time()
112    while evals+1 <= grad_budget:
113        for xb,yb in batches(xtr,ytr,128,SEED+step//100):
114            if evals+1>grad_budget: break
115            model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward(); gn.append(float(flat_grads(model).norm().cpu())); evals+=1
116            with torch.no_grad():
117                for p,v in zip(params,vel): v.sub_(.5*h*p.grad); p.add_(h*v)
118            model.zero_grad(set_to_none=True); loss2=nn.functional.cross_entropy(model(xb),yb); loss2.backward(); gn.append(float(flat_grads(model).norm().cpu())); evals+=1
119            with torch.no_grad():
120                for p,v in zip(params,vel): v.sub_(.5*h*p.grad)
121            step+=1
122            if step%20==0: trace.append((step,evals,float(loss2.detach().cpu())))
123            if evals>=grad_budget: break
124    vl,va=evaluate(model,xva,yva)
125    return {'steps':step,'gradient_evals':evals,'final_train_loss':float(loss2.detach().cpu()),'val_loss':vl,'val_accuracy':va,'mean_grad_norm':float(np.mean(gn)),'time_sec':time.time()-t0,'trace':trace}
126
127def train_standard(kind,xtr,ytr,xva,yva,grad_budget=240):
128    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); model=MLP().to(device)
129    xtr=xtr.to(device);ytr=ytr.to(device);xva=xva.to(device);yva=yva.to(device); params=list(model.parameters())
130    opt=torch.optim.SGD(params,lr=.08,momentum=.9) if kind=='sgd' else torch.optim.AdamW(params,lr=.003,weight_decay=0.)
131    evals=step=0;gn=[];trace=[];t0=time.time()
132    while evals<grad_budget:
133        for xb,yb in batches(xtr,ytr,128,SEED+step//100):
134            if evals>=grad_budget: break
135            opt.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward(); gn.append(float(flat_grads(model).norm().cpu())); opt.step(); evals+=1;step+=1
136            if step%40==0:trace.append((step,evals,float(loss.detach().cpu())))
137    vl,va=evaluate(model,xva,yva)
138    return {'steps':step,'gradient_evals':evals,'final_train_loss':float(loss.detach().cpu()),'val_loss':vl,'val_accuracy':va,'mean_grad_norm':float(np.mean(gn)),'time_sec':time.time()-t0,'trace':trace}
139
140def main():
141    seed_all(); toy=harmonic_check(); xtr,ytr,xva,yva=make_data()
142    results={'harmonic_check':toy,'device':str(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))}
143    for k in ['sgd','adamw','leapfrog']:
144        seed_all(); results[k]=train_leapfrog(xtr,ytr,xva,yva) if k=='leapfrog' else train_standard(k,xtr,ytr,xva,yva)
145    OUT.write_text(json.dumps(results,indent=2)); print(json.dumps(results,indent=2))
146if __name__=='__main__': main()