1import json, math, random
 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)
 9
10def kernel_np(x, y, d, s, sigma, c=None, eps=1e-8):
11    if c is None: c = np.zeros(x.shape[-1])
12    r = np.linalg.norm(x-y, axis=-1) + eps
13    ux = np.linalg.norm(x-c, axis=-1) + eps
14    uy = np.linalg.norm(y-c, axis=-1) + eps
15    m = np.minimum(np.minimum(ux/r, uy/r), 1.0)
16    return r**(s-d) * m**(-sigma)
17
18def math_checks():
19    d, s, sigma = 3, 0.8, 0.35
20    x = np.array([[0.31, 0.22, 0.17], [0.47, -0.19, 0.28], [0.18, 0.41, -0.23]])
21    y = np.array([[-0.24, 0.16, 0.29], [-0.13, 0.39, -0.31], [-0.36, -0.11, 0.21]])
22    k1 = kernel_np(x, y, d, s, sigma)
23    scale = 3.0
24    k2 = kernel_np(scale*x, scale*y, d, s, sigma)
25    hom_err = float(np.max(np.abs(k2/(k1*scale**(s-d))-1)))
26    xa, ya = np.array([[1e-5, 0.0]]), np.array([[0.8, 0.0]])
27    plain = (np.linalg.norm(xa-ya, axis=1)+1e-8)**(s-d)
28    full = kernel_np(xa, ya, d, s, sigma)
29    amplification = float(full[0]/plain[0])
30    predicted = float(((np.linalg.norm(xa[0])+1e-8)/(np.linalg.norm(xa[0]-ya[0])+1e-8))**(-sigma))
31    radii = np.logspace(-7, -2, 3000)
32    integ = np.trapz(radii**(d-1-sigma), radii)
33    expected = (1e-2**(d-sigma)-1e-7**(d-sigma))/(d-sigma)
34    return {'homogeneity_max_relative_error': hom_err,
35            'near_origin_amplification': amplification,
36            'predicted_amplification': predicted,
37            'origin_integral_relative_error': float(abs(integ/expected-1)),
38            'admissible_d3_demo': bool(0 < s < d-2*sigma)}
39
40class AttnRegressor(nn.Module):
41    def __init__(self, biased, n=64, dim=32, s=.8, sigma=.35):
42        super().__init__(); self.biased=biased; self.n=n; self.dim=dim
43        self.inp=nn.Linear(2, dim); self.q=nn.Linear(dim, dim); self.k=nn.Linear(dim, dim)
44        self.v=nn.Linear(dim, dim); self.out=nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, 1))
45        self.s=s; self.sigma=sigma
46    def forward(self, xy):
47        z=self.inp(xy); q=self.q(z); k=self.k(z); v=self.v(z)
48        logits=q@k.transpose(-1,-2)/math.sqrt(self.dim)
49        if self.biased:
50            diff=xy[:,:,None,:]-xy[:,None,:,:]
51            r=torch.sqrt((diff*diff).sum(-1)+1e-8)+1e-5
52            u=torch.sqrt((xy*xy).sum(-1)+1e-8)+1e-5
53            m=torch.minimum(torch.minimum(u[:,:,None]/r,u[:,None,:]/r),torch.ones_like(r))
54            K=r**(self.s-2)*m**(-self.sigma)
55            logits=logits+0.7*torch.log(K+1e-8)
56        a=torch.softmax(logits, -1)
57        return self.out(a@v)
58
59def make_data(num, n=64):
60    xy=np.random.uniform(-1,1,(num,n,2)).astype('float32')
61    r=np.sqrt((xy*xy).sum(-1)+1e-5)
62    # clipped cusp with a smooth multiscale component
63    target=np.minimum(1.0, r**(-0.65))/2 + .15*np.sin(8*xy[:,:,0])*np.cos(6*xy[:,:,1])
64    return torch.tensor(xy), torch.tensor(target[:,:,None], dtype=torch.float32)
65
66def run():
67    checks=math_checks(); device='cuda' if torch.cuda.is_available() else 'cpu'
68    try:
69        train_x, train_y=make_data(256); test_x, test_y=make_data(64)
70        train_x,train_y,test_x,test_y=[z.to(device) for z in (train_x,train_y,test_x,test_y)]
71        results={}
72        for biased in (False, True):
73            torch.manual_seed(SEED); model=AttnRegressor(biased).to(device)
74            opt=torch.optim.Adam(model.parameters(),lr=3e-3); loss_fn=nn.MSELoss()
75            curve=[]
76            for step in range(180):
77                idx=torch.randint(0,train_x.shape[0],(16,),device=device)
78                loss=loss_fn(model(train_x[idx]),train_y[idx])
79                opt.zero_grad(); loss.backward(); opt.step()
80                if step in (0,29,89,179):
81                    with torch.no_grad(): curve.append(float(loss_fn(model(test_x),test_y).cpu()))
82            with torch.no_grad(): final=float(loss_fn(model(test_x),test_y).cpu())
83            results['idea' if biased else 'baseline']={'test_mse_curve':curve,'final_test_mse':final}
84    except Exception as e:
85        device='cpu'; results={'error':repr(e)}
86    out={'seed':SEED,'device':device,'checks':checks,'results':results}
87    Path('results.json').write_text(json.dumps(out,indent=2))
88    print(json.dumps(out,indent=2))
89if __name__=='__main__': run()