1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=2711
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9device='cuda' if torch.cuda.is_available() else 'cpu'
 10try:
 11    if device=='cuda': torch.cuda.empty_cache()
 12except Exception:
 13    device='cpu'
 14
 15def true_x(u):
 16    return 0.2 + 0.8*u + 0.08*torch.sin(4*math.pi*u)
 17
 18def cost(x,u):
 19    return (x-0.72)**2 + 0.035*(u-0.15)**2
 20
 21def grid_opt(xfun, n=20001):
 22    u=torch.linspace(0,1,n,device=device)
 23    with torch.no_grad():
 24        j=cost(xfun(u),u); k=int(torch.argmin(j))
 25    return float(u[k]), float(j[k])
 26
 27def stability_sweep():
 28    H=3.7; rows=[]
 29    for rho in np.linspace(.05,.7,14):
 30        e=1.0
 31        for _ in range(80): e=(1-rho*H)*e
 32        rows.append({'rho':float(rho),'rhoH':float(rho*H),'error80':float(abs(e)),
 33                     'stable_observed':bool(abs(e)<1e-4),
 34                     'stable_predicted':bool(abs(1-rho*H)<1)})
 35    return {'H':H,'predicted_rho_boundary_2_over_H':2/H,'rows':rows}
 36
 37def minima_count(eps,k=6*math.pi,n=20001):
 38    u=np.linspace(0,1,n); us=.63
 39    j=(u-us)**2+eps*np.cos(k*u)
 40    idx=np.where((j[1:-1]<j[:-2])&(j[1:-1]<j[2:]))[0]+1
 41    return int(len(idx)),u[idx].tolist(),float(u[np.argmin(j)])
 42
 43def minima_sweep():
 44    k=6*math.pi; predicted=2/(k*k)
 45    epses=np.array([0,.001,.003,.005,.01,.02,.04,.08])
 46    return {'k':k,'predicted_curvature_threshold':predicted,
 47            'rows':[{'eps':float(e),'count':minima_count(e,k)[0],
 48                     'global_u':minima_count(e,k)[2]} for e in epses]}
 49
 50class Surrogate(nn.Module):
 51    def __init__(self):
 52        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 53    def forward(self,u): return self.net(u)
 54
 55def decision_unroll(model, starts, steps=35, lr=.08):
 56    u=starts.clone().requires_grad_(True)
 57    for _ in range(steps):
 58        j=cost(model(u[:,None]).squeeze(1),u)
 59        g=torch.autograd.grad(j.sum(),u,create_graph=True)[0]
 60        u=torch.clamp(u-lr*g,0,1)
 61    j=cost(model(u[:,None]).squeeze(1),u)
 62    return u,j
 63
 64def train(decision=False, seed=2711, epochs=1800):
 65    torch.manual_seed(seed)
 66    model=Surrogate().to(device)
 67    # Sparse observations intentionally leave decision-relevant interpolation freedom.
 68    ud=torch.tensor([[0.0],[.22],[.48],[.76],[1.0]],dtype=torch.float32,device=device)
 69    yd=true_x(ud)
 70    ut,jt=grid_opt(true_x)
 71    ustar=torch.tensor(ut,device=device); jstar=torch.tensor(jt,device=device)
 72    opt=torch.optim.Adam(model.parameters(),lr=.008)
 73    starts=torch.tensor([.05,.25,.5,.75,.95],device=device)
 74    for ep in range(epochs):
 75        pred=model(ud).squeeze(1); loss=((pred-yd)**2).mean()
 76        if decision:
 77            us,js=decision_unroll(model,starts,steps=18,lr=.10)
 78            # soft selection: lowest-start solution is used as differentiable proxy
 79            q=torch.argmin(js.detach()); ld=(us[q]-ustar)**2 + 2.0*(js[q]-jstar)**2
 80            loss=loss+3.0*ld
 81        opt.zero_grad(); loss.backward(); opt.step()
 82    return model,ut,jt
 83
 84def assess(model, ut, jt):
 85    us=np.linspace(0,1,101); u=torch.tensor(us,dtype=torch.float32,device=device)
 86    with torch.no_grad(): j=cost(model(u[:,None]).squeeze(1),u).cpu().numpy()
 87    ids=np.where((j[1:-1]<j[:-2])&(j[1:-1]<j[2:]))[0]+1
 88    # multistart projected gradient, no graph needed
 89    sols=[]
 90    for s in np.linspace(.01,.99,100):
 91        z=torch.tensor([s],dtype=torch.float32,device=device,requires_grad=True)
 92        for _ in range(120):
 93            jj=cost(model(z),z); g=torch.autograd.grad(jj,z)[0]
 94            with torch.no_grad(): z.clamp_(0,1); z.sub_(.08*g).clamp_(0,1)
 95            z.requires_grad_(True)
 96        sols.append(float(z.detach()))
 97    distinct=[]
 98    for z in sorted(sols):
 99        if not distinct or abs(z-distinct[-1])>.025: distinct.append(z)
100    k=int(np.argmin(j));
101    return {'grid_opt_u':float(us[k]),'grid_opt_j':float(j[k]),'displacement':abs(float(us[k])-ut),
102            'grid_local_minima':int(len(ids)),'multistart_distinct':len(distinct),
103            'multistart_solutions':distinct,'data_mse':float(((model(torch.tensor([[0.],[.22],[.48],[.76],[1.]],device=device))-true_x(torch.tensor([[0.],[.22],[.48],[.76],[1.]],device=device)))**2).mean().detach().cpu())}
104
105def main():
106    ut,jt=grid_opt(true_x)
107    base,_a,_b=train(False); idea,_a,_b=train(True)
108    out={'device':device,'true_optimum':{'u':ut,'j':jt},'math_stability':stability_sweep(),'math_extra_minima':minima_sweep(),
109         'baseline':assess(base,ut,jt),'decision_aware':assess(idea,ut,jt)}
110    with open('results.json','w') as f: json.dump(out,f,indent=2)
111    print(json.dumps(out,indent=2))
112if __name__=='__main__': main()