tramdag 0.2.0__py3-none-any.whl

tramdag/__init__.py

@@ -0,0 +1,22 @@
+"""tramdag — Interpretable Neural Causal Models (TRAM-DAGs) in PyTorch.
+
+One triangular normalizing flow (built on `zuko <https://zuko.readthedocs.io>`_)
+whose Jacobian sparsity is the causal DAG: fit once on observational data,
+then answer observational (L1), interventional (L2) and counterfactual (L3)
+queries. Reference: Sick & Dürr, *Interpretable Neural Causal Models with
+TRAM-DAGs*, CLeaR 2025 (arXiv:2503.16206).
+
+Conventional import alias::
+
+    import tramdag as td
+
+    flow = td.CausalFlowDAG(spec)
+    td.simulations.REGISTRY          # synthetic DGPs with known ground truth
+"""
+
+from . import simulations
+from .flow import CausalFlowDAG
+from .spec import ContinuousNode, OrdinalNode
+
+__all__ = ["CausalFlowDAG", "ContinuousNode", "OrdinalNode", "simulations"]
+__version__ = "0.2.0"

tramdag/conditioners.py

@@ -0,0 +1,74 @@
+"""Conditioner networks for the per-edge term types.
+
+Architectures replicate the original Keras/TF implementation's defaults
+(``tram_models.py`` in https://github.com/tensorchiefs/tram-dag)
+so that fitted models are directly comparable:
+
+- ``LinearShift``        — Linear(n, 1, bias=False)            (term "ls")
+- ``ComplexShift``       — 64-128-64 ReLU MLP -> 1, no bias    (term "cs",
+                            original ``ComplexShiftDefaultTabular``)
+- ``ComplexIntercept``   — 8-8 ReLU MLP -> n_params, no bias   (term "ci",
+                            original ``ComplexInterceptDefaultTabular``)
+- ``SimpleIntercept``    — free parameter vector (no parent dependence)
+
+Parent features follow the original implementation's encoding: continuous parents enter raw
+(one column), ordinal parents are one-hot encoded (``levels`` columns).
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+
+class SimpleIntercept(nn.Module):
+    """Free (data-independent) transform parameters, broadcast over the batch."""
+
+    def __init__(self, n_params: int):
+        super().__init__()
+        self.theta = nn.Parameter(torch.zeros(n_params))
+
+    def forward(self, n: int) -> Tensor:
+        return self.theta.unsqueeze(0).expand(n, -1)
+
+
+class ComplexIntercept(nn.Module):
+    """Transform parameters as a function of the (joint) ci-parent features."""
+
+    def __init__(self, n_features: int, n_params: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_features, 8), nn.ReLU(),
+            nn.Linear(8, 8), nn.ReLU(),
+            nn.Linear(8, n_params, bias=False),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x)
+
+
+class LinearShift(nn.Module):
+    def __init__(self, n_features: int):
+        super().__init__()
+        self.fc = nn.Linear(n_features, 1, bias=False)
+
+    @property
+    def weight(self) -> Tensor:
+        return self.fc.weight.squeeze(0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.fc(x).squeeze(-1)
+
+
+class ComplexShift(nn.Module):
+    def __init__(self, n_features: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_features, 64), nn.ReLU(),
+            nn.Linear(64, 128), nn.ReLU(),
+            nn.Linear(128, 64), nn.ReLU(),
+            nn.Linear(64, 1, bias=False),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x).squeeze(-1)

tramdag/flow.py

@@ -0,0 +1,427 @@
+"""CausalFlowDAG — a single triangular normalizing flow on a user-defined DAG.
+
+The flow maps iid standard-logistic latents ``U`` to the observed variables ``X``
+in topological order; its Jacobian sparsity is exactly the DAG adjacency. The
+joint log-likelihood decomposes per node, so one optimizer fits all nodes at once.
+
+Causal queries:
+    flow.sample(n)                    observational sampling
+    flow.sample(n, do={"T": 1})       interventional sampling (graph mutilation)
+    u = flow.abduct(df)               Pearl step 1 (latents from observations)
+    flow.sample(do={"T": 1}, u=u)     Pearl steps 2+3 (counterfactuals)
+    flow.pmf(df, node, do=...)        analytic per-row interventional PMF
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import Tensor, nn
+
+from .conditioners import ComplexIntercept, ComplexShift, LinearShift, SimpleIntercept
+from .spec import (ContinuousNode, NodeSpec, OrdinalNode, spec_from_dict,
+                   spec_to_dict, validate_and_sort)
+from .transforms import (StandardLogistic, make_univariate_transform,
+                         ordinal_abduct, ordinal_log_prob, ordinal_pmf,
+                         ordinal_sample)
+
+__all__ = ["CausalFlowDAG"]
+
+
+class _Node(nn.Module):
+    """One dimension of the flow: intercept (transform params) + additive shifts."""
+
+    def __init__(self, name: str, node: NodeSpec, spec: dict[str, NodeSpec]):
+        super().__init__()
+        self.name = name
+        self.kind = node.kind
+        self.parents = dict(node.parents)
+        self.ci_parents = [p for p, t in node.parents.items() if t == "ci"]
+
+        if isinstance(node, ContinuousNode):
+            self.ut = make_univariate_transform(node.transform, **node.transform_kwargs)
+            n_params = self.ut.n_params
+            self.levels = None
+        else:
+            self.ut = None
+            self.levels = node.levels
+            n_params = node.levels - 1
+
+        def width(parent: str) -> int:
+            pn = spec[parent]
+            return pn.levels if isinstance(pn, OrdinalNode) else 1
+
+        if self.ci_parents:
+            self.intercept = ComplexIntercept(sum(width(p) for p in self.ci_parents), n_params)
+        else:
+            self.intercept = SimpleIntercept(n_params)
+
+        self.shifts = nn.ModuleDict()
+        for parent, term in node.parents.items():
+            if term == "ls":
+                self.shifts[parent] = LinearShift(width(parent))
+            elif term == "cs":
+                self.shifts[parent] = ComplexShift(width(parent))
+
+    def theta_shift(self, feats: dict[str, Tensor], n: int) -> tuple[Tensor, Tensor]:
+        """Transform parameters (n, P) and total shift (n,) from parent features."""
+        if self.ci_parents:
+            theta = self.intercept(torch.cat([feats[p] for p in self.ci_parents], dim=1))
+        else:
+            theta = self.intercept(n)
+        shift = torch.zeros(n, dtype=theta.dtype, device=theta.device)
+        for parent, module in self.shifts.items():
+            shift = shift + module(feats[parent])
+        return theta, shift
+
+
+class CausalFlowDAG(nn.Module):
+    """A causal normalizing flow defined by ``spec = {name: NodeSpec}``."""
+
+    def __init__(self, spec: dict[str, NodeSpec], device: str = "cpu"):
+        super().__init__()
+        self.spec = spec
+        self.order = validate_and_sort(spec)
+        self.nodes = nn.ModuleDict({name: _Node(name, spec[name], spec) for name in self.order})
+        self.device = torch.device(device)
+        self.history: dict = {"train": [], "val": [], "lr": [], "time": []}
+        self.to(self.device)
+
+    # ------------------------------------------------------------------ data
+    def _encode_parent(self, name: str, values: Tensor) -> Tensor:
+        """Encode a node's values for use as a parent feature (original TRAM-DAG convention:
+        continuous raw (n, 1); ordinal one-hot (n, levels))."""
+        node = self.spec[name]
+        if isinstance(node, OrdinalNode):
+            return torch.nn.functional.one_hot(
+                values.long(), num_classes=node.levels).to(values.dtype)
+        return values.view(-1, 1)
+
+    def _tensorize(self, df: pd.DataFrame) -> dict[str, Tensor]:
+        out = {}
+        for name in self.order:
+            vals = torch.as_tensor(
+                df[name].to_numpy(dtype=np.float32), device=self.device)
+            out[name] = vals
+        return out
+
+    def _features(self, values: dict[str, Tensor]) -> dict[str, Tensor]:
+        return {name: self._encode_parent(name, vals) for name, vals in values.items()}
+
+    # ------------------------------------------------------------- likelihood
+    def node_log_prob(self, values: dict[str, Tensor],
+                      nodes: list[str] | None = None) -> dict[str, Tensor]:
+        """Per-node log-likelihood contributions, each (n,).
+
+        ``nodes`` restricts computation to a subset (used to skip frozen nodes
+        during training — valid because the per-node losses are independent)."""
+        feats = self._features(values)
+        n = next(iter(values.values())).shape[0]
+        out = {}
+        for name in (self.order if nodes is None else nodes):
+            node = self.nodes[name]
+            theta, shift = node.theta_shift(feats, n)
+            x = values[name]
+            if node.kind == "continuous":
+                z0, ladj = node.ut.forward(theta, x)
+                z = z0 + shift
+                out[name] = StandardLogistic.log_prob(z) + ladj
+            else:
+                out[name] = ordinal_log_prob(theta, shift, x)
+        return out
+
+    def log_prob(self, df: pd.DataFrame) -> Tensor:
+        """Joint log-likelihood log p(x) per row, shape (n,)."""
+        per_node = self.node_log_prob(self._tensorize(df))
+        return torch.stack(list(per_node.values()), dim=0).sum(dim=0)
+
+    def nll(self, df: pd.DataFrame) -> dict[str, float]:
+        """Mean negative log-likelihood per node (diagnostic)."""
+        with torch.no_grad():
+            per_node = self.node_log_prob(self._tensorize(df))
+        return {k: float(-v.mean()) for k, v in per_node.items()}
+
+    # ------------------------------------------------------------------- fit
+    def _set_ranges(self, train_df: pd.DataFrame) -> None:
+        """Train 5%/95% quantiles -> transform domain (the original implementation's min_max scaling)."""
+        for name in self.order:
+            node = self.nodes[name]
+            if node.kind == "continuous" and not node.ut._fitted:
+                q = train_df[name].quantile([0.05, 0.95])
+                node.ut.set_range(q.iloc[0], q.iloc[1])
+
+    def fit(self, train_df: pd.DataFrame, val_df: pd.DataFrame | None = None,
+            epochs: int = 500, learning_rate: float = 1e-2, batch_size: int = 512,
+            verbose: int = 50, seed: int | None = None,
+            restore_best: bool = False, schedule: str | None = None,
+            plateau_patience: int = 15, freeze_patience: int | None = None,
+            min_delta: float = 1e-4) -> "CausalFlowDAG":
+        """Jointly fit all nodes by maximum likelihood.
+
+        By default training keeps the **final** (converged) weights, so an
+        all-``ls`` model trained to convergence reproduces the classical maximum
+        likelihood estimate exactly (e.g. matches ``statsmodels``/``polr``).
+
+        The optimizer holds one parameter group per node. Because the joint NLL
+        decomposes per node with independent gradients, per-node learning rates
+        and freezing are exactly equivalent to independent per-node training.
+
+        Args:
+            val_df: optional held-out set, used only for monitoring (and for
+                ``restore_best``, ``schedule="plateau"`` and ``freeze_patience``).
+                If omitted, the training set is used for the validation metric.
+            restore_best: if True, snapshot each node's best-validation weights
+                during training and restore them at the end (mild early-stopping
+                regularization, the original implementation's convention). This makes the fit
+                *not* the training-data MLE, so leave it False for an exact
+                classical comparison. Default False.
+            schedule: learning-rate schedule. ``None`` = constant (the classic
+                behavior); ``"onecycle"`` = ``OneCycleLR`` (warmup to
+                ``learning_rate``, then anneal; stepped per batch);
+                ``"cosine"`` = ``CosineAnnealingLR`` over ``epochs``;
+                ``"plateau"`` = **per-node** decay: a node's lr is multiplied by
+                0.3 whenever its own validation NLL hasn't improved by
+                ``min_delta`` for ``plateau_patience`` epochs (floor 1e-3 ×
+                ``learning_rate``).
+            freeze_patience: if set, a node whose validation NLL hasn't improved
+                by ``min_delta`` for this many epochs is **frozen** — excluded
+                from the loss and backward pass (a real compute saving, since
+                per-node losses are independent). When every node is frozen the
+                fit returns early. Freeze epochs are recorded in
+                ``history["frozen"]``.
+
+        Calling ``fit`` again continues training (e.g. a second phase with a
+        lower learning rate); freezing state does not carry across calls.
+        """
+        if schedule not in (None, "onecycle", "cosine", "plateau"):
+            raise ValueError(f"unknown schedule {schedule!r}")
+        if seed is not None:
+            torch.manual_seed(seed)
+        self._set_ranges(train_df)
+
+        train_vals = self._tensorize(train_df)
+        val_vals = self._tensorize(val_df) if val_df is not None else train_vals
+        n = len(train_df)
+        steps_per_epoch = (n + batch_size - 1) // batch_size
+
+        opt = torch.optim.Adam(
+            [{"params": list(self.nodes[name].parameters()), "lr": learning_rate,
+              "node": name} for name in self.order])
+        sched = None
+        if schedule == "onecycle":
+            sched = torch.optim.lr_scheduler.OneCycleLR(
+                opt, max_lr=learning_rate, total_steps=epochs * steps_per_epoch)
+        elif schedule == "cosine":
+            sched = torch.optim.lr_scheduler.CosineAnnealingLR(
+                opt, T_max=epochs, eta_min=learning_rate * 1e-3)
+
+        if restore_best and not hasattr(self, "_best"):
+            self._best = {name: (float("inf"), None) for name in self.order}
+        best = self._best if restore_best else None
+        # per-node plateau/freeze bookkeeping (local to this fit call)
+        node_best = {name: float("inf") for name in self.order}
+        node_bad = {name: 0 for name in self.order}
+        frozen: set[str] = set()
+        t0 = time.perf_counter()
+        t_offset = self.history["time"][-1] if self.history.get("time") else 0.0
+        prev_train: dict[str, float] = {}
+
+        for epoch in range(epochs):
+            self.train()
+            active = [name for name in self.order if name not in frozen]
+            perm = torch.randperm(n, device=self.device)
+            train_acc = {name: prev_train.get(name, float("nan"))
+                         for name in frozen}
+            train_acc.update({name: 0.0 for name in active})
+            for start in range(0, n, batch_size):
+                idx = perm[start:start + batch_size]
+                batch = {k: v[idx] for k, v in train_vals.items()}
+                per_node = self.node_log_prob(batch, nodes=active)
+                node_nlls = {k: -v.mean() for k, v in per_node.items()}
+                loss = torch.stack(list(node_nlls.values())).sum()
+                opt.zero_grad()
+                loss.backward()
+                opt.step()
+                if schedule == "onecycle":
+                    sched.step()
+                w = len(idx) / n
+                for k, v in node_nlls.items():
+                    train_acc[k] += float(v.detach()) * w
+            if schedule == "cosine":
+                sched.step()
+            prev_train = train_acc
+
+            self.eval()
+            with torch.no_grad():
+                val_per_node = {k: float(-v.mean())
+                                for k, v in self.node_log_prob(val_vals).items()}
+            self.history["train"].append(train_acc)
+            self.history["val"].append(val_per_node)
+            self.history.setdefault("lr", []).append(
+                max(g["lr"] for g in opt.param_groups))
+            self.history.setdefault("time", []).append(
+                t_offset + time.perf_counter() - t0)
+
+            # per-node improvement tracking (plateau decay + freezing)
+            for g in opt.param_groups:
+                name = g["node"]
+                if name in frozen:
+                    continue
+                if val_per_node[name] < node_best[name] - min_delta:
+                    node_best[name] = val_per_node[name]
+                    node_bad[name] = 0
+                else:
+                    node_bad[name] += 1
+                if (schedule == "plateau" and node_bad[name] > 0
+                        and node_bad[name] % plateau_patience == 0):
+                    g["lr"] = max(g["lr"] * 0.3, learning_rate * 1e-3)
+                # under "plateau", only freeze nodes whose lr has already been
+                # decayed substantially — otherwise a node can freeze while a
+                # smaller lr would still make progress toward the optimum
+                lr_decayed = (schedule != "plateau"
+                              or g["lr"] <= learning_rate * 1e-2 * (1 + 1e-9))
+                if (freeze_patience is not None and lr_decayed
+                        and node_bad[name] >= freeze_patience):
+                    frozen.add(name)
+                    self.history.setdefault("frozen", {}).setdefault(
+                        name, len(self.history["val"]))  # 1-based global epoch
+
+            if restore_best:
+                for name in self.order:
+                    if val_per_node[name] < best[name][0]:
+                        best[name] = (val_per_node[name],
+                                      copy.deepcopy(self.nodes[name].state_dict()))
+
+            if verbose and (epoch % verbose == 0 or epoch == epochs - 1):
+                tot_t = sum(train_acc.values())
+                tot_v = sum(val_per_node.values())
+                print(f"[epoch {epoch + 1:5d}/{epochs}] train NLL {tot_t:.4f}  "
+                      f"val NLL {tot_v:.4f}"
+                      + (f"  frozen {sorted(frozen)}" if frozen else ""))
+
+            if len(frozen) == len(self.order):  # everything converged
+                if verbose:
+                    print(f"[epoch {epoch + 1:5d}] all nodes frozen — stopping.")
+                break
+
+        if restore_best:  # restore per-node best-validation weights
+            for name, (_, state) in best.items():
+                if state is not None:
+                    self.nodes[name].load_state_dict(state)
+        self.eval()
+        return self
+
+    # ------------------------------------------------------- causal queries
+    @torch.no_grad()
+    def sample(self, n: int | None = None, *, do: dict[str, float] | None = None,
+               u: pd.DataFrame | None = None, seed: int | None = None) -> pd.DataFrame:
+        """Sample from the (optionally mutilated) flow.
+
+        Args:
+            n: number of samples (ignored if ``u`` is given).
+            do: interventions {node: value}; intervened nodes are clamped and
+                their parent dependence removed (graph mutilation).
+            u: latent variables (as returned by :meth:`abduct`). If given, they
+                are pushed through the flow — together with ``do`` this yields
+                counterfactuals (Pearl's abduction -> action -> prediction).
+        """
+        do = do or {}
+        gen = None
+        if seed is not None:
+            gen = torch.Generator(device=self.device).manual_seed(seed)
+
+        if u is not None:
+            n = len(u)
+            u_vals = {name: torch.as_tensor(u[name].to_numpy(dtype=np.float32, copy=True),
+                                            device=self.device) for name in self.order}
+        elif n is not None:
+            u_vals = {name: StandardLogistic.sample((n,), device=self.device)
+                      if gen is None else
+                      StandardLogistic.icdf(torch.rand((n,), device=self.device, generator=gen))
+                      for name in self.order}
+        else:
+            raise ValueError("Provide either n or u.")
+
+        values: dict[str, Tensor] = {}
+        for name in self.order:
+            if name in do:
+                values[name] = torch.full((n,), float(do[name]), device=self.device)
+                continue
+            node = self.nodes[name]
+            feats = self._features({p: values[p] for p in node.parents})
+            theta, shift = node.theta_shift(feats, n)
+            z = u_vals[name]
+            if node.kind == "continuous":
+                values[name] = node.ut.inverse(theta, z - shift)
+            else:
+                values[name] = ordinal_sample(theta, shift, z)
+        return pd.DataFrame({k: v.cpu().numpy() for k, v in values.items()})
+
+    @torch.no_grad()
+    def abduct(self, df: pd.DataFrame, seed: int | None = None) -> pd.DataFrame:
+        """Pearl abduction: recover the latent variables ``u`` from observations.
+
+        Continuous nodes are inverted exactly (``u = h(x) + shift``); for ordinal
+        nodes the latent is only interval-identified, so it is sampled from the
+        standard logistic truncated to the observed level's interval.
+        """
+        gen = None
+        if seed is not None:
+            gen = torch.Generator(device=self.device).manual_seed(seed)
+        values = self._tensorize(df)
+        feats = self._features(values)
+        n = len(df)
+        u = {}
+        for name in self.order:
+            node = self.nodes[name]
+            theta, shift = node.theta_shift(feats, n)
+            x = values[name]
+            if node.kind == "continuous":
+                z0, _ = node.ut.forward(theta, x)
+                u[name] = z0 + shift
+            else:
+                u[name] = ordinal_abduct(theta, shift, x, generator=gen)
+        return pd.DataFrame({k: v.cpu().numpy() for k, v in u.items()})
+
+    @torch.no_grad()
+    def pmf(self, df: pd.DataFrame, node: str, do: dict[str, float] | None = None) -> np.ndarray:
+        """Analytic class probabilities (n, levels) for an ordinal node, with the
+        node's parents taken from ``df`` after applying ``do`` overrides."""
+        if not isinstance(self.spec[node], OrdinalNode):
+            raise ValueError(f"pmf() requires an ordinal node, '{node}' is continuous.")
+        df_local = df.copy()
+        for col, val in (do or {}).items():
+            df_local[col] = val
+        nd = self.nodes[node]
+        values = {p: torch.as_tensor(df_local[p].to_numpy(dtype=np.float32),
+                                     device=self.device) for p in nd.parents}
+        feats = self._features(values)
+        theta, shift = nd.theta_shift(feats, len(df_local))
+        return ordinal_pmf(theta, shift).cpu().numpy()
+
+    # ------------------------------------------------------------------- io
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"spec": spec_to_dict(self.spec),
+                    "state_dict": self.state_dict(),
+                    "history": self.history}, path)
+
+    @classmethod
+    def load(cls, path: str | Path, device: str = "cpu") -> "CausalFlowDAG":
+        ckpt = torch.load(path, map_location=device, weights_only=False)
+        flow = cls(spec_from_dict(ckpt["spec"]), device=device)
+        for name in flow.order:  # mark transforms as fitted before loading buffers
+            node = flow.nodes[name]
+            if node.kind == "continuous":
+                node.ut._fitted = True
+        flow.load_state_dict(ckpt["state_dict"])
+        flow.history = ckpt.get("history", {"train": [], "val": []})
+        flow.eval()
+        return flow

tramdag/simulations/__init__.py

@@ -0,0 +1,23 @@
+"""Synthetic-cohort generators for tramdag.
+
+Each scenario is one module exposing a numpy-only SCM generator class with known
+causal ground truth. New scenarios register here so experiments/tests can look
+them up by name. Frozen CSVs live under ``data/<name>/`` and are a contract —
+regenerate only deliberately via each module's CLI.
+"""
+
+from .carefl import Carefl4
+from .magic_mrclean import MagicMrClean
+from .triangle import TriangleContinuous, TriangleMixed
+from .vaca import VacaTriangle
+
+REGISTRY = {
+    "magic-mrclean": MagicMrClean,
+    "triangle": TriangleContinuous,
+    "triangle-mixed": TriangleMixed,
+    "vaca": VacaTriangle,
+    "carefl": Carefl4,
+}
+
+__all__ = ["MagicMrClean", "TriangleContinuous", "TriangleMixed",
+           "VacaTriangle", "Carefl4", "REGISTRY"]

tramdag/simulations/carefl.py

@@ -0,0 +1,125 @@
+"""The CAREFL counterfactual benchmark DGP (TRAM-DAG paper App. C.2).
+
+Originally from Khemakhem et al. (2021, CAREFL Fig. 5), used in the paper to
+benchmark TRAM-DAG's L3 (counterfactual) queries (Fig. 6)::
+
+    x1, x2 ~ Laplace(0, 1/sqrt(2))
+    x3 = x1 + 0.5 x2^3 + Laplace(0, 1/sqrt(2))
+    x4 = -x2 + 0.5 x1^2 + Laplace(0, 1/sqrt(2))
+
+All-continuous additive-noise SCM, so individual counterfactuals are **analytic**
+via noise abduction (no Monte Carlo): eps3 = x3 - x1 - 0.5 x2^3 and
+eps4 = x4 + x2 - 0.5 x1^2 are recovered exactly, then the mutilated SCM is
+re-evaluated. The paper picks the observation ``X_OBS = (2.00, 1.50, 0.81, -0.28)``
+and sweeps two queries for alpha in [-3, 3]:
+
+    (i)  x3^cf  given do(x2 = alpha):  x1 + 0.5 alpha^3 + eps3
+    (ii) x4^cf  given do(x1 = alpha): -x2 + 0.5 alpha^2 + eps4
+
+CLI::
+
+    uv run python -m tramdag.simulations.carefl --out data/carefl --seed 42
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+X_OBS = {"x1": 2.00, "x2": 1.50, "x3": 0.81, "x4": -0.28}  # the paper's observation
+ALPHA_GRID = np.round(np.linspace(-3.0, 3.0, 61), 4)
+_SCALE = 1.0 / np.sqrt(2.0)
+
+
+@dataclass
+class Carefl4:
+    """SCM generator for the 4-variable CAREFL benchmark."""
+
+    seed: int = 42
+
+    def draw_latents(self, n: int, rng: np.random.Generator) -> dict[str, np.ndarray]:
+        return {k: rng.laplace(loc=0.0, scale=_SCALE, size=n)
+                for k in ["x1", "x2", "x3", "x4"]}
+
+    def simulate(self, n: int | None = None, *, rng: np.random.Generator | None = None,
+                 do: dict[str, float] | None = None,
+                 latents: dict[str, np.ndarray] | None = None) -> pd.DataFrame:
+        do = do or {}
+        if latents is None:
+            if n is None:
+                raise ValueError("provide either n or latents")
+            rng = rng or np.random.default_rng(self.seed)
+            latents = self.draw_latents(n, rng)
+        n = len(latents["x1"])
+
+        def clamp_or(name, value):
+            return np.full(n, float(do[name])) if name in do else value
+
+        x1 = clamp_or("x1", latents["x1"])
+        x2 = clamp_or("x2", latents["x2"])
+        x3 = clamp_or("x3", x1 + 0.5 * x2**3 + latents["x3"])
+        x4 = clamp_or("x4", -x2 + 0.5 * x1**2 + latents["x4"])
+        return pd.DataFrame({"x1": x1, "x2": x2, "x3": x3, "x4": x4})
+
+    # ----------------------------------------------------------------- datasets
+    def observational(self, n: int, seed_offset: int = 0) -> pd.DataFrame:
+        rng = np.random.default_rng(self.seed + 1 + seed_offset)
+        return self.simulate(n, rng=rng)
+
+    # -------------------------------------------------------------- ground truth
+    @staticmethod
+    def abduct_noise(obs: dict[str, float] | pd.DataFrame) -> dict[str, np.ndarray]:
+        """Exact noise values consistent with an observation (vectorized)."""
+        x1, x2 = np.asarray(obs["x1"], float), np.asarray(obs["x2"], float)
+        x3, x4 = np.asarray(obs["x3"], float), np.asarray(obs["x4"], float)
+        return {"x1": x1, "x2": x2,
+                "x3": x3 - x1 - 0.5 * x2**3,
+                "x4": x4 + x2 - 0.5 * x1**2}
+
+    def true_counterfactual(self, obs: dict[str, float],
+                            do: dict[str, float]) -> dict[str, float]:
+        """Analytic counterfactual of a single observation under ``do``."""
+        eps = self.abduct_noise({k: np.atleast_1d(v) for k, v in obs.items()})
+        cf = self.simulate(do=do, latents=eps)
+        return {k: float(cf[k].iloc[0]) for k in cf}
+
+    def true_cf_curves(self, obs: dict[str, float] = X_OBS,
+                       alphas: np.ndarray = ALPHA_GRID) -> dict:
+        """The paper's two Fig.-6 curves, analytic."""
+        x3_cf = [self.true_counterfactual(obs, {"x2": a})["x3"] for a in alphas]
+        x4_cf = [self.true_counterfactual(obs, {"x1": a})["x4"] for a in alphas]
+        return {"x_obs": dict(obs), "alphas": [float(a) for a in alphas],
+                "x3_cf_do_x2": x3_cf, "x4_cf_do_x1": x4_cf}
+
+
+# --------------------------------------------------------------------------- CLI
+def main(argv: list[str] | None = None) -> None:
+    p = argparse.ArgumentParser(description="Generate the CAREFL benchmark data.")
+    p.add_argument("--out", type=Path, default=Path("data/carefl"))
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--n-obs", type=int, default=5000)
+    args = p.parse_args(argv)
+
+    gen = Carefl4(seed=args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+    obs = gen.observational(args.n_obs)
+    obs.to_csv(args.out / "obs.csv", index=False)
+
+    truth = {"source": "arXiv:2503.16206 App. C.2 (orig. Khemakhem 2021 Fig. 5)",
+             "seed": args.seed, "n_obs": args.n_obs,
+             "scm": {"x1": "Laplace(0, 1/sqrt(2))", "x2": "Laplace(0, 1/sqrt(2))",
+                     "x3": "x1 + 0.5*x2^3 + Laplace", "x4": "-x2 + 0.5*x1^2 + Laplace"},
+             **gen.true_cf_curves()}
+    (args.out / "truth.json").write_text(json.dumps(truth, indent=2) + "\n")
+    print(f"[carefl] n={len(obs)}  cf sanity: "
+          f"x3_cf(do x2=1.5)={gen.true_counterfactual(X_OBS, {'x2': 1.5})['x3']:.3f} "
+          f"(factual {X_OBS['x3']})")
+
+
+if __name__ == "__main__":
+    main()