1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1305
  6np.random.seed(SEED); random.seed(SEED)
  7
  8# Overdamped bistable diffusion: dx=b(x)dt+sqrt(eps)*G dW,
  9# U=(x^2-1)^2/4, b=-U', G=sqrt(2), hence a=2.
 10def U(x): return 0.25*(x*x-1.0)**2
 11def drift(x): return x-x**3
 12def exact_V(x): return U(x) # V(-1)=0
 13def exact_grad_V(x): return x**3-x
 14
 15
 16def hj_check():
 17    x=np.linspace(-1.5,1.5,2001)
 18    b=drift(x); a=2.0; vp=exact_grad_V(x)
 19    r=b*vp+0.5*a*vp*vp
 20    # The nonzero numerical error is only floating point cancellation.
 21    return {"max_abs_residual":float(np.max(np.abs(r))),
 22            "barrier_predicted":float(exact_V(0)-exact_V(-1)),
 23            "gradient_relation_max_error":float(np.max(np.abs(vp+2*b/a)))}
 24
 25
 26def pinn_fit():
 27    # Fit a scalar neural quasipotential to HJ residual plus anchor and
 28    # positivity. Drift/covariance are deliberately estimated from bins.
 29    import torch
 30    torch.manual_seed(SEED)
 31    device="cuda" if torch.cuda.is_available() else "cpu"
 32    try:
 33        dtype=torch.float32
 34        centers=np.linspace(-1.45,1.45,81)
 35        # Synthetic increment samples emulate measured optimizer increments.
 36        rng=np.random.default_rng(SEED); dt=.01; eps=.08; n=50000
 37        x=rng.uniform(-1.45,1.45,n)
 38        dx=drift(x)*dt+np.sqrt(eps*2*dt)*rng.normal(size=n)
 39        bid=np.digitize(x,centers)-1
 40        bh=[]; ah=[]; cc=[]
 41        for i,c in enumerate(centers):
 42            q=bid==i
 43            if q.sum()>50:
 44                bh.append(dx[q].mean()/dt); ah.append(dx[q].var()/dt/eps); cc.append(c)
 45        # Estimated a is 2, as it should be; interpolation in training grid.
 46        z=torch.tensor(np.array(cc)[:,None],device=device,dtype=torch.float32)
 47        bt=torch.tensor(np.array(bh),device=device,dtype=torch.float32); at=torch.tensor(np.array(ah),device=device,dtype=torch.float32)
 48        net=torch.nn.Sequential(torch.nn.Linear(1,32),torch.nn.Tanh(),torch.nn.Linear(32,32),torch.nn.Tanh(),torch.nn.Linear(32,1)).to(device)
 49        opt=torch.optim.Adam(net.parameters(),lr=.01)
 50        for _ in range(1800):
 51            zz=z.clone().requires_grad_(True); v=net(zz).squeeze(1)
 52            g=torch.autograd.grad(v.sum(),zz,create_graph=True)[0].squeeze(1)
 53            res=bt*g+.5*at*g*g
 54            pos=torch.relu(-v).square().mean()
 55            # Anchor at the left attractor; positivity is meaningful for this reference.
 56            loss=res.square().mean()+8*v[torch.argmin(torch.abs(z[:,0]+1))].square()+.2*pos
 57            opt.zero_grad(); loss.backward(); opt.step()
 58        with torch.no_grad():
 59            pred=net(torch.tensor([[0.]],device=device,dtype=torch.float32)).item()-net(torch.tensor([[-1.]],device=device,dtype=torch.float32)).item()
 60        # Evaluate an independent collocation residual using autodiff.
 61        zz=torch.tensor(np.linspace(-1.4,1.4,401)[:,None],device=device,dtype=torch.float32,requires_grad=True)
 62        vv=net(zz); gg=torch.autograd.grad(vv.sum(),zz)[0].squeeze(1)
 63        rr=drift(zz.detach().cpu().numpy().ravel())*gg.detach().cpu().numpy()+gg.detach().cpu().numpy()**2
 64        return {"device":device,"estimated_covariance_mean":float(np.mean(ah)),
 65                "pinn_barrier":float(pred),"pinn_independent_rms_HJ":float(np.sqrt(np.mean(rr**2))),
 66                "exact_barrier":.25}
 67    except Exception as e:
 68        return {"device":"cpu-fallback","error":str(e)}
 69
 70
 71def transition_sweep():
 72    # Vectorized Euler simulation, with reset after every left-to-right crossing.
 73    # Rate = crossings / time spent in left basin; this gives an empirical
 74    # renewal escape rate and tests log(rate)~ -DeltaV/eps.
 75    rng=np.random.default_rng(SEED+1); dt=.005; ntraj=500; steps=30000
 76    epses=np.array([.06,.075,.09,.11,.14,.18])
 77    rates=[]; counts=[]; exposures=[]
 78    for eps in epses:
 79        x=np.full(ntraj,-1.0); count=0; exposure=0.0
 80        for _ in range(steps):
 81            left=x<0
 82            exposure += left.sum()*dt
 83            noise=np.sqrt(eps*2*dt)*rng.normal(size=ntraj)
 84            x += drift(x)*dt+noise
 85            hit=left & (x>=0)
 86            count += int(hit.sum())
 87            x[hit]=-1.0
 88            # keep all trajectories in the renewal experiment in the left basin
 89            x[x>0]=-1.0
 90        rates.append(count/max(exposure,1e-12)); counts.append(count); exposures.append(exposure)
 91    rates=np.array(rates); inv=1/epses
 92    slope,intercept=np.polyfit(inv,np.log(rates),1)
 93    # fit only rates with observed events (all should have events here)
 94    # Second parameter sweep: scale a=2*c. The same HJ equation predicts
 95    # DeltaV(c)=0.25/c, so the fitted slope versus 1/eps is -0.25/c.
 96    scale_rows=[]
 97    for c in np.array([0.5, 1.0, 2.0]):
 98        ee=np.array([0.08,0.11,0.15,0.20,0.28])
 99        xx=np.full(300,-1.0); ct=0; ex=0.0
100        for _ in range(20000):
101            left=xx<0; ex += left.sum()*dt
102            xx += drift(xx)*dt+np.sqrt(epses[0]*0 + ee[0]*2*c*dt)*rng.normal(size=300)
103            # Use a per-scale fixed epsilon list over independent blocks below.
104            hit=left & (xx>=0); ct += int(hit.sum()); xx[hit]=-1.; xx[xx>0]=-1.
105        # Re-simulate each epsilon; this keeps the fit transparent and seeded.
106        rr=[]
107        for e in ee:
108            xx=np.full(300,-1.0); hits=0; exposure=0.0
109            for _ in range(16000):
110                left=xx<0; exposure += left.sum()*dt
111                xx += drift(xx)*dt+np.sqrt(e*2*c*dt)*rng.normal(size=300)
112                hit=left & (xx>=0); hits += int(hit.sum()); xx[hit]=-1.; xx[xx>0]=-1.
113            rr.append(hits/max(exposure,1e-12))
114        sl=np.polyfit(1/ee,np.log(np.maximum(rr,1e-15)),1)[0]
115        scale_rows.append({"covariance_scale":float(c),"fitted_slope":float(sl),
116                           "predicted_slope":float(-.25/c),
117                           "relative_error":float(abs(sl+.25/c)/(.25/c)),"rates":rr})
118    return {"eps":epses.tolist(),"rates":rates.tolist(),"counts":counts,
119            "slope_log_rate_vs_inverse_eps":float(slope),"predicted_slope":-.25,
120            "relative_slope_error":float(abs(slope+.25)/.25),"intercept":float(intercept),
121            "covariance_scale_sweep":scale_rows}
122
123
124def main():
125    out={"seed":SEED,"math_check":hj_check(),"pinn":pinn_fit(),"transitions":transition_sweep()}
126    Path("results.json").write_text(json.dumps(out,indent=2))
127    print(json.dumps(out,indent=2))
128if __name__=="__main__": main()