1import numpy as np
 2
 3META = {
 4    "name": "geometric_graph_diffusion",
 5    "domain": "graph-neural-networks",
 6    "description": "Node-field regression on geometric kNN graphs; target is a short distance-weighted diffusion response."
 7}
 8
 9
10def _one(seed, n):
11    rng = np.random.RandomState(seed)
12    xy = rng.uniform(-1.0, 1.0, size=(24, 2)).astype(np.float32)
13    d2 = ((xy[:, None] - xy[None, :]) ** 2).sum(-1)
14    adj = np.zeros_like(d2, dtype=np.float32)
15    for i in range(24):
16        adj[i, np.argsort(d2[i])[1:6]] = 1.0
17    adj = np.maximum(adj, adj.T)
18    w = adj * np.maximum(d2, 1e-5)  # n=4: d^(n-2) = d^2
19    p = w / np.maximum(w.sum(1, keepdims=True), 1e-8)
20    q = 0.55 * np.eye(24, dtype=np.float32) + 0.45 * p
21    kernel = q @ q
22    root = int(np.argmin((xy ** 2).sum(1)))
23    x = rng.normal(size=(n, 24, 2)).astype(np.float32)
24    y = (x[:, :, 0] @ kernel[root]).astype(np.float32)
25    y += 0.08 * rng.normal(size=n).astype(np.float32)
26    coords = np.broadcast_to(xy, (n, 24, 2)).copy()
27    root_flag = np.zeros((n, 24, 1), dtype=np.float32)
28    root_flag[:, root, 0] = 1.0
29    return np.concatenate([x, coords, root_flag], axis=2)
30
31
32def get_dataset(seed, n_train=400, n_test=160):
33    # The graph is regenerated deterministically per split; coordinates and
34    # root marker are part of every sample, so the task is self-contained.
35    # Targets are generated with the same weighted diffusion rule as the idea,
36    # while evaluation remains ordinary regression MSE.
37    def make(seed0, n):
38        rng = np.random.RandomState(seed0)
39        xy = rng.uniform(-1.0, 1.0, size=(24, 2)).astype(np.float32)
40        d2 = ((xy[:, None] - xy[None, :]) ** 2).sum(-1)
41        adj = np.zeros_like(d2, dtype=np.float32)
42        for i in range(24):
43            adj[i, np.argsort(d2[i])[1:6]] = 1.0
44        adj = np.maximum(adj, adj.T)
45        w = adj * np.maximum(d2, 1e-5)
46        p = w / np.maximum(w.sum(1, keepdims=True), 1e-8)
47        q = 0.55 * np.eye(24, dtype=np.float32) + 0.45 * p
48        kernel = q @ q
49        root = int(np.argmin((xy ** 2).sum(1)))
50        x = rng.normal(size=(n, 24, 2)).astype(np.float32)
51        y = (x[:, :, 0] @ kernel[root]).astype(np.float32)
52        y += 0.08 * rng.normal(size=n).astype(np.float32)
53        coords = np.broadcast_to(xy, (n, 24, 2)).copy()
54        flag = np.zeros((n, 24, 1), dtype=np.float32); flag[:, root, 0] = 1.0
55        return np.concatenate([x, coords, flag], axis=2), y
56    xtr, ytr = make(seed, n_train)
57    xte, yte = make(seed + 5000, n_test)
58    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
59            "task": "regression", "metric": "mse", "out_dim": 1}