1import numpy as np
 2
 3META = {
 4    "name": "conditional_ot_domains",
 5    "domain": "domain_generalization",
 6    "description": "Three source domains with shared class signal and domain-specific nuisance, plus a shifted target test domain."
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.default_rng(seed)
11    domains = (-2.0, 0.0, 2.0)
12    base = n_train // len(domains)
13    sizes = [base + (i < (n_train % len(domains))) for i in range(len(domains))]
14    xs, ys = [], []
15    for d, ns in zip(domains, sizes):
16        y = rng.integers(0, 2, ns)
17        signal = (2 * y - 1).astype(np.float32)
18        x = np.stack([
19            signal + rng.normal(0, 0.65, ns),
20            d + rng.normal(0, 0.75, ns),
21            np.full(ns, d / 2.0),
22        ], axis=1)
23        xs.append(x.astype(np.float32))
24        ys.append(y.astype(np.int64))
25    xtr, ytr = np.concatenate(xs), np.concatenate(ys)
26    yte = rng.integers(0, 2, n_test)
27    xte = np.stack([
28        (2 * yte - 1) + rng.normal(0, 0.65, n_test),
29        3.5 + rng.normal(0, 0.75, n_test),
30        np.full(n_test, 1.75),
31    ], axis=1).astype(np.float32)
32    return {
33        "xtr": xtr,
34        "ytr": ytr,
35        "xte": xte,
36        "yte": yte.astype(np.int64),
37        "task": "classification",
38        "metric": "cross_entropy",
39        "out_dim": 2,
40    }