1import json, math, random
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED=1234
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8torch.set_num_threads(4)
 9try:
10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11    # Probe CUDA so an unavailable/broken shared device falls back cleanly.
12    if device.type=='cuda':
13        torch.zeros(1,device=device).sum().item()
14except Exception:
15    device=torch.device('cpu')
16
17lam=0.5
18
19def field_np(x):
20    r=np.sqrt(x[:,0]**2+x[:,1]**2)
21    th=np.arctan2(x[:,1],x[:,0])
22    sing=r**lam*np.sin(th/2)
23    smooth=0.2*np.sin(np.pi*x[:,0])*np.sin(np.pi*x[:,1])+0.1*x[:,0]*x[:,1]
24    return sing+smooth
25
26def sample_uniform(n): return np.random.rand(n,2)
27def sample_graded(n):
28    # Concentrate radial mass near the junction while retaining the full square.
29    # Independent Beta(.55,1) coordinates gives a near-corner radial grading.
30    return np.random.beta(.55,1.0,size=(n,2))
31
32class MLP(nn.Module):
33    def __init__(self, enriched=False):
34        super().__init__(); self.enriched=enriched
35        self.net=nn.Sequential(nn.Linear(2,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,1))
36        if enriched: self.a=nn.Parameter(torch.tensor([0.8]))
37    def forward(self,x):
38        z=self.net(x)
39        if self.enriched:
40            r=torch.sqrt((x*x).sum(1)+1e-12); th=torch.atan2(x[:,1],x[:,0])
41            phi=torch.sin(th/2).unsqueeze(1)
42            r=r.unsqueeze(1)
43            # cutoff is one in the relevant near-corner region and smoothly vanishes
44            r0=.7; q=torch.clamp(r/r0,0,1)
45            chi=1-3*q*q+2*q*q*q
46            sing=chi*r.pow(lam)*phi*self.a
47            return sing+(1-chi)*z+chi*z*r.pow(lam)
48        return z
49
50def train(enriched, sampler, n=1600, steps=1400):
51    x=torch.tensor(sampler(n),dtype=torch.float32,device=device)
52    y=torch.tensor(field_np(x.detach().cpu().numpy()),dtype=torch.float32,device=device).view(-1,1)
53    model=MLP(enriched).to(device)
54    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
55    for i in range(steps):
56        # fixed points makes the equal-evaluation comparison deterministic
57        pred=model(x); loss=((pred-y)**2).mean()
58        opt.zero_grad(); loss.backward(); opt.step()
59    return model, float(loss.detach().cpu())
60
61def evaluate(model):
62    # fixed dense test grid, plus a corner-only region
63    g=np.linspace(.001,1,121); xx,yy=np.meshgrid(g,g); X=np.c_[xx.ravel(),yy.ravel()]
64    with torch.no_grad(): p=model(torch.tensor(X,dtype=torch.float32,device=device)).cpu().numpy().ravel()
65    y=field_np(X)
66    rel=np.linalg.norm(p-y)/np.linalg.norm(y)
67    mask=np.sqrt((X*X).sum(1))<.12
68    corner=np.linalg.norm(p[mask]-y[mask])/np.linalg.norm(y[mask])
69    return float(rel),float(corner),int(mask.sum())
70
71def math_check():
72    # Along a fixed ray the singular part must have log-log slope lambda.
73    r=np.logspace(-6,-1,80); X=np.c_[r,r]
74    s=np.sqrt(2*r*r)**lam*np.sin(np.pi/8)
75    slope=np.polyfit(np.log(r),np.log(np.abs(s)),1)[0]
76    # Smooth remainder has bounded first radial derivative numerically at the corner.
77    # Compare its radial increments at two scales (increments should be O(r)).
78    def rem(t): return .2*np.sin(np.pi*t/np.sqrt(2))**2+.1*(t/np.sqrt(2))**2
79    inc1=abs(rem(2e-3)-rem(1e-3)); inc2=abs(rem(4e-3)-rem(2e-3))
80    return {'loglog_slope':float(slope),'expected_lambda':lam,'slope_abs_error':float(abs(slope-lam)), 'smooth_increment_ratio':float(inc2/inc1)}
81
82if __name__=='__main__':
83    check=math_check()
84    results={}
85    configs=[('baseline_uniform',False,sample_uniform),('baseline_graded',False,sample_graded),('enriched_uniform',True,sample_uniform),('enriched_graded',True,sample_graded)]
86    for name,en,samp in configs:
87        m,trainloss=train(en,samp)
88        rel,corner,count=evaluate(m)
89        results[name]={'train_mse':trainloss,'relative_l2':rel,'corner_relative_l2':corner,'corner_test_points':count}
90    # Verify grading actually changes allocation without changing point count.
91    u=sample_uniform(100000); g=sample_graded(100000)
92    grading={'n_uniform':len(u),'n_graded':len(g),'fraction_r_lt_.12_uniform':float((np.linalg.norm(u,axis=1)<.12).mean()),'fraction_r_lt_.12_graded':float((np.linalg.norm(g,axis=1)<.12).mean())}
93    out={'device':str(device),'math_check':check,'grading_check':grading,'results':results,'seed':SEED,'train_points':1600,'steps':1400}
94    print(json.dumps(out,indent=2))