tramdag 0.2.0__tar.gz

tramdag-0.2.0/.claude/scheduled_tasks.lock

@@ -0,0 +1 @@
+{"sessionId":"ea0171fe-94d9-4044-b1c8-81400b3676ad","pid":50879,"procStart":"Thu Jun 11 09:52:35 2026","acquiredAt":1781179094640}

tramdag-0.2.0/.claude/settings.local.json

@@ -0,0 +1,26 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(git add *)",
+      "Bash(git commit -m ' *)",
+      "Bash(git checkout *)",
+      "Bash(uv run *)",
+      "Bash(uv venv *)",
+      "Bash(uv pip *)",
+      "Bash(/tmp/zdtest/bin/python -c ' *)",
+      "Bash(MPLBACKEND=Agg uv run python notebooks/demo_tram_dag_colab.py)",
+      "Bash(uvx jupytext *)",
+      "Bash(git check-ignore *)",
+      "Bash(git commit -q -m 'fit\\(\\): lr schedules + per-node freezing \\(defaults unchanged\\) *)",
+      "Bash(git commit -q -m 'Training-speed benchmark + report *)",
+      "Bash(git commit -q -m 'Colab GPU demo \\(bimodal VACA benchmark\\) + README badge *)",
+      "Bash(git push *)",
+      "Bash(gh auth *)",
+      "Bash(gh api *)",
+      "Bash(git commit -q -m 'Address review: option guide in report + notebooks README *)",
+      "Bash(gh pr *)",
+      "Bash(git pull *)",
+      "Bash(git branch *)"
+    ]
+  }
+}

tramdag-0.2.0/.gitignore

@@ -0,0 +1,24 @@
+# Python / uv
+.venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+*.egg-info/
+
+# Experiment outputs — regenerable, and (for the clinical 'magic' source) derived
+# from patient data, so never committed. The synthetic data/ folder IS tracked.
+results/
+
+# R
+.Rhistory
+.RData
+
+# Notebooks: jupytext py:percent files are the source of truth;
+# generated .ipynb (with embedded image outputs) stay out of git.
+notebooks/*.ipynb
+# exception: the Colab demo needs a tracked (output-stripped) ipynb for the
+# "Open in Colab" badge; regenerate with `uvx jupytext --to ipynb <demo>.py`
+!notebooks/demo_tram_dag_colab.ipynb
+
+# build artifacts
+dist/

tramdag-0.2.0/CHANGELOG.md

@@ -0,0 +1,93 @@
+# Changelog
+
+## 0.2.0 (2026-06-12)
+
+First PyPI release: `pip install tramdag`.
+
+### Changed (naming & packaging)
+
+- **Renamed**: Python package `zuko_dag` → **`tramdag`** (conventional alias
+  `import tramdag as td`); GitHub repo `tram-dag-zuko` → `tensorchiefs/tramdag`
+  (old URLs redirect). The package implements TRAM-DAGs; zuko names the backend.
+  No API changes; old checkpoints still load. References to the original
+  Keras/TF implementation (tensorchiefs/tram-dag) reworded to avoid
+  self-reference.
+- **MIT license** added; PyPI metadata (authors, urls, classifiers); runtime
+  dependencies trimmed to `torch`, `zuko`, `numpy`, `pandas` (pytest/scipy/
+  statsmodels/scikit-learn/matplotlib moved to the `dev` dependency group).
+- **README rewritten method-first**: the repo is the reference implementation of
+  the CLeaR 2025 paper (arXiv:2503.16206); the stroke analysis is the case study
+  (arXiv:2606.12623) with its detail moved to `docs/stroke-case-study.md`.
+  Citation BibTeX added for both papers.
+
+### Added
+
+- **`fit(schedule=..., freeze_patience=...)`** — learning-rate schedules and
+  per-node early stopping (defaults unchanged). The optimizer now holds one
+  param group per node; `schedule="plateau"` decays each node's lr off its own
+  validation NLL, and `freeze_patience` drops converged nodes from the loss
+  (real FLOP savings — per-node gradients are independent) with early exit when
+  all nodes froze. Also `"onecycle"`/`"cosine"`. Benchmarks + recommendation in
+  `docs/training-speed.md` (`experiments/bench_training.py`): plateau+freeze
+  matches the hand-tuned two-phase recipe's time-to-accuracy with **no budget
+  tuning and ~3× less total compute**; full-batch LBFGS solves the classical
+  all-`ls` MLE in <2 s (2/3 seeds). Existing defaults intentionally untouched.
+- **Colab demo** `notebooks/demo_tram_dag_colab.py` (+ tracked output-stripped
+  `.ipynb` for the badge): the paper's bimodal VACA benchmark fitted live
+  (cuda/cpu auto-detect), L1 pairs plot, analytic do-checks, per-individual
+  counterfactuals vs DGP truth, GPU-vs-CPU race.
+
+- **The TRAM-DAG paper's DGPs** (Sick & Dürr, CLeaR 2025, arXiv:2503.16206) as
+  simulation registry families, each a numpy-only SCM with known/analytic ground
+  truth + frozen n=5000 CSVs (`data/<name>/`, the test contract) and CLIs:
+  - `simulations/triangle.py` — `TriangleContinuous` (§6.1: logistic-latent TRAM
+    DGP, h₂=5x₂+2x₁, h₃=0.63x₃−0.2x₁−f(x₂)) and `TriangleMixed` (§6.2: ordinal x₃,
+    θ=(−2, 0.42, 1.02)); f variants `linear`/`cubic`/`exp`/`atan`/`sin`; supports
+    array-valued `do` (C.4 soft interventions).
+  - `simulations/vaca.py` — `VacaTriangle` (App. C.1 bimodal Gaussian L1/L2
+    benchmark vs CNF).
+  - `simulations/carefl.py` — `Carefl4` (App. C.2 Laplace SCM; **analytic**
+    counterfactuals via `abduct_noise`/`true_counterfactual`).
+- `experiments/paper_{triangle,triangle_mixed,vaca,carefl}.py` (+ `paper_common.py`)
+  — replicate the paper's figures: coefficient trajectories (Fig. 14/15/19), CS-curve
+  recovery (Fig. 7), L1/L2 distribution overlays (Fig. 4/5/9/16/20), counterfactual
+  curves at the paper's x_obs (Fig. 6), and the C.4 odds-ratio check (OR ≈ 7.4).
+- `tests/test_paper_dgps.py` — generator pinning (KS TRAM-identities, frozen-CSV
+  contract, analytic ground truth) + flow recovery (coefficients with the ordinal
+  sign-flip, CS curve, VACA do-moments, CAREFL counterfactual MAE).
+
+### Changed (behavior)
+
+- **`CausalFlowDAG.fit(..., restore_best=False)` is now the default.** Training keeps
+  the **final converged weights** instead of restoring per-node best-validation
+  weights. Rationale:
+  - *Least surprise* — `fit()` returns the model you trained, not a silently
+    swapped earlier epoch.
+  - *Exact classical comparison* — an all-`ls` model trained to convergence is now
+    exactly the maximum-likelihood (proportional-odds) estimate, matching
+    `statsmodels` `OrderedModel` and R `MASS::polr` to ~1e-3 (see
+    `experiments/validate_ls.py`, `tests/test_simulations.py::test_all_ls_flow_is_exact_mle`).
+    This was **not achievable before**: best-validation restoration pinned the fit
+    off the training optimum.
+  - Early stopping is now an explicit, opt-in regularization choice.
+
+  To restore the previous behavior, pass `restore_best=True`.
+
+  **Note for flexible (`ci`/`cs`) models:** their MLE *overfits the observational
+  confounding*, so they need `restore_best=True` to recover the causal effect (lower
+  validation NLL confirms it generalizes better). `experiments/run_experiment`
+  therefore defaults `restore_best` per style — off for all-`ls`, on for flexible.
+
+### Added
+
+- `src/tramdag/simulations/magic_mrclean.py` — synthetic stroke cohort (SCM with
+  known ground truth); `ls`/`nl` variants; CLI to (re)generate frozen CSVs.
+- `data/magic-mrclean/` — frozen public CSVs + `fit_ls.R` classical R reference and
+  committed `ref_ls/` outputs. The public, reproducible substitute for the private
+  clinical data.
+- `experiments/common.py::load_data(source)` — switch between `"magic"` (private) and
+  `"magic-mrclean/{ls,nl}"` (synthetic, default).
+- `experiments/sim_flow.py` — known-truth recovery storyline; `validate_ls.py`
+  rewritten as a spot-on flow-vs-MLE-vs-R comparison.
+- `tests/test_simulations.py` — generator, known-truth recovery, the all-`ls`
+  spot-on MLE check, and the Python-vs-R regression.

tramdag-0.2.0/CLAUDE.md

@@ -0,0 +1,113 @@
+# CLAUDE.md — working context for tramdag
+
+## What this is
+
+A causal normalizing-flow implementation of **TRAM-DAG** (transformation models on a
+DAG) built on [zuko](https://zuko.readthedocs.io/stable/). One triangular flow from iid
+standard-logistic latents to the observed variables; Jacobian sparsity = the DAG.
+Supports the do-operator, Pearl abduction (counterfactuals), analytic interventional
+PMFs, and per-node configurable monotone transforms (Bernstein / RQ-spline / affine).
+
+Origin: extracted from the private `tensorchiefs/tram-dag-stroke` paper repo (as
+`zuko_dag`; renamed to `tramdag` in June 2026, repo `tensorchiefs/tramdag`). The paper analyzed the MAGIC stroke cohort against the MR CLEAN RCT;
+that **clinical data is NOT in this repo** and never should be. The synthetic
+`data/magic-mrclean/` cohort is the public stand-in (same schema, known ground truth).
+
+## Commands
+
+```bash
+uv sync                          # install (uv.lock pinned: zuko, torch, statsmodels, ...)
+uv run pytest tests/ -q          # full suite ~11 min; tests/test_flow.py alone ~20 s
+cd experiments
+uv run python sim_flow.py nl     # headline storyline (all-ls vs flexible vs known truth)
+uv run python validate_ls.py     # spot-on flow == statsmodels == R polr check
+uv run python paper_triangle.py atan cs   # TRAM-DAG paper replications (paper_*.py)
+```
+
+Experiments default to the synthetic data (`magic-mrclean/nl`). The `magic` source
+(private clinical data) only works inside the original paper monorepo.
+
+## Architecture (src/tramdag/)
+
+- `spec.py` — user-facing DAG spec: `{name: ContinuousNode|OrdinalNode}`, each node
+  declares `parents={parent: term}` with term ∈ `ls` (linear shift), `cs` (complex
+  shift MLP), `ci` (complex intercept — transform params from parents; multiple ci
+  parents feed ONE joint network).
+- `transforms.py` — monotone 1-D transforms wrapping zuko (`BernsteinUT`, `SplineUT`,
+  `AffineUT`; pre-scaled from train 5%/95% quantiles to [-5,5], expanding-bracket
+  bisection inverse) + the ordinal ordered-logit transform
+  (`P(Y<=k) = sigmoid(theta_k - shift)`, cutpoints `[t0, t0+cumsum(exp(...))]`).
+- `conditioners.py` — ls/cs/ci networks (widths replicate the original Keras/TF implementation).
+- `flow.py` — `CausalFlowDAG`: `fit`, `sample(n, do=, u=)`, `abduct`, `pmf`,
+  `log_prob`, `save/load`. NLL decomposes per node → one Adam fits all nodes jointly.
+- `simulations/` — numpy-only SCM generators with known ground truth, looked up via
+  `REGISTRY`; each module has a CLI that regenerates its frozen `data/<name>/` CSVs:
+  `magic_mrclean.py` (stroke SCM, `ls`/`nl`), `triangle.py` (paper §6 continuous +
+  ordinal triangles, f variants linear/cubic/exp/atan/sin), `vaca.py` (App. C.1
+  bimodal L1/L2 benchmark), `carefl.py` (App. C.2 Laplace SCM, **analytic**
+  counterfactuals).
+
+## Conventions that matter (easy to get wrong)
+
+- **Latent scale**: continuous `z = h(x) + shift` (shifts ADDED); ordinal
+  `P(Y<=k) = sigmoid(theta_k − shift)` (shift SUBTRACTED). Both follow the original TRAM-DAG
+  conventions; tests pin them.
+- **Parent encoding**: continuous parents enter RAW (no standardization); ordinal
+  parents one-hot (all levels). With cutpoints, only shift *differences* between
+  one-hot levels are identified — compare `w[k] − w[0]` against classical references.
+- **Ordinal log-prob is computed in log-space** (`logsigmoid` + stable `log1mexp`,
+  better-conditioned side chosen per element). The naive sigmoid difference saturates
+  in float32 → *exactly zero* gradients → a node can freeze at init forever. Do not
+  "simplify" it back.
+- **Seeding**: weight init happens at construction — call `torch.manual_seed` BEFORE
+  `CausalFlowDAG(spec)`, not just in `fit`.
+- **`fit(restore_best=False)` is the default** (keeps final converged weights = exact
+  MLE; an all-`ls` model then matches statsmodels/R-polr to ~1e-3). `restore_best=True`
+  = per-node best-validation restoration (early stopping). Key empirical finding:
+  **flexible (ci/cs) models overfit observational confounding at the MLE and need
+  `restore_best=True` to recover the causal effect; all-`ls` models don't.**
+  `run_experiment` defaults per style. See CHANGELOG.md.
+
+## Ground truth & reference numbers (seed 7 synthetic data)
+
+- `data/magic-mrclean/{ls,nl}/truth.json` — true ATE from the SCM: `ls` +0.132,
+  `nl` +0.104; naive confounded contrast +0.26/+0.30.
+- `nl` storyline: all-`ls` flow ≈ +0.076 (biased — can't extrapolate the age-fading
+  treatment effect to the younger RCT population), flexible flow ≈ +0.10 (recovers).
+- Spot-on check (`ls` variant, full-data, restore_best=False): flow = statsmodels =
+  R polr at Age 0.0526, NIHSSa 0.1630, T −0.9424; ATE +0.1429 vs +0.1428.
+- R reference: `data/magic-mrclean/fit_ls.R` (needs `tram`, `MASS`); its committed
+  `ref_ls/` outputs let tests run without R.
+- Original clinical-data numbers (context only, not reproducible here): TRAM-DAG
+  nihss6 +0.108, md_dag_ls +0.054, MR CLEAN RCT +0.135 [0.057, 0.213].
+- **Paper DGPs** (seed 42, arXiv:2503.16206): `triangle` true coefficients β12=+2,
+  β13=−0.2 (+0.3 on x2 for `linear`); a fitted `cs` learns −f(x2)+const.
+  `triangle-mixed` cutpoints θ=(−2, 0.42, 1.02); **ordinal sign flip**: the paper
+  ADDS the ordinal shift, the flow SUBTRACTS → fitted weights −0.2 / +0.3; the C.4
+  odds-ratio check gives OR ≈ e² ≈ 7.4. `vaca`: E[x3|do(x2=a)] = −0.25 + 0.25a
+  (do(x2=−3) is off-manifold extrapolation — looser tolerance). `carefl`:
+  counterfactuals are analytic (`Carefl4.true_counterfactual`); the paper's x_obs has
+  a ~4σ abducted noise, so tests score 300 typical rows instead of that single point.
+
+## Testing policy
+
+- Frozen CSVs in `data/` (`magic-mrclean`, `triangle*`, `vaca`, `carefl`) are a
+  contract — **never regenerate silently**; a new seed/equations → new folder
+  (sim2-style), regenerate `ref_ls/` with R where applicable, update
+  truth-dependent tests. `test_paper_dgps.py::test_frozen_csv_contract` pins the
+  paper-DGP CSVs to their generators bit-exactly.
+- Fit tests for the paper DGPs train on **regenerated n=20k** (deterministic
+  `observational(n, seed_offset=100)`), not the frozen n=5k — β13 multiplies the
+  low-variance x1 ∈ [0.25, 0.73] and is too weakly identified at n=5k.
+- New causal features should be validated against the simulator's known truth
+  (`MagicMrClean.true_ate`, `counterfactual_pair` gives true individual
+  counterfactuals via shared latents).
+
+## Roadmap notes
+
+- ~~Generalize `simulations/` registry beyond the stroke DAG~~ — done for the
+  TRAM-DAG paper's DGPs (triangle/triangle-mixed/vaca/carefl, June 2026). Still
+  open: hidden confounding à la DeCaFlow.
+- ~~Package for PyPI~~ — published as `tramdag` 0.2.0 (June 2026); release flow:
+  bump version in pyproject + `__init__`, `uv build`, `uv publish` (Oliver's
+  PyPI token), CHANGELOG section.

tramdag-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Oliver Dürr, Beate Sick (tensorchiefs)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tramdag-0.2.0/PKG-INFO

@@ -0,0 +1,206 @@
+Metadata-Version: 2.4
+Name: tramdag
+Version: 0.2.0
+Summary: Interpretable Neural Causal Models (TRAM-DAGs) in PyTorch: one causal normalizing flow for observational, interventional and counterfactual queries
+Project-URL: Homepage, https://github.com/tensorchiefs/tramdag
+Project-URL: Method paper (CLeaR 2025), https://arxiv.org/abs/2503.16206
+Project-URL: Stroke case study, https://arxiv.org/abs/2606.12623
+Author: Beate Sick
+Author-email: Oliver Dürr <oliver.duerr@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: causal-inference,counterfactuals,interpretability,normalizing-flows,structural-causal-models,transformation-models
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: torch>=2.0
+Requires-Dist: zuko>=1.3
+Description-Content-Type: text/markdown
+
+# tramdag — Interpretable Neural Causal Models (TRAM-DAGs) in PyTorch
+
+[![Open the demo in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tensorchiefs/tramdag/blob/main/notebooks/demo_tram_dag_colab.ipynb)
+[![PyPI](https://img.shields.io/pypi/v/tramdag)](https://pypi.org/project/tramdag/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+**TRAM-DAGs** model each variable of a structural causal model with a
+(transformation-model) flow: one triangular normalizing flow from iid
+standard-logistic latents to the observed variables, whose Jacobian sparsity is
+exactly your causal DAG. Fit it **once** on observational data and answer all
+three rungs of Pearl's causal hierarchy — observational (L1), interventional
+(L2, the do-operator), and counterfactual (L3, Pearl abduction) — while keeping
+**interpretable effects**: every linear-shift coefficient is a log-odds ratio,
+exactly as in classical proportional-odds models.
+
+> Beate Sick & Oliver Dürr, *Interpretable Neural Causal Models with TRAM-DAGs*,
+> CLeaR 2025 ([arXiv:2503.16206](https://arxiv.org/abs/2503.16206)).
+> This repo is the reference implementation (PyTorch, built on
+> [zuko](https://zuko.readthedocs.io/stable/)); all of the paper's experiments are
+> replicated here with pinned tests.
+
+**5-minute showcase**: the Colab badge above fits the paper's bimodal benchmark
+live (GPU-ready) and walks L1 → L2 → L3, every answer checked against analytic
+ground truth. Didactic walkthrough of the model:
+[`notebooks/intro_tram_dag.py`](notebooks/intro_tram_dag.py).
+
+## Install
+
+```bash
+pip install tramdag            # PyPI
+uv sync                        # or: dev setup from a clone (tests, experiments)
+```
+
+## 30 seconds of API
+
+```python
+import tramdag as td
+from tramdag import CausalFlowDAG, ContinuousNode, OrdinalNode
+
+spec = {                                  # the spec IS the labelled DAG
+    "Age":     ContinuousNode(),
+    "mRS_pre": OrdinalNode(levels=6, parents={"Age": "ci"}),
+    "NIHSSa":  ContinuousNode(parents={"Age": "ci", "mRS_pre": "ls"}),
+    "T":       OrdinalNode(levels=2,
+                           parents={"Age": "ci", "mRS_pre": "ls", "NIHSSa": "cs"}),
+    "mRS_3m":  OrdinalNode(levels=7,
+                           parents={"Age": "ci", "mRS_pre": "ls",
+                                    "NIHSSa": "cs", "T": "ls"}),
+}
+flow = CausalFlowDAG(spec)                # validates acyclicity, builds the flow
+
+# self-stopping training: per-node plateau lr decay + freezing of converged
+# nodes (exact, since the per-node NLLs have independent gradients);
+# see docs/training-speed.md for benchmarks and the classic two-phase recipe
+flow.fit(train_df, val_df, epochs=4000, learning_rate=1e-2,
+         schedule="plateau", plateau_patience=30, freeze_patience=120)
+
+flow.log_prob(df)                          # L1: joint log-likelihood per row
+flow.sample(1000)                          # L1: observational sampling
+flow.sample(1000, do={"T": 1})             # L2: interventional (graph mutilation)
+flow.pmf(df, node="mRS_3m", do={"T": 1})   # L2: analytic interventional PMF
+
+u  = flow.abduct(df)                       # L3 step 1: latents from observations
+cf = flow.sample(do={"T": 1}, u=u)         # L3 steps 2+3: counterfactuals
+
+flow.save("flow.pt"); flow = CausalFlowDAG.load("flow.pt")
+
+td.simulations.REGISTRY                    # synthetic DGPs with known ground truth
+```
+
+## The model in one table
+
+Per node, the transformation is additive on the latent (log-odds) scale —
+`u = h(x; θ) + Σ β·x_pa + Σ g(x_pa)` — and each parent edge declares how it enters:
+
+| edge term | meaning | interpretability |
+|---|---|---|
+| `ls` | linear shift `β·x_pa` | `exp(β)` is an odds ratio — one number per edge |
+| `cs` | complex shift `g(x_pa)` (MLP), still additive | odds-ratio *function*, plot `g` |
+| `ci` | complex intercept: the transform's parameters depend on the parents (several `ci` parents feed one joint network) | maximal flexibility, interactions |
+
+Continuous nodes carry a monotone 1-D transform (`bernstein` — TRAM-faithful
+default, `spline`, `affine`; `ContinuousNode(transform=..., transform_kwargs=...)`);
+ordinal nodes an ordered-logit head `P(x ≤ k) = σ(θ_k − shift)`. Abduction is exact
+for continuous nodes and truncated-logistic for ordinal ones, so
+`flow.sample(u=flow.abduct(df))` reproduces `df` exactly / level-exactly.
+
+## Validation (all pinned by tests)
+
+- **Paper replication** — every experiment of the CLeaR paper is a registry
+  family (numpy-only SCM + frozen CSVs + replication script):
+
+  | family | paper | demonstrates |
+  |---|---|---|
+  | `triangle` (`linear`,`atan`,`sin`) | §6.1 | LS coefficient recovery (β = 2, −0.2, +0.3), CS curve ≡ −f(x₂), non-monotone f |
+  | `triangle-mixed` (`linear`,`exp`) | §6.2 | mixed data L1/L2 + the C.4 odds-ratio check (OR ≈ 7.4) |
+  | `vaca` | §5.1–5.2 | the bimodal L1 case a default CNF misses; L2 `p(x₃ \| do(x₂))` |
+  | `carefl` | §5.3 | L3 counterfactual curves vs **analytic** truth |
+
+  ```bash
+  cd experiments && uv run python paper_triangle.py atan cs   # etc., see paper_*.py
+  ```
+
+  Sign note: ordinal shifts are *subtracted* here but *added* in the paper, so
+  fitted ordinal weights are the paper's with flipped sign (`truth.json` records
+  both conventions per family).
+
+- **Exact classical equivalence** — an all-`ls` flow trained to convergence *is*
+  the proportional-odds MLE: coefficients match `statsmodels` **and** R
+  `MASS::polr` to ~4 decimals (`experiments/validate_ls.py`, R reference committed
+  under `data/magic-mrclean/*/ref_ls/`).
+
+- **Training speed** — schedules, per-node freezing, LBFGS and device benchmarks:
+  [`docs/training-speed.md`](docs/training-speed.md).
+
+## Case study: individualized treatment effects in stroke
+
+The method's flagship application estimates individualized thrombectomy effects
+from the observational MAGIC cohort with external validation against the
+MR CLEAN trial:
+
+> Dürr, Herzog, Bühler, Wegener & Sick, *Estimating Individualized Treatment
+> Effects in Acute Ischemic Stroke with Causal Transformation Models (TRAM-DAG)*
+> ([arXiv:2606.12623](https://arxiv.org/abs/2606.12623)).
+
+The clinical data is private and **never** part of this repo. Its public
+stand-in is `data/magic-mrclean/` — a fully synthetic cohort with the same
+schema and **known ground truth** (true ATE, true individual counterfactuals),
+including an `nl` variant where an all-`ls` model is provably misspecified:
+
+| `nl` variant | ATE | vs true **+0.104** |
+|---|---|---|
+| naive observational contrast | +0.303 | confounded (overstates 2.9×) |
+| all-`ls` flow | +0.076 | undershoots (misses the age-varying effect) |
+| flexible (`ci`/`cs`) flow | +0.101 | **recovers the truth** |
+
+Full storyline, clinical-data context, R cross-check and reading notes:
+[`docs/stroke-case-study.md`](docs/stroke-case-study.md).
+
+## Layout
+
+```
+src/tramdag/            spec.py transforms.py conditioners.py flow.py
+                        simulations/   (magic_mrclean, triangle, vaca, carefl + CLIs)
+data/                   frozen synthetic CSVs + truth.json — a test contract
+experiments/            stroke pipeline, paper replications, training benchmark
+notebooks/              intro (didactic) + Colab demo   (jupytext .py — see README there)
+tests/                  66 tests: unit, known-truth recovery, R regression
+docs/                   training-speed.md, stroke-case-study.md
+```
+
+Implementation conventions (latent-scale signs, raw/one-hot parent encoding,
+log-space ordinal likelihood, seeding) are documented in
+[`CLAUDE.md`](CLAUDE.md) and pinned by tests.
+
+## Citation
+
+If you use `tramdag`, please cite the method paper:
+
+```bibtex
+@inproceedings{sick2025tramdag,
+  title     = {Interpretable Neural Causal Models with TRAM-DAGs},
+  author    = {Sick, Beate and D{\"u}rr, Oliver},
+  booktitle = {Proceedings of the 4th Conference on Causal Learning and Reasoning (CLeaR)},
+  series    = {Proceedings of Machine Learning Research},
+  volume    = {275},
+  year      = {2025},
+}
+```
+
+For the stroke application (and the `magic-mrclean` cohort design) additionally:
+
+```bibtex
+@article{duerr2026stroke,
+  title  = {Estimating Individualized Treatment Effects in Acute Ischemic Stroke
+            with Causal Transformation Models (TRAM-DAG): A Multi-Centre
+            Observational Study with External RCT Validation},
+  author = {D{\"u}rr, Oliver and Herzog, Lisa and B{\"u}hler, Pascal and
+            Wegener, Susanne and Sick, Beate},
+  journal = {arXiv preprint arXiv:2606.12623},
+  year   = {2026},
+}
+```