Skip to content

Commit 754bade

Browse files
TomeHirataclaude
andauthored
JOSS submission (#17)
* docs: draft for JOSS * yodate paper.md * Address PR review comments and update for latest codebase - Add "(RCTs, also known as A/B tests)" clarification per reviewer feedback - Add DoWhy and EconML comparison in Statement of Need section - Add imperfect compliance paper reference for LDTE estimators - Update CAR paper to ICML'25 proceedings format with updated author list - Update multi-task learning paper author list (add Shunsuke Uto) - Fix Fisher citation key to match bibliography 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Restructure paper to comply with JOSS submission guidelines - Add required "State of the Field" section comparing with SciPy, DoWhy, EconML, and causal-curve packages - Add required "Software Design" section explaining class architecture - Add required "Research Impact Statement" with ICML publications and industry usage - Add required "AI Usage Disclosure" section - Remove detailed "Features" section (API docs belong in documentation) - Add SciPy and causal-curve references to bibliography - Ensure word count is within 750-1750 range (776 words) - Add VanderPlas to typos allowlist (author name in SciPy citation) References: - JOSS submission guidelines: https://joss.readthedocs.io/en/latest/paper.html - JOSS review checklist: https://github.com/openjournals/joss/blob/main/docs/review_checklist.md 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Address Copilot review comments - Fix author metadata: make Tomu Hirata the single corresponding author and remove conflicting equal-contrib flag - Fix NumPy citation: replace Unicode en-dash (–) with BibTeX format (--) for compatibility across citation pipelines 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Move paper to paper/ subdirectory and add tutorial figures - Move paper.md and paper.bib into paper/ to avoid top-level clutter - Add hillstorm_dte.png and oregon_ldte_costs_comparison.png from tutorials - Add figure captions showing simple vs ML-adjusted comparison for both the Hillstrom email campaign and Oregon Health Insurance Experiment - Cross-reference figures in Research Impact Statement - Add dataset citations (Hillstrom 2008, Finkelstein et al. 2012) - Add GitHub Action (draft-pdf.yml) to auto-build JOSS PDF on push using paper-path: paper/paper.md Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Pin GitHub Actions to commit SHAs for supply chain security Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Update Databricks affiliation to Databricks, Inc., United States Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Add explicit permissions to draft-pdf workflow Addresses CodeQL security alert: restrict GITHUB_TOKEN to read-only contents access following least-privilege principle. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Fix figure paths for inara PDF build inara/pandoc runs from the repo root, so relative image paths in paper/paper.md resolve against root, not paper/. Copy PNGs to root before the build step so pandoc finds them. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Fix multiply-defined LaTeX labels for figures Pandoc extracts \label{} from caption text and also auto-generates a figure label, causing the label to be defined twice. Remove explicit \label{} from captions and replace \autoref{} with plain text references. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Fix figure 1 filename: hillstorm_dte.png (not hillstrom) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Remove unnecessary cp step in draft-pdf workflow inara CDs into the paper/ directory before running pandoc, so images in paper/ are already on the resource path. No copying needed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Add citations per reviewer feedback - Add Econometric Reviews publication (Oka et al. 2025) - Add NeurIPS 2025 reference for imperfect compliance paper - Add ABEMA field experiment paper (Yasui et al. 2026, Japanese Economic Review) as industry application evidence Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
1 parent f54ef49 commit 754bade

6 files changed

Lines changed: 267 additions & 0 deletions

File tree

.github/workflows/draft-pdf.yml

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
on: [push]
2+
3+
permissions:
4+
contents: read
5+
6+
jobs:
7+
paper:
8+
runs-on: ubuntu-latest
9+
name: Paper Draft
10+
steps:
11+
- name: Checkout
12+
uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
13+
- name: Build draft PDF
14+
uses: openjournals/openjournals-draft-action@85a18372e48f551d8af9ddb7a747de685fbbb01c # master
15+
with:
16+
journal: joss
17+
paper-path: paper/paper.md
18+
- name: Upload
19+
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
20+
with:
21+
name: paper
22+
path: paper/paper.pdf

_typos.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,7 @@
22
# Allow dataset-specific column names from Hillstrom dataset
33
womens = "womens"
44
mens = "mens"
5+
6+
[default.extend-identifiers]
7+
# Author name in SciPy citation
8+
VanderPlas = "VanderPlas"

paper/hillstorm_dte.png

180 KB
Loading
67.6 KB
Loading

paper/paper.bib

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
@misc{byambadalai2024estimatingdistributionaltreatmenteffects,
2+
title={Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction},
3+
author={Undral Byambadalai and Tatsushi Oka and Shota Yasui},
4+
year={2024},
5+
eprint={2407.16037},
6+
archivePrefix={arXiv},
7+
primaryClass={econ.EM},
8+
url={https://arxiv.org/abs/2407.16037},
9+
}
10+
11+
@book{fisher1935design,
12+
title={The Design of Experiments},
13+
author={Fisher, Ronald A.},
14+
year={1935},
15+
publisher={Oliver and Boyd}
16+
}
17+
18+
@ARTICLE{2020NumPy-Array,
19+
author = {Harris, Charles R. and Millman, K. Jarrod and
20+
van der Walt, Stéfan J and Gommers, Ralf and
21+
Virtanen, Pauli and Cournapeau, David and
22+
Wieser, Eric and Taylor, Julian and Berg, Sebastian and
23+
Smith, Nathaniel J. and Kern, Robert and Picus, Matti and
24+
Hoyer, Stephan and van Kerkwijk, Marten H. and
25+
Brett, Matthew and Haldane, Allan and
26+
Fernández del Río, Jaime and Wiebe, Mark and
27+
Peterson, Pearu and Gérard-Marchant, Pierre and
28+
Sheppard, Kevin and Reddy, Tyler and Weckesser, Warren and
29+
Abbasi, Hameer and Gohlke, Christoph and
30+
Oliphant, Travis E.},
31+
title = {Array programming with {NumPy}},
32+
journal = {Nature},
33+
year = {2020},
34+
volume = {585},
35+
pages = {357--362},
36+
doi = {10.1038/s41586-020-2649-2}
37+
}
38+
39+
@article{scikit-learn,
40+
title={Scikit-learn: Machine Learning in {P}ython},
41+
author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
42+
and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
43+
and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
44+
Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
45+
journal={Journal of Machine Learning Research},
46+
volume={12},
47+
pages={2825--2830},
48+
year={2011}
49+
}
50+
51+
@inproceedings{byambadalai2025efficientestimationdistributionaltreatment,
52+
title={On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization},
53+
author={Undral Byambadalai and Tomu Hirata and Tatsushi Oka and Shota Yasui},
54+
booktitle={Proceedings of the 42nd International Conference on Machine Learning},
55+
year={2025},
56+
series={ICML'25},
57+
url={https://arxiv.org/abs/2506.05945}
58+
}
59+
60+
@misc{hirata2025efficientscalableestimationdistributional,
61+
title={Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks},
62+
author={Tomu Hirata and Undral Byambadalai and Tatsushi Oka and Shota Yasui and Shunsuke Uto},
63+
year={2025},
64+
eprint={2507.07738},
65+
archivePrefix={arXiv},
66+
primaryClass={econ.EM},
67+
url={https://arxiv.org/abs/2507.07738}
68+
}
69+
70+
@misc{byambadalai2025imperfectcompliance,
71+
title={Beyond the Average: Distributional Causal Inference under Imperfect Compliance},
72+
author={Undral Byambadalai and Tomu Hirata and Tatsushi Oka and Shota Yasui},
73+
year={2025},
74+
eprint={2509.15594},
75+
archivePrefix={arXiv},
76+
primaryClass={econ.EM},
77+
url={https://arxiv.org/abs/2509.15594}
78+
}
79+
80+
@article{oka2025regression,
81+
title={Regression Adjustment for Estimating Distributional Treatment Effects in Randomized Controlled Trials},
82+
author={Oka, Tatsushi and Yasui, Shota and Hayakawa, Yuta and Byambadalai, Undral},
83+
journal={Econometric Reviews},
84+
year={2025},
85+
doi={10.1080/07474938.2025.2538843}
86+
}
87+
88+
@article{yasui2026abema,
89+
title={Distributional Treatment Effects of Content Promotion: Evidence from an {ABEMA} Field Experiment},
90+
author={Yasui, Shota and Oka, Tatsushi and Byambadalai, Undral and Oishi, Yuki},
91+
journal={The Japanese Economic Review},
92+
year={2026},
93+
doi={10.1007/s42973-026-00234-y}
94+
}
95+
96+
@article{dowhy,
97+
title={DoWhy: An End-to-End Library for Causal Inference},
98+
author={Sharma, Amit and Kiciman, Emre},
99+
journal={arXiv preprint arXiv:2011.04216},
100+
year={2020},
101+
url={https://arxiv.org/abs/2011.04216}
102+
}
103+
104+
@inproceedings{econml,
105+
title={Causal Inference and Machine Learning in Practice with {EconML} and {CausalML}: Industrial Use Cases at {Microsoft}, {TripAdvisor}, {Uber}},
106+
author={Battocchi, Keith and Dillon, Eleanor and Hei, Maggie and Lewis, Greg and Ling, Miruna and Rao, Jing and Shyr, Dennis and Syrgkanis, Vasilis},
107+
booktitle={Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery \& Data Mining},
108+
pages={4072--4073},
109+
year={2021},
110+
organization={ACM}
111+
}
112+
113+
@ARTICLE{2020SciPy-NMeth,
114+
author = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and
115+
Haberland, Matt and Reddy, Tyler and Cournapeau, David and
116+
Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and
117+
Bright, Jonathan and van der Walt, St{\'e}fan J. and
118+
Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and
119+
Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and
120+
Kern, Robert and Larson, Eric and Carey, C J and
121+
Polat, {\.I}lhan and Feng, Yu and Moore, Eric W. and
122+
VanderPlas, Jake and Laxalde, Denis and Perktold, Josef and
123+
Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and
124+
Harris, Charles R. and Archibald, Anne M. and
125+
Ribeiro, Ant{\^o}nio H. and Pedregosa, Fabian and
126+
van Mulbregt, Paul and {SciPy 1.0 Contributors}},
127+
title = {{SciPy} 1.0: Fundamental Algorithms for Scientific Computing in Python},
128+
journal = {Nature Methods},
129+
year = {2020},
130+
volume = {17},
131+
pages = {261--272},
132+
doi = {10.1038/s41592-019-0686-2}
133+
}
134+
135+
@misc{hillstrom2008,
136+
title={The MineThatData E-Mail Analytics And Data Mining Challenge},
137+
author={Hillstrom, Kevin},
138+
year={2008},
139+
url={https://blog.minethatdata.com/2008/03/minethatdata-e-mail-analytics-and-data.html}
140+
}
141+
142+
@article{finkelstein2012,
143+
title={The Oregon Health Insurance Experiment: Evidence from the First Year},
144+
author={Finkelstein, Amy and Taubman, Sarah and Wright, Bill and Bernstein, Mira and Gruber, Jonathan and Newhouse, Joseph P. and Allen, Heidi and Baicker, Katherine and {Oregon Health Study Group}},
145+
journal={The Quarterly Journal of Economics},
146+
volume={127},
147+
number={3},
148+
pages={1057--1106},
149+
year={2012},
150+
doi={10.1093/qje/qjs020}
151+
}
152+
153+
@article{kobrosly2020causalcurve,
154+
title={causal-curve: A Python Causal Inference Package to Estimate Causal Dose-Response Curves},
155+
author={Kobrosly, Roni W.},
156+
journal={Journal of Open Source Software},
157+
volume={5},
158+
number={52},
159+
pages={2523},
160+
year={2020},
161+
doi={10.21105/joss.02523}
162+
}

paper/paper.md

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
---
2+
title: 'dte_adj: A Python Package for Estimating Distributional Treatment Effects in Randomized Experiments'
3+
tags:
4+
- Python
5+
- randomized experiments
6+
- causal inference
7+
- distributional treatment effects
8+
- machine learning
9+
- variance reduction
10+
authors:
11+
- name: Tomu Hirata
12+
orcid: 0009-0006-3140-291X
13+
corresponding: true
14+
affiliation: "1, 3"
15+
- name: Undral Byambadalai
16+
affiliation: 1
17+
- name: Tatsushi Oka
18+
affiliation: "1, 2"
19+
- name: Shota Yasui
20+
affiliation: 1
21+
affiliations:
22+
- name: CyberAgent, Inc., Japan
23+
index: 1
24+
- name: Keio University, Japan
25+
index: 2
26+
- name: Databricks, Inc., United States
27+
index: 3
28+
date: 24 August 2025
29+
bibliography: paper.bib
30+
---
31+
32+
# Summary
33+
34+
`dte_adj` is a Python package for estimating distributional treatment effects (DTEs) in randomized experiments (RCTs, also known as A/B tests). Unlike traditional approaches that focus on average treatment effects, `dte_adj` enables researchers to analyze the full distributional impact of interventions across different outcome levels. The package implements machine learning-enhanced regression adjustment methods for variance reduction, supports multiple experimental designs including simple randomization, covariate-adaptive randomization, and settings with imperfect compliance, and provides a scikit-learn compatible API with comprehensive functionality for computing distribution functions, probability treatment effects, and quantile treatment effects with confidence intervals.
35+
36+
# Statement of Need
37+
38+
Randomized experiments have been fundamental to scientific inquiry since @fisher1935design, providing the gold standard for causal inference. While most experimental analyses focus on average treatment effects (ATEs), many research questions require understanding how treatments affect the entire distribution of outcomes. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask. For example, a policy intervention might have no effect on average income while substantially reducing poverty rates at lower quantiles, or a medical treatment might benefit patients at the tails of the distribution differently than those near the median.
39+
40+
Despite the growing importance of distributional analysis in economics, medicine, and technology, the Python ecosystem lacks comprehensive tools for DTE estimation with modern variance reduction techniques. Researchers often resort to basic empirical CDFs or manual implementations that lack statistical rigor. `dte_adj` fills this gap by providing a unified framework for distributional treatment effect analysis that integrates state-of-the-art machine learning methods for improved precision, rigorous confidence interval construction, and support for complex experimental designs.
41+
42+
# State of the Field
43+
44+
Several Python packages address causal inference, but none focus on distributional treatment effects with machine learning-based variance reduction:
45+
46+
- **SciPy** [@2020SciPy-NMeth]: Provides basic empirical cumulative distribution functions but offers no functionality for treatment effect estimation or confidence interval construction in experimental settings.
47+
- **DoWhy** [@dowhy]: Focuses on causal graph-based inference and average treatment effects, without distributional analysis capabilities.
48+
- **EconML** [@econml]: Incorporates machine learning for heterogeneous treatment effect estimation (CATE) but does not address distributional effects.
49+
- **causal-curve** [@kobrosly2020causalcurve]: Estimates dose-response curves but targets continuous treatments rather than distributional outcomes.
50+
51+
In the R ecosystem, packages like `qte` provide quantile treatment effect estimation but lack machine learning integration for variance reduction. `dte_adj` uniquely combines: (1) distributional treatment effect estimation across the full outcome distribution, (2) machine learning-enhanced regression adjustment for precision gains, and (3) support for multiple experimental designs including covariate-adaptive randomization and imperfect compliance settings.
52+
53+
# Software Design
54+
55+
`dte_adj` follows a class-based architecture with a template method pattern, where a base class defines the algorithm structure and subclasses implement design-specific computations:
56+
57+
- **`SimpleDistributionEstimator`** and **`AdjustedDistributionEstimator`**: For simple randomized experiments, implementing methods from @byambadalai2024estimatingdistributionaltreatmenteffects.
58+
- **`SimpleStratifiedDistributionEstimator`** and **`AdjustedStratifiedDistributionEstimator`**: For covariate-adaptive randomization designs, implementing methods from @byambadalai2025efficientestimationdistributionaltreatment.
59+
- **`SimpleLocalDistributionEstimator`** and **`AdjustedLocalDistributionEstimator`**: For settings with imperfect compliance, implementing methods from @byambadalai2025imperfectcompliance.
60+
61+
All estimators implement a consistent API with three primary methods: `predict_dte()` for distributional treatment effects, `predict_pte()` for probability treatment effects over intervals, and `predict_qte()` for quantile treatment effects. The adjusted estimators use K-fold cross-fitting to prevent overfitting and support both single-task and multi-task learning modes [@hirata2025efficientscalableestimationdistributional] for computational efficiency. Bootstrap methods provide confidence intervals with multiple variance estimation approaches.
62+
63+
![Distributional treatment effects for the Hillstrom email marketing dataset [@hillstrom2008], comparing Women's vs Men's email campaigns. The simple estimator (left, purple) and ML-adjusted estimator (right, green) show that adjustment substantially tightens confidence bands, demonstrating the variance reduction benefit of regression adjustment.](hillstorm_dte.png)
64+
65+
![Local distributional treatment effects for emergency department costs in the Oregon Health Insurance Experiment [@finkelstein2012], estimated using `SimpleLocalDistributionEstimator` (left) and `AdjustedLocalDistributionEstimator` (right). Health insurance coverage shifts the distribution of ED costs, with ML adjustment again yielding narrower confidence intervals.](oregon_ldte_costs_comparison.png)
66+
67+
# Research Impact Statement
68+
69+
The methods implemented in `dte_adj` have been published across machine learning and econometrics venues: ICML 2024 [@byambadalai2024estimatingdistributionaltreatmenteffects], Econometric Reviews [@oka2025regression], ICML 2025 [@byambadalai2025efficientestimationdistributionaltreatment], and NeurIPS 2025 [@byambadalai2025imperfectcompliance]. The package has been applied in industry settings, including analyzing the distributional impact of content promotion on user engagement at ABEMA, a major video streaming platform [@yasui2026abema]. The documentation includes tutorials demonstrating applications to the Hillstrom email marketing dataset (Figure 1) and the Oregon Health Insurance Experiment (Figure 2), facilitating adoption by researchers in economics, marketing, and healthcare.
70+
71+
# AI Usage Disclosure
72+
73+
Generative AI tools (Claude) were used to assist with documentation writing and code review during development. All AI-generated content was reviewed and validated by the human authors.
74+
75+
# Acknowledgements
76+
77+
We thank CyberAgent, Inc. for supporting this research and the open-source community for valuable feedback during development.
78+
79+
# References

0 commit comments

Comments
 (0)