1import numpy as np
 2
 3META = {"name": "torus_cayley_graph", "domain": "graph-nn", "description": "Nodewise regression on labeled Z_6 x Z_6 Cayley torus graphs; targets are a fixed local relation convolution."}
 4N = 6
 5EDGES, RELS = [], []
 6for x in range(N):
 7    for y in range(N):
 8        u = x * N + y
 9        for r, (dx, dy) in enumerate(((1, 0), (-1, 0), (0, 1), (0, -1))):
10            v = ((x + dx) % N) * N + ((y + dy) % N)
11            EDGES.append((u, v)); RELS.append(r)
12EDGES = np.asarray(EDGES, dtype=np.int64)
13RELS = np.asarray(RELS, dtype=np.int64)
14COEF = np.asarray([1.0, -0.7, 0.45, -1.15], dtype=np.float32)
15
16def get_dataset(seed, n_train, n_test):
17    rng = np.random.RandomState(seed)
18    def make(n):
19        x = rng.normal(size=(n, N*N, 4)).astype(np.float32)
20        y = np.zeros((n, N*N, 1), dtype=np.float32)
21        for (u, v), r in zip(EDGES, RELS):
22            y[:, v, 0] += COEF[r] * x[:, u, 0]
23        y += rng.normal(0, 0.03, y.shape).astype(np.float32)
24        return x, y
25    xtr, ytr = make(n_train); xte, yte = make(n_test)
26    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
27            "task": "regression", "metric": "mse", "out_dim": 1}