UNIT — replication package

Reproduction code and frozen data for the simulation study of two-stage G-estimation under no essential heterogeneity, where the Stage-1 CATE learner is varied across a four-rung ladder (Ridge → tuned Random Forest → neural T-learner → neural shared-representation learner / TARNet).

Everything needed to rebuild every table and figure in the paper from committed data is here. The package is frozen at 200 seeds per cell for all methods.

Layout

src/rpmdeep/         simulation engine: DGPs, estimators, two-stage pipeline,
                     scenario "pool" runners
third_party/catenets vendored CATENets (neural CATE learners; BSD-3, see its LICENSE)
scripts/
  make_tables.py     aggregate frozen data -> results/tables/{summary_canonical.json,
                     cells.csv, tables.tex}   (single source of truth)
  plot_stage1_quality.py  -> figures/stage1_{nrmse,corr}.pdf
  verify.py          integrity + reproducibility check
results/
  pool_scenarioa/ pool_scenariob/ pool_scenario_cd/ pool_scenario_ext/ pool_tnet/
                     the frozen per-seed records (200 seeds/cell)
  tables/            reference table artifacts (regenerated by make_tables.py)
  MANIFEST.json      sha256 of every input + the certified table outputs
figures/             reference figures (regenerated by plot_stage1_quality.py)

Design: 7 scenarios (A, B, C, D, Dsmall, E, Wtau) × 5 sample sizes (500–10000) × 4 learners, 200 seeds each.

Quick start — rebuild the paper's numbers from frozen data

Only NumPy is needed for the tables (no jax/CATENets):

pip install numpy matplotlib            # or: pip install -r requirements.txt
python scripts/make_tables.py           # results/tables/{summary_canonical.json,cells.csv,tables.tex}
python scripts/plot_stage1_quality.py   # figures/stage1_{nrmse,corr}.pdf
python scripts/verify.py                # check integrity + byte-for-byte reproduction

verify.py confirms (1) every committed per-seed record matches its recorded sha256 and (2) re-running make_tables.py reproduces the table.

Regenerating the raw data from scratch

This requires the full environment, including jax/CATENets:

pip install -r requirements.txt
pip install -e .                        # installs the rpmdeep package
pip install -e third_party/catenets     # neural CATE learners

Each scenario pool is an independent, crash-resilient runner that writes one JSON checkpoint per seed (re-running skips seeds already on disk):

python -m rpmdeep.studies.pool_scenarioa
python -m rpmdeep.studies.pool_scenariob
python -m rpmdeep.studies.pool_scenario_cd      # scenarios C, D, E
python -m rpmdeep.studies.pool_scenario_ext     # Dsmall, Wtau
python -m rpmdeep.studies.pool_tnet             # neural T-learner (separate arms)

The committed results/pool_*/ directories are the first 200 seeds (by seed id) of these runs; make_tables.py deterministically consumes them.

Dependencies

See requirements.txt. Core table rebuild: NumPy (+ Matplotlib for figures). From-scratch simulation additionally needs scikit-learn, jax/jaxlib, and the vendored CATENets.

License & attribution

This project is released under the MIT License (LICENSE). It vendors CATENets under third_party/catenets/, which is distributed under the BSD-3-Clause License (see third_party/catenets/LICENSE).

Citation

If you use this package, please cite the accompanying article and this archive (machine-readable metadata in CITATION.cff):

Faleh, R., Morelli, S., & Brandt, H. (2026). Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity. arXiv:2607.10540. https://arxiv.org/abs/2607.10540

An archived release with a persistent DOI is available on Zenodo (DOI: to be added).