1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 892
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float64)
  9try:
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device = torch.device('cpu')
 13
 14
 15def lap_scalar(z, xy):
 16    g = torch.autograd.grad(z.sum(), xy, create_graph=True, retain_graph=True)[0]
 17    vals = []
 18    for j in range(2):
 19        gj = torch.autograd.grad(g[:, j].sum(), xy, create_graph=True, retain_graph=True)[0][:, j]
 20        vals.append(gj)
 21    return vals[0] + vals[1]
 22
 23
 24def q_residuals(q, xy):
 25    qr, qi = q[:, 0], q[:, 1]
 26    gr = torch.autograd.grad(qr.sum(), xy, create_graph=True, retain_graph=True)[0]
 27    gi = torch.autograd.grad(qi.sum(), xy, create_graph=True, retain_graph=True)[0]
 28    lqr, lqi = lap_scalar(qr, xy), lap_scalar(qi, xy)
 29    er = gr[:, 0]**2 - gi[:, 0]**2 + gr[:, 1]**2 - gi[:, 1]**2
 30    ei = 2 * (gr[:, 0]*gi[:, 0] + gr[:, 1]*gi[:, 1])
 31    return lqr, lqi, er, ei
 32
 33
 34def q_pert(xy, eps):
 35    x, y = xy[:, 0], xy[:, 1]
 36    return torch.stack((x + eps*x*x, y), dim=1)
 37
 38
 39def toy_verification():
 40    a = torch.linspace(-1, 1, 31, device=device)
 41    X, Y = torch.meshgrid(a, a, indexing='ij')
 42    base = torch.stack((X.flatten(), Y.flatten()), 1)
 43    p = base.clone().requires_grad_(True)
 44    z = q_pert(p, 0.0)
 45    vals = q_residuals(z, p)
 46    exact_max = float(torch.max(torch.stack([v.abs().max() for v in vals])))
 47    epslist = np.array([1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1])
 48    rows = []
 49    for e in epslist:
 50        p = base.clone().requires_grad_(True)
 51        vals = q_residuals(q_pert(p, float(e)), p)
 52        h = torch.sqrt(vals[0]**2 + vals[1]**2).mean().item()
 53        en = torch.sqrt(vals[2]**2 + vals[3]**2).mean().item()
 54        rows.append((float(e), h, en))
 55    # Predictions: harmonic residual is identically zero for this perturbation;
 56    # null-gradient residual is 4 eps*x + 4 eps^2*x^2, hence RMS scales linearly.
 57    earr = np.array([r[0] for r in rows]); nerr = np.array([r[2] for r in rows])
 58    slope = float(np.polyfit(np.log(earr), np.log(nerr), 1)[0])
 59    pred_coeff = 4.0 * math.sqrt(float((base[:,0]**2).mean()))
 60    observed_coeff = float(np.polyfit(earr[:4], nerr[:4], 1)[0])
 61    return {'exact_identity_max_abs': exact_max, 'sweep': [dict(eps=e,harmonic=h,null_gradient=n) for e,h,n in rows],
 62            'predicted_null_rms_linear_coefficient': pred_coeff, 'observed_small_eps_coefficient': observed_coeff,
 63            'observed_loglog_scaling_exponent': slope, 'predictions': {'harmonic stays zero': True, 'null residual O(eps)': True}}
 64
 65
 66class MLP(nn.Module):
 67    def __init__(self, din, width=24):
 68        super().__init__()
 69        self.net = nn.Sequential(nn.Linear(din,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,1))
 70    def forward(self,x): return self.net(x).squeeze(-1)
 71
 72
 73class Ansatz(nn.Module):
 74    def __init__(self, width=20):
 75        super().__init__()
 76        self.q = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,2))
 77        self.v = MLP(2,width)
 78        with torch.no_grad():
 79            self.q[-1].weight.mul_(0.02)
 80            self.q[-1].bias.copy_(torch.tensor([0.,0.]))
 81    def forward(self,xy):
 82        raw = self.q(xy)
 83        # affine canonical coordinate plus trainable residual
 84        q = torch.stack((xy[:,0]+raw[:,0], xy[:,1]+raw[:,1]),1)
 85        return self.v(q), q
 86
 87
 88def lap_u(u, xy): return lap_scalar(u, xy)
 89
 90def train_model(kind, steps=500, n=96):
 91    model = MLP(2,24) if kind=='baseline' else Ansatz(20)
 92    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 93    for step in range(steps):
 94        xy=(torch.rand(n,2,device=device)*2-1).requires_grad_(True)
 95        if kind=='baseline': u=model(xy); q=None
 96        else: u,q=model(xy)
 97        # target harmonic solution u=x^2-y^2 with boundary supervision
 98        target=xy[:,0]**2-xy[:,1]**2
 99        pde=lap_u(u,xy)
100        loss=(pde**2).mean()+10*((u-target)**2).mean()
101        if q is not None:
102            qr,qi,er,ei=q_residuals(q,xy)
103            ramp=min(1.,step/50.)
104            loss=loss+ramp*(qr.square().mean()+qi.square().mean()+er.square().mean()+ei.square().mean())
105        opt.zero_grad(); loss.backward(); opt.step()
106    with torch.no_grad():
107        test=(torch.rand(512,2,device=device)*2-1)
108        target=test[:,0]**2-test[:,1]**2
109        if kind=='baseline': pred=model(test); q=None
110        else: pred,q=model(test)
111        mse=((pred-target)**2).mean().item()
112    test=test.requires_grad_(True)
113    if kind=='baseline': pred=model(test)
114    else: pred,q=model(test)
115    residual=lap_u(pred,test).abs().mean().item()
116    qres=None
117    if q is not None:
118        rr=q_residuals(q,test); qres=float(torch.sqrt(sum(x.square() for x in rr)).mean())
119    return {'test_mse':mse,'mean_abs_laplacian':residual,'coordinate_constraint_rms':qres}
120
121
122def main():
123    verify=toy_verification()
124    results={}
125    for k in ['baseline','ansatz']:
126        try: results[k]=train_model(k)
127        except RuntimeError:
128            global device
129            device=torch.device('cpu'); results[k]=train_model(k)
130    out={'device':str(device),'verification':verify,'comparison':results}
131    print(json.dumps(out,indent=2))
132    with open('results.json','w') as f: json.dump(out,f,indent=2)
133
134if __name__=='__main__': main()