tramdag.simulations

 71@dataclass
 72class MagicMrClean:
 73    """SCM generator for the synthetic stroke cohort.
 74
 75    Args:
 76        variant: ``"ls"`` (all linear shifts) or ``"nl"`` (mild non-linearities).
 77        seed: master seed; each draw uses an independent child stream.
 78    """
 79
 80    variant: str = "nl"
 81    seed: int = 7
 82
 83    def __post_init__(self):
 84        """Validate the variant and set the switch for the non-linear terms.
 85
 86        Returns
 87        -------
 88        None
 89
 90        Raises
 91        ------
 92        ValueError
 93            If ``variant`` is neither ``"ls"`` nor ``"nl"``.
 94        """
 95        if self.variant not in ("ls", "nl"):
 96            raise ValueError(f"variant must be 'ls' or 'nl', got {self.variant!r}")
 97        self.nl = float(self.variant == "nl")  # 0.0 disables the ★ terms
 98
 99    # ------------------------------------------------------------------ latents
100    def draw_latents(self, n: int, rng: np.random.Generator) -> dict[str, np.ndarray]:
101        """Draw the latent noise of every variable.
102
103        Parameters
104        ----------
105        n : int
106            Number of rows to draw.
107        rng : np.random.Generator
108            Random source.
109
110        Returns
111        -------
112        dict[str, np.ndarray]
113            One array of length ``n`` per variable.
114        """
115        return {k: _logistic(rng, n) for k in COLUMNS}
116
117    # --------------------------------------------------------------------- SCM
118    def simulate(
119        self,
120        n: int | None = None,
121        *,
122        rng: np.random.Generator | None = None,
123        randomize_T: bool = False,
124        population: str = "obs",
125        do: dict[str, float] | None = None,
126        latents: dict[str, np.ndarray] | None = None,
127    ) -> pd.DataFrame:
128        """Forward-sample the SCM.
129
130        Args:
131            randomize_T: if True, assign T ~ Bernoulli(0.5) independently of its
132                parents (the RCT design); otherwise T follows the confounded
133                observational mechanism.
134            population: covariate population. ``"obs"`` is the full observational
135                cohort; ``"rct"`` mimics trial inclusion — a **younger** enrolled
136                population (age location shifted down). Only the ``Age`` source
137                marginal differs; all structural equations are unchanged. With the
138                heterogeneous ``nl`` treatment effect this fit-vs-eval shift is
139                what biases an all-``ls`` model (it cannot extrapolate ``tau(Age)``
140                from the older obs cohort to the younger trial). The ``ls`` DGP has
141                constant ``tau`` and is therefore unaffected.
142            do: hard interventions {node: value} (graph mutilation); the node is
143                clamped and its structural equation skipped.
144            latents: reuse a fixed latent draw (for counterfactuals / paired
145                interventions). If given, ``n`` and ``rng`` are ignored.
146        """
147        do = do or {}
148        if latents is None:
149            if n is None:
150                raise ValueError("provide either n or latents")
151            rng = rng or np.random.default_rng(self.seed)
152            latents = self.draw_latents(n, rng)
153        else:
154            n = len(next(iter(latents.values())))
155        nl = self.nl
156        age_loc = 73.0 - (9.0 if population == "rct" else 0.0)  # trial enrolls younger
157
158        # --- Age: location-scale logistic (sd ~ 13.8), clipped to a plausible range
159        if "Age" in do:
160            Age = np.full(n, float(do["Age"]))
161        else:
162            Age = np.clip(age_loc + 7.6 * latents["Age"], 20.0, 103.0)
163        a = (Age - 73.0) / 10.0  # standardized age
164        relu_a = np.maximum(a, 0.0)
165
166        # --- mRS_pre: pre-stroke disability, worse (and ★ accelerating) with age
167        if "mRS_pre" in do:
168            mRS_pre = np.full(n, float(do["mRS_pre"]))
169        else:
170            eta_pre = 0.55 * a + nl * (0.18 * relu_a**2)
171            mRS_pre = _ordinal(eta_pre, CUTS_PRE, latents["mRS_pre"]).astype(float)
172
173        # --- NIHSSa: stroke severity; linear shift on parents, ★ age^2 in nl.
174        #     Free monotone marginal map keeps it in the realistic [6, 42] range.
175        if "NIHSSa" in do:
176            NIHSSa = np.full(n, float(do["NIHSSa"]))
177        else:
178            shift_nih = 0.45 * a + 0.25 * mRS_pre + nl * (0.20 * relu_a**2)
179            lat_nih = shift_nih + 0.85 * latents["NIHSSa"]
180            NIHSSa = np.clip(6.0 + 36.0 * _sigmoid((lat_nih - 1.75) / 1.5), 6.0, 42.0)
181        s = (NIHSSa - 15.0) / 6.0  # standardized severity
182
183        # --- T: thrombectomy assignment. Observational mechanism is confounded by
184        #     age and severity; ★ in nl it is smoothly withheld from the very old.
185        if "T" in do:
186            T = np.full(n, float(do["T"]))
187        elif randomize_T:
188            T = (latents["T"] > 0.0).astype(float)  # Bernoulli(0.5), latent-driven
189        else:
190            logit_T = (
191                1.9
192                - 0.45 * np.maximum(a + 0.5, 0.0)
193                - 0.30 * np.maximum(s - 1.0, 0.0)
194                + nl * (-1.3 * _sigmoid((Age - 82.0) / 4.0))
195            )
196            T = (latents["T"] > -logit_T).astype(float)  # P(T=1) = sigmoid(logit_T)
197
198        # --- mRS_3m: 3-month outcome. Treatment lowers the latent (better outcome);
199        #     ★ in nl the benefit tau(Age) fades to ~0 in the elderly.
200        if "mRS_3m" in do:
201            mRS_3m = np.full(n, float(do["mRS_3m"]))
202        else:
203            tau = -0.85 + nl * (0.85 - 0.85 * _sigmoid((78.0 - Age) / 6.0))
204            zeta = 0.85 * s + 0.55 * a + 0.45 * mRS_pre + tau * T
205            mRS_3m = _ordinal(zeta, CUTS_Y, latents["mRS_3m"]).astype(float)
206
207        return pd.DataFrame(
208            {"Age": Age, "mRS_pre": mRS_pre, "NIHSSa": NIHSSa, "T": T, "mRS_3m": mRS_3m}
209        )
210
211    # ----------------------------------------------------------------- datasets
212    def observational(self, n: int, seed_offset: int = 0) -> pd.DataFrame:
213        """Draw an observational sample.
214
215        Parameters
216        ----------
217        n : int
218            Number of rows.
219        seed_offset : int, optional
220            Added to the generator seed, by default ``0``.
221
222        Returns
223        -------
224        pd.DataFrame
225            The sample.
226        """
227        rng = np.random.default_rng(self.seed + 1 + seed_offset)
228        return self.simulate(n, rng=rng)
229
230    def rct(self, n: int, seed_offset: int = 0) -> pd.DataFrame:
231        """Draw a randomized-trial sample.
232
233        Treatment is randomized and the covariates follow the trial population, so
234        the contrast this yields is unconfounded.
235
236        Parameters
237        ----------
238        n : int
239            Number of rows.
240        seed_offset : int, optional
241            Added to the generator seed, by default ``0``.
242
243        Returns
244        -------
245        pd.DataFrame
246            The sample.
247        """
248        rng = np.random.default_rng(self.seed + 1001 + seed_offset)
249        return self.simulate(n, rng=rng, randomize_T=True, population="rct")
250
251    # -------------------------------------------------------------- ground truth
252    def true_ate(self, n: int = 500_000, on: str = "rct") -> dict:
253        """Estimate the true ATE of T on ``P(mRS_3m <= 2)`` by Monte Carlo.
254
255        Both arms use the same latent draw, so the result is the do-effect and
256        carries none of the confounding in T.
257
258        ``on`` selects the covariate population the ATE is averaged over:
259        ``"rct"`` (default) mirrors the **younger trial** population that
260        :meth:`rct` enrols and that ``evaluate_rct`` scores on; ``"obs"`` the
261        observational cohort.
262        """
263        rng = np.random.default_rng(self.seed + 9001)
264        latents = self.draw_latents(n, rng)
265        d0 = self.simulate(latents=latents, population=on, do={"T": 0})
266        d1 = self.simulate(latents=latents, population=on, do={"T": 1})
267        good0 = float((d0["mRS_3m"] <= 2).mean())
268        good1 = float((d1["mRS_3m"] <= 2).mean())
269        # naive (observational, confounded) contrast for contrast with the truth
270        obs = self.observational(n, seed_offset=777)
271        naive = float((obs.loc[obs["T"] == 1, "mRS_3m"] <= 2).mean()) - float(
272            (obs.loc[obs["T"] == 0, "mRS_3m"] <= 2).mean()
273        )
274        return {
275            "p_good_do_T0": good0,
276            "p_good_do_T1": good1,
277            "ate_population": on,
278            "true_ate": good1 - good0,
279            "naive_obs_diff": naive,
280            "mc_n": n,
281        }
282
283    def counterfactual_pair(
284        self, n: int, do: dict[str, float], seed_offset: int = 0
285    ) -> tuple[pd.DataFrame, pd.DataFrame]:
286        """Draw a factual sample and its counterfactual under ``do``.
287
288        Both share the same latents, so the pair gives true individual
289        counterfactuals. Real data cannot supply these. Use them to score the
290        abduction of the flow.
291        """
292        rng = np.random.default_rng(self.seed + 2 + seed_offset)
293        latents = self.draw_latents(n, rng)
294        factual = self.simulate(latents=latents)
295        cf = self.simulate(latents=latents, do=do)
296        return factual, cf

178class TriangleContinuous(_TriangleBase):
179    """Paper Sec. 6.1: the all-continuous triangle.
180
181    ``h(x3|x1,x2) = 0.63 x3 - 0.2 x1 - f(x2)``.
182    """
183
184    family = "continuous"
185
186    def _x3(self, x1, x2, do, latents):
187        if "x3" in do:
188            return _clamp(do["x3"], len(x1))
189        return (latents["x3"] + 0.2 * x1 + self.f_callable(x2)) / 0.63
190
191    def paper_truth(self) -> dict:
192        """State the true parameters of this data-generating process.
193
194        Returns
195        -------
196        dict
197            The coefficients and transformation functions used to generate the
198            data, in the notation of the paper.
199        """
200        t = {
201            "beta12": 2.0,
202            "beta13": -0.2,
203            "h2": "5*x2 + 2*x1",
204            "h3": f"0.63*x3 - 0.2*x1 - ({F_VARIANTS[self.f][1]})",
205        }
206        if self.f == "linear":
207            t["beta23"] = 0.3
208        return t

211class TriangleMixed(_TriangleBase):
212    """Paper Sec. 6.2: the triangle with an ordinal x3.
213
214    x3 has 4 levels, stored as 0 to 3, with
215    ``level = #{k : u3 > theta_k + 0.2 x1 + f(x2)}`` and
216    ``theta = (-2, 0.42, 1.02)``.
217    """
218
219    family = "mixed"
220    theta = THETA_MIXED
221
222    def _x3(self, x1, x2, do, latents):
223        if "x3" in do:
224            return _clamp(do["x3"], len(x1))
225        cuts = self.theta[None, :] + (0.2 * x1 + self.f_callable(x2))[:, None]
226        return (latents["x3"][:, None] > cuts).sum(axis=1).astype(float)
227
228    def true_pmf(self, x1: np.ndarray, x2: np.ndarray) -> np.ndarray:
229        """Analytic (n, 4) class probabilities given parents."""
230        shift = 0.2 * np.asarray(x1, float) + self.f_callable(np.asarray(x2, float))
231        cuts = self.theta[None, :] + shift[:, None]
232        cdf = 1.0 / (1.0 + np.exp(-cuts))
233        cdf = np.concatenate(
234            [np.zeros((len(cdf), 1)), cdf, np.ones((len(cdf), 1))], axis=1
235        )
236        return np.diff(cdf, axis=1)
237
238    def paper_truth(self) -> dict:
239        """State the true parameters of this data-generating process.
240
241        Returns
242        -------
243        dict
244            The coefficients, cutpoints and shift functions used to generate the
245            data, in the notation of the paper.
246        """
247        t = {
248            "beta13": 0.2,
249            "theta": self.theta.tolist(),
250            "h3": f"theta_k + 0.2*x1 + ({F_VARIANTS[self.f][1]})",
251            "levels": 4,
252            "level_offset": "paper counts 1..4, stored 0..3",
253            "zuko_sign": -1,
254        }
255        if self.f == "linear":
256            t["beta23"] = -0.3
257        return t
258
259    def zuko_expectations(self) -> dict:
260        """Give the same truth in the conventions of ``CausalFlowDAG``.
261
262        The ordinal shift is subtracted here and added in the paper, so the
263        expected weights carry the opposite sign.
264
265        Returns
266        -------
267        dict
268            Expected fitted values, keyed by parameter name.
269        """
270        exp = {
271            "w_x2_from_x1": 2.0,
272            "w_x3_from_x1": -0.2,
273            "theta": self.theta.tolist(),
274            "cs_curve": "-f(x2) + const",
275        }
276        if self.f == "linear":
277            exp["w_x3_from_x2"] = 0.3
278        return exp

 37@dataclass
 38class VacaTriangle:
 39    """SCM generator for the VACA bimodal triangle."""
 40
 41    seed: int = 42
 42
 43    def draw_latents(self, n: int, rng: np.random.Generator) -> dict[str, np.ndarray]:
 44        """Draw the latent noise of every variable.
 45
 46        Parameters
 47        ----------
 48        n : int
 49            Number of rows to draw.
 50        rng : np.random.Generator
 51            Random source.
 52
 53        Returns
 54        -------
 55        dict[str, np.ndarray]
 56            One array of length ``n`` per variable.
 57        """
 58        return {
 59            "x1_mix": rng.uniform(size=n),
 60            "x1_a": rng.normal(size=n),  # N(-2, sqrt(1.5)) branch
 61            "x1_b": rng.normal(size=n),  # N(1.5, 1) branch
 62            "x2": rng.normal(size=n),
 63            "x3": rng.normal(size=n),
 64        }
 65
 66    def simulate(
 67        self,
 68        n: int | None = None,
 69        *,
 70        rng: np.random.Generator | None = None,
 71        do: dict[str, float] | None = None,
 72        latents: dict[str, np.ndarray] | None = None,
 73    ) -> pd.DataFrame:
 74        """Simulate the SCM, with optional interventions and reused latents.
 75
 76        Parameters
 77        ----------
 78        n : int | None, optional
 79            Number of rows, by default ``None``. Then ``latents`` sets the count.
 80        rng : np.random.Generator | None, optional
 81            Random source, by default ``None``.
 82        do : dict[str, float] | None, optional
 83            Variables to hold at a fixed value, by default ``None``.
 84        latents : dict[str, np.ndarray] | None, optional
 85            Latent values to reuse, by default ``None``. Then they are drawn fresh.
 86
 87        Returns
 88        -------
 89        pd.DataFrame
 90            One column per variable.
 91        """
 92        do = do or {}
 93        if latents is None:
 94            if n is None:
 95                raise ValueError("provide either n or latents")
 96            rng = rng or np.random.default_rng(self.seed)
 97            latents = self.draw_latents(n, rng)
 98        n = len(latents["x2"])
 99
100        if "x1" in do:
101            x1 = np.full(n, float(do["x1"]))
102        else:
103            x1 = np.where(
104                latents["x1_mix"] < 0.5,
105                -2.0 + np.sqrt(1.5) * latents["x1_a"],
106                1.5 + 1.0 * latents["x1_b"],
107            )
108        x2 = np.full(n, float(do["x2"])) if "x2" in do else -x1 + latents["x2"]
109        x3 = (
110            np.full(n, float(do["x3"]))
111            if "x3" in do
112            else x1 + 0.25 * x2 + latents["x3"]
113        )
114        return pd.DataFrame({"x1": x1, "x2": x2, "x3": x3})
115
116    # ----------------------------------------------------------------- datasets
117    def observational(self, n: int, seed_offset: int = 0) -> pd.DataFrame:
118        """Draw an observational sample.
119
120        Parameters
121        ----------
122        n : int
123            Number of rows.
124        seed_offset : int, optional
125            Added to the generator seed, by default ``0``.
126
127        Returns
128        -------
129        pd.DataFrame
130            The sample.
131        """
132        rng = np.random.default_rng(self.seed + 1 + seed_offset)
133        return self.simulate(n, rng=rng)
134
135    def interventional(
136        self, n: int, do: dict[str, float], seed_offset: int = 0
137    ) -> pd.DataFrame:
138        """Draw a sample under an intervention.
139
140        Parameters
141        ----------
142        n : int
143            Number of rows.
144        do : dict[str, float]
145            Variables to hold at a fixed value.
146        seed_offset : int, optional
147            Added to the generator seed, by default ``0``.
148
149        Returns
150        -------
151        pd.DataFrame
152            The sample.
153        """
154        rng = np.random.default_rng(self.seed + 501 + seed_offset)
155        return self.simulate(n, rng=rng, do=do)
156
157    # -------------------------------------------------------------- ground truth
158    def true_moments(self, mc_n: int = 1_000_000) -> dict:
159        """Observational moments + the analytic moments of x3 under do(x2 = a).
160
161        Under do(x2=a): x3 = x1 + 0.25 a + N(0,1), so E = E[x1] + 0.25 a and
162        Var = Var[x1] + 1 — exact, but MC values are stored too (same estimator
163        a test would use).
164        """
165        mu1 = 0.5 * (-2.0) + 0.5 * 1.5
166        var1 = 0.5 * (1.5 + (-2.0 - mu1) ** 2) + 0.5 * (1.0 + (1.5 - mu1) ** 2)
167        obs = self.observational(mc_n, seed_offset=777)
168        out = {
169            "mc_n": mc_n,
170            "obs_mean": {c: float(obs[c].mean()) for c in obs},
171            "obs_std": {c: float(obs[c].std()) for c in obs},
172            "do_x2": {},
173        }
174        for a in DO_X2_VALUES:
175            out["do_x2"][str(a)] = {
176                "mean_x3_analytic": mu1 + 0.25 * a,
177                "std_x3_analytic": float(np.sqrt(var1 + 1.0)),
178            }
179        return out

 40@dataclass
 41class Carefl4:
 42    """SCM generator for the 4-variable CAREFL benchmark."""
 43
 44    seed: int = 42
 45
 46    def draw_latents(self, n: int, rng: np.random.Generator) -> dict[str, np.ndarray]:
 47        """Draw the latent noise of every variable.
 48
 49        Parameters
 50        ----------
 51        n : int
 52            Number of rows to draw.
 53        rng : np.random.Generator
 54            Random source.
 55
 56        Returns
 57        -------
 58        dict[str, np.ndarray]
 59            One array of length ``n`` per variable.
 60        """
 61        return {
 62            k: rng.laplace(loc=0.0, scale=_SCALE, size=n)
 63            for k in ["x1", "x2", "x3", "x4"]
 64        }
 65
 66    def simulate(
 67        self,
 68        n: int | None = None,
 69        *,
 70        rng: np.random.Generator | None = None,
 71        do: dict[str, float] | None = None,
 72        latents: dict[str, np.ndarray] | None = None,
 73    ) -> pd.DataFrame:
 74        """Simulate the SCM, with optional interventions and reused latents.
 75
 76        Parameters
 77        ----------
 78        n : int | None, optional
 79            Number of rows, by default ``None``. Then ``latents`` sets the count.
 80        rng : np.random.Generator | None, optional
 81            Random source, by default ``None``.
 82        do : dict[str, float] | None, optional
 83            Variables to hold at a fixed value, by default ``None``.
 84        latents : dict[str, np.ndarray] | None, optional
 85            Latent values to reuse, by default ``None``. Then they are drawn fresh.
 86
 87        Returns
 88        -------
 89        pd.DataFrame
 90            One column per variable.
 91        """
 92        do = do or {}
 93        if latents is None:
 94            if n is None:
 95                raise ValueError("provide either n or latents")
 96            rng = rng or np.random.default_rng(self.seed)
 97            latents = self.draw_latents(n, rng)
 98        n = len(latents["x1"])
 99
100        def clamp_or(name, value):
101            return np.full(n, float(do[name])) if name in do else value
102
103        x1 = clamp_or("x1", latents["x1"])
104        x2 = clamp_or("x2", latents["x2"])
105        x3 = clamp_or("x3", x1 + 0.5 * x2**3 + latents["x3"])
106        x4 = clamp_or("x4", -x2 + 0.5 * x1**2 + latents["x4"])
107        return pd.DataFrame({"x1": x1, "x2": x2, "x3": x3, "x4": x4})
108
109    # ----------------------------------------------------------------- datasets
110    def observational(self, n: int, seed_offset: int = 0) -> pd.DataFrame:
111        """Draw an observational sample.
112
113        Parameters
114        ----------
115        n : int
116            Number of rows.
117        seed_offset : int, optional
118            Added to the generator seed, by default ``0``.
119
120        Returns
121        -------
122        pd.DataFrame
123            The sample.
124        """
125        rng = np.random.default_rng(self.seed + 1 + seed_offset)
126        return self.simulate(n, rng=rng)
127
128    # -------------------------------------------------------------- ground truth
129    @staticmethod
130    def abduct_noise(obs: dict[str, float] | pd.DataFrame) -> dict[str, np.ndarray]:
131        """Exact noise values consistent with an observation (vectorized)."""
132        x1, x2 = np.asarray(obs["x1"], float), np.asarray(obs["x2"], float)
133        x3, x4 = np.asarray(obs["x3"], float), np.asarray(obs["x4"], float)
134        return {
135            "x1": x1,
136            "x2": x2,
137            "x3": x3 - x1 - 0.5 * x2**3,
138            "x4": x4 + x2 - 0.5 * x1**2,
139        }
140
141    def true_counterfactual(
142        self, obs: dict[str, float], do: dict[str, float]
143    ) -> dict[str, float]:
144        """Analytic counterfactual of a single observation under ``do``."""
145        eps = self.abduct_noise({k: np.atleast_1d(v) for k, v in obs.items()})
146        cf = self.simulate(do=do, latents=eps)
147        return {k: float(cf[k].iloc[0]) for k in cf}
148
149    def true_cf_curves(
150        self, obs: dict[str, float] = X_OBS, alphas: np.ndarray = ALPHA_GRID
151    ) -> dict:
152        """Compute the two analytic counterfactual curves of paper Fig. 6.
153
154        Returns
155        -------
156        dict
157            The observation, the intervention grid, and the counterfactual
158            values of x3 under do(x2) and of x4 under do(x1).
159        """
160        x3_cf = [self.true_counterfactual(obs, {"x2": a})["x3"] for a in alphas]
161        x4_cf = [self.true_counterfactual(obs, {"x1": a})["x4"] for a in alphas]
162        return {
163            "x_obs": dict(obs),
164            "alphas": [float(a) for a in alphas],
165            "x3_cf_do_x2": x3_cf,
166            "x4_cf_do_x1": x4_cf,
167        }