1import numpy as np
 2
 3META = {
 4    "name": "categorical_energy",
 5    "domain": "discrete_energy_sampling",
 6    "description": "Synthetic categorical-vector classification with local and pairwise energy structure.",
 7}
 8
 9
10def _make(seed, n):
11    rng = np.random.default_rng(seed)
12    d, k = 8, 3
13    x = rng.integers(0, k, size=(n, d))
14    unary = np.array([[0.0, 0.45, 0.9], [0.1, 0.0, 0.55], [0.5, 0.15, 0.0]], dtype=np.float32)
15    e = unary[x, np.arange(d) % 3].sum(axis=1)
16    e += 0.55 * (x[:, :-1] == x[:, 1:]).sum(axis=1)
17    e += rng.normal(0.0, 0.18, n)
18    y = (e < np.median(e)).astype(np.int64)
19    oh = np.eye(k, dtype=np.float32)[x].reshape(n, d * k)
20    return oh, y
21
22
23def get_dataset(seed, n_train, n_test):
24    xtr, ytr = _make(int(seed) + 11, n_train)
25    xte, yte = _make(int(seed) + 29, n_test)
26    return {
27        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
28        "task": "classification", "metric": "err",
29        "input_shape": (24,), "out_dim": 2,
30    }