1import numpy as np
 2
 3META = {
 4    "name": "minkowski_composition",
 5    "domain": "structured_convex_composition",
 6    "description": "Regression from two direction-indexed convex latent tuples to their Minkowski-additive composition."
 7}
 8
 9N_DIR = 8
10D = 2
11ANGLES = np.linspace(0.0, 2.0 * np.pi, N_DIR, endpoint=False)
12
13
14def _tuple(rng):
15    center = rng.uniform(-1.0, 1.0, size=2)
16    radii = rng.uniform(0.25, 1.25, size=2)
17    phase = rng.uniform(-0.35, 0.35)
18    ang = ANGLES + phase
19    return center[None, :] + np.stack([radii[0] * np.cos(ang), radii[1] * np.sin(ang)], axis=1)
20
21
22def _make(seed, n):
23    rng = np.random.RandomState(seed)
24    x = np.empty((n, 2 * N_DIR * D), dtype=np.float32)
25    y = np.empty((n, N_DIR * D), dtype=np.float32)
26    for i in range(n):
27        a, b = _tuple(rng), _tuple(rng)
28        x[i] = np.concatenate([a.reshape(-1), b.reshape(-1)])
29        y[i] = (a + b).reshape(-1)
30    return x, y
31
32
33def get_dataset(seed, n_train, n_test):
34    xtr, ytr = _make(seed, n_train)
35    xte, yte = _make(seed + 5000, n_test)
36    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
37            "task": "regression", "metric": "mse", "out_dim": N_DIR * D}