1import numpy as np
 2
 3META = {"name": "latent_mixture_transport", "domain": "diffusion-sampling", "description": "Regression from noisy source Gaussian-mixture states to component-specific transported Gaussian targets."}
 4
 5def get_dataset(seed, n_train, n_test):
 6    rng = np.random.default_rng(seed)
 7    d = 8
 8    def make(n):
 9        z = rng.integers(0, 2, size=n)
10        x = rng.normal(size=(n, d)).astype(np.float32)
11        x[:, 0] += np.where(z == 0, -2.0, 2.0)
12        x[:, 1] += np.where(z == 0, 1.0, -1.0)
13        y = np.empty((n, 2), dtype=np.float32)
14        y[:, 0] = 0.9*x[:, 0] + 0.25*x[:, 2] + np.where(z == 0, -2.5, 2.5)
15        y[:, 1] = -0.65*x[:, 1] + 0.35*x[:, 3] + np.where(z == 0, 1.8, -1.8)
16        y += rng.normal(scale=np.where(z[:, None] == 0, 0.65, 1.15), size=(n, 2)).astype(np.float32)
17        return x, y
18    xtr, ytr = make(n_train)
19    xte, yte = make(n_test)
20    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
21            "task": "regression", "metric": "mse", "input_shape": (d,), "out_dim": 2}