1import numpy as np
 2
 3META = {
 4    "name": "simplicial_cochain_regression",
 5    "domain": "simplicial_geometry_pde_like",
 6    "description": "Regression from noisy node and edge cochains on a filled oriented triangle; targets are compatible cochain observables."
 7}
 8
 9D0 = np.array([[-1., 1., 0.], [-1., 0., 1.], [0., -1., 1.]])
10D1 = np.array([[1., -1., 1.]])
11
12
13def get_dataset(seed, n_train, n_test):
14    def make(n, s):
15        rng = np.random.RandomState(s)
16        v = rng.normal(size=(n, 3, 2)).astype(np.float32)
17        e = np.einsum('ev,nvc->nec', D0, v).astype(np.float32)
18        x0 = v + 0.65 * rng.normal(size=v.shape).astype(np.float32)
19        x1 = e + 0.65 * rng.normal(size=e.shape).astype(np.float32)
20        x = np.concatenate([x0, x1], axis=1)
21        y = (v[:, :, 0].mean(axis=1) + 0.35 * e[:, :, 1].mean(axis=1)
22             + 0.15 * np.tanh(v[:, :, 1].sum(axis=1))).astype(np.float32)
23        return x, y[:, None]
24    xtr, ytr = make(n_train, seed)
25    xte, yte = make(n_test, seed + 5000)
26    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
27            "task": "regression", "metric": "mse", "out_dim": 1}