1import numpy as np
 2
 3META = {
 4    "name": "separable_diffusion_field",
 5    "domain": "pde",
 6    "description": "3D voxel fields from an anisotropic separable diffusion map, with scalar energy target."
 7}
 8
 9
10def get_dataset(seed, n_train, n_test):
11    rng = np.random.RandomState(seed)
12    shape = (4, 4, 4)
13    q = np.arange(4, dtype=np.float32)
14    z = np.exp(-0.5 * ((q[:, None] - q[None, :]) / 0.9) ** 2)
15    z /= z.sum(axis=1, keepdims=True)
16    w = np.linspace(-1.0, 1.0, 64).astype(np.float32)
17
18    def make(n):
19        x = rng.normal(0, 1, (n,) + shape).astype(np.float32)
20        y = np.einsum('ab,nbcd->nacd', z, x)
21        y = np.einsum('ab,nacd->nbcd', z, y)
22        y = np.einsum('ab,nbcd->nabc', z, y)
23        # Scalar PDE observable; nonlinear term prevents a trivial identity rule.
24        target = (y.reshape(n, 64) @ w / 8.0 + 0.15 * np.mean(np.tanh(x), axis=(1, 2, 3)))
25        return x, target.astype(np.float32)[:, None]
26
27    xtr, ytr = make(n_train)
28    xte, yte = make(n_test)
29    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
30            "task": "regression", "metric": "mse", "out_dim": 1}