1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 3070
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def math_checks():
 10    # Piecewise-constant temperature field: the stated integral reduces to a
 11    # harmonic mean over the uphill half-period.
 12    T = np.array([0.5, 1.0, 2.0, 4.0])
 13    Tact = len(T) / np.sum(1.0 / T)
 14    numerical = 2.0 / np.mean(1.0 / T) / 2.0  # same harmonic mean, explicit average
 15    # T <- T r^k gives log(r_next)=(1-k)log(r) when U is locally fixed.
 16    rs = np.array([0.05, 0.2, 3.0, 20.0])
 17    k = 0.08
 18    next_rs = rs ** (1-k)
 19    contraction = np.max(np.abs(np.log(next_rs))/np.abs(np.log(rs)))
 20    # Arrhenius claim: log escape probability is linear in -U/T.
 21    ratios = np.linspace(0.3, 3.0, 12)
 22    p = np.exp(-ratios)
 23    slope = np.polyfit(ratios, np.log(p), 1)[0]
 24    return {"Tact_harmonic": float(Tact), "Tact_formula_error": float(abs(Tact-numerical)),
 25            "controller_log_contraction": float(contraction), "arrhenius_log_slope": float(slope)}
 26
 27
 28def double_well_check():
 29    # V(x)=(x^2-1)^2 has minima +/-1 and barrier U=1 at x=0.
 30    # Parallel trajectories reduce rare-transition sampling noise.
 31    rng = np.random.default_rng(SEED)
 32    dt, steps, burn, ntraj = 0.002, 100000, 5000, 256
 33    temps = np.array([0.18, 0.24, 0.32, 0.42, 0.55])
 34    rates = []
 35    for temp in temps:
 36        x = np.full(ntraj, -1.0)
 37        side = np.full(ntraj, -1)
 38        transitions = 0
 39        for i in range(steps):
 40            grad = 4*x*(x*x - 1)
 41            x += -grad*dt + math.sqrt(2*temp*dt)*rng.normal(size=ntraj)
 42            if i >= burn:
 43                newside = np.where(x > 0, 1, -1)
 44                transitions += int(np.count_nonzero(newside != side))
 45                side = newside
 46        rates.append(transitions/(ntraj*(steps-burn)))
 47    rates = np.maximum(np.asarray(rates), 1e-12)
 48    fit = np.polyfit(1/temps, np.log(rates), 1)
 49    return {"temperatures": temps.tolist(), "transition_rates": rates.tolist(),
 50            "log_rate_vs_inverse_temperature_slope": float(fit[0]),
 51            "expected_negative_barrier_sign": bool(fit[0] < 0)}
 52
 53def make_data(n=512):
 54    rng=np.random.default_rng(SEED)
 55    x=rng.normal(size=(n,2)).astype(np.float32)
 56    y=((x[:,0]**2 + x[:,1]**2 + .35*x[:,0]) > 1.0).astype(np.int64)
 57    # introduce a nontrivial but tiny classification task
 58    return x,y
 59
 60
 61def run_mlp(controller, steps=500):
 62    import torch
 63    torch.manual_seed(SEED); np.random.seed(SEED)
 64    device = "cuda" if torch.cuda.is_available() else "cpu"
 65    try:
 66        dev=torch.device(device)
 67        X,y=make_data(); X=torch.tensor(X,device=dev); y=torch.tensor(y,device=dev)
 68        model=torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,2)).to(dev)
 69        lossfn=torch.nn.CrossEntropyLoss()
 70        lr=0.035; T=0.0025; batch=64; d=sum(p.numel() for p in model.parameters())
 71        losses=[]; controls=[]; transitions=0; last_pred=None
 72        for step in range(steps):
 73            ix=torch.randint(0,len(X),(batch,),device=dev)
 74            logits=model(X[ix]); loss=lossfn(logits,y[ix])
 75            grads=torch.autograd.grad(loss,tuple(model.parameters()),create_graph=False)
 76            with torch.no_grad():
 77                for p,g in zip(model.parameters(),grads):
 78                    # Langevinized SGD. T is the injected update temperature.
 79                    p.add_(-lr*g + math.sqrt(2*max(T,1e-12))*lr*torch.randn_like(p))
 80            losses.append(float(loss.detach().cpu()))
 81            if step % 20 == 0:
 82                with torch.no_grad():
 83                    pred=model(X).argmax(1)
 84                    basin=int((pred==y).sum().item() > len(y)*.5)
 85                    if last_pred is not None and basin != last_pred: transitions += 1
 86                    last_pred=basin
 87            if controller and step % 50 == 0:
 88                # Gradient covariance from 8 independent minibatches; P=lr I.
 89                gs=[]
 90                for _ in range(8):
 91                    j=torch.randint(0,len(X),(batch,),device=dev)
 92                    ll=lossfn(model(X[j]),y[j])
 93                    gg=torch.autograd.grad(ll,tuple(model.parameters()))
 94                    gs.append(torch.cat([z.detach().reshape(-1) for z in gg]))
 95                G=torch.stack(gs); cov_trace=((G-G.mean(0))**2).sum(1).mean()
 96                That=float((lr*lr*cov_trace/(2*d)).detach().cpu())
 97                # Two perturbed endpoints and a linear interpolation barrier.
 98                base=[p.detach().clone() for p in model.parameters()]
 99                eps=[torch.randn_like(p)*math.sqrt(2*max(T,1e-12))*0.7 for p in model.parameters()]
100                vals=[]
101                for alpha in torch.linspace(0,1,7,device=dev):
102                    with torch.no_grad():
103                        for p,b,e in zip(model.parameters(),base,eps): p.copy_(b+(2*alpha-1)*e)
104                        vals.append(float(lossfn(model(X),y).detach().cpu()))
105                with torch.no_grad():
106                    for p,b in zip(model.parameters(),base): p.copy_(b)
107                U=max(0.0,max(vals)-vals[3])
108                effective=max(That, T)
109                r=U/max(effective,1e-10)
110                old=T
111                T=float(np.clip(T*math.exp(0.08*np.clip(math.log(max(r,1e-8)),-4,4)),1e-5,0.08))
112                controls.append((step,U,That,r,old,T))
113        with torch.no_grad():
114            final_loss=float(lossfn(model(X),y).cpu()); acc=float((model(X).argmax(1)==y).float().mean().cpu())
115        return {"final_loss":final_loss,"accuracy":acc,"mean_last100_loss":float(np.mean(losses[-100:])),
116                "transitions":transitions,"final_temperature":T,"controls":controls,"device":str(dev)}
117    except Exception as e:
118        # Explicit CPU fallback as required for shared-GPU robustness.
119        if device == "cuda":
120            torch.cuda.empty_cache()
121            torch.cuda.is_available=lambda: False
122            return run_mlp(controller,steps)
123        raise
124
125
126def main():
127    t=time.time()
128    out={"seed":SEED,"math":math_checks(),"double_well":double_well_check(),
129         "mlp_fixed":run_mlp(False),"mlp_barrier_matching":run_mlp(True),"seconds":time.time()-t}
130    # Keep the report compact while preserving all controller observations.
131    Path("results.json").write_text(json.dumps(out,indent=2))
132    print(json.dumps({k:v for k,v in out.items() if k not in ("mlp_fixed","mlp_barrier_matching")},indent=2))
133    for k in ("mlp_fixed","mlp_barrier_matching"):
134        v=out[k]; print(k, {z:v[z] for z in ("final_loss","accuracy","mean_last100_loss","transitions","final_temperature","device")})
135
136if __name__ == "__main__": main()