1import numpy as np
 2
 3META = {
 4    'name': 'soft_graph_connectivity',
 5    'domain': 'graph-nn',
 6    'description': 'Synthetic soft graph edge-weight regression with eight-node adjacency structure.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    def make(rng, n):
11        x = np.empty((n, 8, 2), dtype=np.float32)
12        y = np.empty((n, 1), dtype=np.float32)
13        ii, jj = np.triu_indices(8, 1)
14        for b in range(n):
15            ang = np.linspace(0, 2*np.pi, 8, endpoint=False) + rng.normal(0, .06, 8)
16            rad = 1.0 + rng.normal(0, .05, 8)
17            pts = np.stack((rad*np.cos(ang), rad*np.sin(ang)), axis=1)
18            pts += rng.normal(0, .035, pts.shape)
19            w = np.zeros((8, 8), dtype=np.float32)
20            for i in range(8):
21                for j in range(i):
22                    d = np.linalg.norm(pts[i] - pts[j])
23                    ring = min((i-j) % 8, (j-i) % 8) == 1
24                    value = np.exp(-d*d/1.8) * (1.0 if ring else .42)
25                    w[i, j] = w[j, i] = max(.01, value + rng.normal(0, .012))
26            x[b] = pts
27            y[b, 0] = w[ii, jj].mean()
28        return x, y
29    xtr, ytr = make(np.random.RandomState(seed), n_train)
30    xte, yte = make(np.random.RandomState(seed + 5000), n_test)
31    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
32            'task': 'regression', 'metric': 'mse', 'out_dim': 1}