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