1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED=17
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type=='cuda':
 12        torch.cuda.manual_seed_all(SEED)
 13        # force a tiny allocation to detect unusable CUDA contexts
 14        torch.zeros(1, device=device)
 15except Exception:
 16    device=torch.device('cpu')
 17
 18a=0.75
 19s=0.10
 20lam_g=0.01
 21EPS=1e-6
 22
 23class MLP(nn.Module):
 24    def __init__(self, width=48, depth=3):
 25        super().__init__()
 26        layers=[nn.Linear(1,width), nn.Tanh()]
 27        for _ in range(depth-1): layers += [nn.Linear(width,width), nn.Tanh()]
 28        layers += [nn.Linear(width,1)]
 29        self.net=nn.Sequential(*layers)
 30        for m in self.modules():
 31            if isinstance(m,nn.Linear):
 32                nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
 33    def forward(self,x): return self.net(x)
 34
 35def d_exact(x): return torch.minimum(x,1-x)
 36def target(x):
 37    # Smooth quotient times the fractional boundary factor.
 38    d=d_exact(x)
 39    return d.pow(a)*(1.0+0.35*torch.sin(2*math.pi*x)+0.15*x)
 40def weighted_grad(v,x):
 41    grad=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
 42    d=d_exact(x).clamp_min(EPS)
 43    return d.pow(1-a+s)*grad
 44
 45def train(kind, steps=1800, n=96):
 46    torch.manual_seed(SEED)
 47    model=MLP().to(device)
 48    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 49    hist=[]; t0=time.time()
 50    for it in range(steps):
 51        x=torch.rand(n,1,device=device)
 52        # avoid the cusp of min(x,1-x) at the midpoint for this gradient regularizer
 53        x.requires_grad_(kind=='idea')
 54        y=target(x)
 55        raw=model(x)
 56        if kind=='idea':
 57            u=d_exact(x).clamp_min(EPS).pow(a)*raw
 58            wg=weighted_grad(raw,x)
 59            loss=((u-y)**2).mean()+lam_g*(wg**2).mean()
 60        else:
 61            u=raw
 62            # Standard free-output PINN-style boundary penalty.
 63            xb=torch.tensor([[0.0],[1.0]],device=device)
 64            loss=((u-y)**2).mean()+10.0*(model(xb)**2).mean()
 65        opt.zero_grad(); loss.backward(); opt.step()
 66        if it in (0,99,499,999,1799): hist.append(float(loss.detach().cpu()))
 67    with torch.no_grad():
 68        xx=torch.linspace(1e-5,1-1e-5,4000,device=device).view(-1,1)
 69        pred=(d_exact(xx).clamp_min(EPS).pow(a)*model(xx) if kind=='idea' else model(xx))
 70        yy=target(xx)
 71        rel=float(torch.sqrt(((pred-yy)**2).mean())/torch.sqrt((yy**2).mean()))
 72        strip=((xx<0.05)|(xx>0.95))
 73        striperr=float(torch.sqrt(((pred[strip]-yy[strip])**2).mean())/torch.sqrt((yy[strip]**2).mean()))
 74        maxb=float(torch.max(torch.abs(pred[strip]-yy[strip])))
 75    return {'loss_checkpoints':hist,'relative_L2':rel,'boundary_strip_relative_L2':striperr,'boundary_strip_max_abs':maxb,'seconds':time.time()-t0}
 76
 77def math_check():
 78    # For smooth v, the claimed weighted derivative should vanish at boundary,
 79    # whereas u'=a*d^(a-1)*v has the fractional singular scaling.
 80    x=torch.logspace(-6,-2,80,dtype=torch.float64)
 81    v=1+0.4*x
 82    vp=torch.full_like(x,0.4)
 83    up=a*x**(a-1)*v+x**a*vp
 84    w=x**(1-a+s)*vp
 85    # Fit log-log slopes; expected u' slope a-1 and weighted smooth term slope 1-a+s.
 86    su=float(np.polyfit(np.log(x.numpy()),np.log(torch.abs(up).numpy()),1)[0])
 87    sw=float(np.polyfit(np.log(x.numpy()),np.log(torch.abs(w).numpy()),1)[0])
 88    ratio=float((w[-1]/w[0]).item())
 89    return {'unweighted_uprime_log_slope':su,'expected_unweighted_slope':a-1,
 90            'weighted_vprime_log_slope':sw,'expected_weighted_slope':1-a+s,
 91            'weighted_growth_across_range':ratio,'weighted_decreases_to_boundary':bool(w[0]<w[-1])}
 92
 93def main():
 94    out={'device':str(device),'seed':SEED,'a':a,'s':s,'lambda_g':lam_g,
 95         'math_check':math_check()}
 96    out['baseline']=train('baseline'); out['idea']=train('idea')
 97    out['idea_minus_baseline_relative_L2']=out['idea']['relative_L2']-out['baseline']['relative_L2']
 98    out['idea_minus_baseline_boundary_strip']=out['idea']['boundary_strip_relative_L2']-out['baseline']['boundary_strip_relative_L2']
 99    Path('results.json').write_text(json.dumps(out,indent=2))
100    print(json.dumps(out,indent=2))
101if __name__=='__main__': main()