1import numpy as np
 2
 3META = {
 4    "name": "stochastic_intervention_scm",
 5    "domain": "causal-world-model",
 6    "description": "Random intervention X causes mediator M and outcome Y; tests joint modeling under intervention shift and mediator selection."
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.RandomState(seed)
11    # Context is observed but does not explain away the stochastic intervention.
12    c_train = rng.randn(n_train, 1).astype("float32")
13    c_test = rng.randn(n_test, 1).astype("float32")
14    x_train = (rng.randn(n_train, 1) + 0.35*c_train).astype("float32")
15    x_test = (2.0*rng.randn(n_test, 1) + 0.35*c_test).astype("float32")
16    m_train = (x_train + 0.3*c_train + rng.randn(n_train, 1)).astype("float32")
17    m_test = (x_test + 0.3*c_test + rng.randn(n_test, 1)).astype("float32")
18    y_train = (x_train + m_train + 0.5*c_train + rng.randn(n_train, 1)).astype("float32")
19    y_test = (x_test + m_test + 0.5*c_test + rng.randn(n_test, 1)).astype("float32")
20    # x is supplied as an observed feature for training/evaluation; the task target is y.
21    return {"xtr": np.concatenate([c_train, x_train, m_train], 1),
22            "ytr": y_train[:, 0], "xte": np.concatenate([c_test, x_test, m_test], 1),
23            "yte": y_test[:, 0], "task": "regression", "metric": "mse"}