1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED=2433
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device='cuda' if torch.cuda.is_available() else 'cpu'
 11try:
 12    if device=='cuda': torch.cuda.empty_cache()
 13except Exception:
 14    device='cpu'
 15
 16a=0.7; sigma=0.35; beta=0.8; dt=0.04
 17
 18def exact_V(x): return x*x
 19def drift(x): return -a*x
 20def exact_r(x): return (beta+2*a-sigma*sigma)*x*x-sigma*sigma
 21
 22def generator_torch(v,x):
 23    g=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
 24    h=torch.autograd.grad(g.sum(),x,create_graph=True)[0]
 25    return drift(x)*g+0.5*sigma*sigma*h
 26
 27def analytic_generator(x):
 28    # For V=x^2: f V' + .5 sigma^2 V'' = -2a x^2 + sigma^2.
 29    return -2*a*x*x+sigma*sigma
 30
 31def math_checks():
 32    xs=np.linspace(-.9,.9,101)
 33    x=torch.tensor(xs[:,None],dtype=torch.float64,requires_grad=True)
 34    ad=float(np.max(np.abs(generator_torch(x*x,x).detach().numpy().ravel()-analytic_generator(xs))))
 35    # Antithetic Monte Carlo makes the leading O(dt) Euler bias measurable.
 36    rng=np.random.default_rng(SEED); x0=.63; hs=[.04,.02,.01,.005]; errs=[]
 37    for h in hs:
 38        z=rng.normal(size=1000000); z=np.concatenate([z,-z])
 39        xn=x0+drift(x0)*h+sigma*np.sqrt(h)*z
 40        errs.append(float(abs(np.mean((xn*xn-x0*x0)/h)-analytic_generator(np.array([x0]))[0])))
 41    return {'autodiff_generator_max_error':ad,'weak_dt':hs,'weak_generator_abs_error':errs,
 42            'weak_error_ratio_last_first':errs[-1]/errs[0],
 43            'predicted_euler_bias_ratio':hs[-1]/hs[0]}
 44
 45class Net(nn.Module):
 46    def __init__(self):
 47        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 48    def forward(self,x): return self.net(x)
 49
 50def make_data(n=512):
 51    rng=np.random.default_rng(SEED+4)
 52    x=rng.uniform(-1,1,(n,1)); z=rng.normal(size=(n,1))
 53    xp=x+drift(x)*dt+sigma*np.sqrt(dt)*z
 54    r=exact_r(x)*dt
 55    return tuple(torch.tensor(q,dtype=torch.float32,device=device) for q in (x,xp,r))
 56
 57def train(mode, lam_final=10., steps=500, threshold=.0007, patience=5, seed=SEED):
 58    torch.manual_seed(seed); np.random.seed(seed)
 59    net=Net().to(device); opt=torch.optim.Adam(net.parameters(),lr=3e-3)
 60    x,xp,r=make_data(); coll=torch.linspace(-1,1,96,device=device).reshape(-1,1)
 61    lam=0.; ramp_steps=[]; ema=None; good=0; history=[]
 62    for k in range(steps):
 63        opt.zero_grad(); vx=net(x); td=(r+math.exp(-beta*dt)*net(xp)-vx).pow(2).mean()
 64        c=coll.detach().clone().requires_grad_(True); vc=net(c)
 65        H=generator_torch(vc,c)-beta*vc+exact_r(c); pde=H.pow(2).mean()
 66        b=net(torch.zeros(16,1,device=device)).pow(2).mean()
 67        with torch.no_grad():
 68            t=float(td.detach()); var=float(vx.var().detach())
 69            ema=t if ema is None else .95*ema+.05*t
 70        if mode=='td': use=0.
 71        elif mode=='fixed': use=lam_final
 72        else:
 73            if ema < threshold and var < .30: good+=1
 74            else: good=0
 75            if good>=patience and lam < lam_final:
 76                lam=.1 if lam==0 else min(lam_final,2*lam); ramp_steps.append(k+1); good=0
 77            use=lam
 78        loss=td+use*pde+2*b; loss.backward(); opt.step()
 79        if k in (0,4,9,24,49,99,199,299,399,499):
 80            with torch.no_grad():
 81                grid=torch.linspace(-1,1,401,device=device).reshape(-1,1)
 82                rmse=(net(grid)-grid.pow(2)).pow(2).mean().sqrt().item()
 83            history.append([k+1,float(td.detach()),float(pde.detach()),rmse,use])
 84    return {'final_td':history[-1][1],'final_pde':history[-1][2],'value_rmse':history[-1][3], 'ramp_steps':ramp_steps,'history':history,'final_lambda':lam}
 85
 86def gate_sweep():
 87    # Prediction: stricter TD threshold cannot ramp earlier; threshold zero never ramps.
 88    out=[]
 89    for th in [0.,.0003,.0007,.002,.02]:
 90        z=train('continuation',steps=180,threshold=th,patience=5)
 91        out.append({'threshold':th,'first_ramp_step':z['ramp_steps'][0] if z['ramp_steps'] else None,'num_ramps':len(z['ramp_steps'])})
 92    return out
 93
 94def lambda_sweep():
 95    # Prediction: lambda=0 is exactly TD-only; fixed strong PDE perturbs TD upward.
 96    out=[]
 97    for mode in ['td','fixed','continuation']:
 98        z=train(mode,steps=300,threshold=(.002 if mode=='continuation' else .0007))
 99        out.append({'mode':mode,'td':z['final_td'],'pde':z['final_pde'],'rmse':z['value_rmse']})
100    return out
101
102def main():
103    out={'device':device,'math':math_checks(),'gate_sweep':gate_sweep(),'matched_training':lambda_sweep()}
104    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
105if __name__=='__main__': main()