diff --git a/_typos.toml b/_typos.toml index 66c04e3..b47387e 100644 --- a/_typos.toml +++ b/_typos.toml @@ -2,3 +2,7 @@ # Allow dataset-specific column names from Hillstrom dataset womens = "womens" mens = "mens" + +[default.extend-identifiers] +# Author name in SciPy citation +VanderPlas = "VanderPlas" diff --git a/paper.bib b/paper.bib new file mode 100644 index 0000000..f6897d4 --- /dev/null +++ b/paper.bib @@ -0,0 +1,128 @@ +@misc{byambadalai2024estimatingdistributionaltreatmenteffects, + title={Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction}, + author={Undral Byambadalai and Tatsushi Oka and Shota Yasui}, + year={2024}, + eprint={2407.16037}, + archivePrefix={arXiv}, + primaryClass={econ.EM}, + url={https://arxiv.org/abs/2407.16037}, +} + +@book{fisher1935design, + title={The Design of Experiments}, + author={Fisher, Ronald A.}, + year={1935}, + publisher={Oliver and Boyd} +} + +@ARTICLE{2020NumPy-Array, + author = {Harris, Charles R. and Millman, K. Jarrod and + van der Walt, Stéfan J and Gommers, Ralf and + Virtanen, Pauli and Cournapeau, David and + Wieser, Eric and Taylor, Julian and Berg, Sebastian and + Smith, Nathaniel J. and Kern, Robert and Picus, Matti and + Hoyer, Stephan and van Kerkwijk, Marten H. and + Brett, Matthew and Haldane, Allan and + Fernández del Río, Jaime and Wiebe, Mark and + Peterson, Pearu and Gérard-Marchant, Pierre and + Sheppard, Kevin and Reddy, Tyler and Weckesser, Warren and + Abbasi, Hameer and Gohlke, Christoph and + Oliphant, Travis E.}, + title = {Array programming with {NumPy}}, + journal = {Nature}, + year = {2020}, + volume = {585}, + pages = {357--362}, + doi = {10.1038/s41586-020-2649-2} +} + +@article{scikit-learn, + title={Scikit-learn: Machine Learning in {P}ython}, + author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V. + and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P. + and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and + Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.}, + journal={Journal of Machine Learning Research}, + volume={12}, + pages={2825--2830}, + year={2011} +} + +@inproceedings{byambadalai2025efficientestimationdistributionaltreatment, + title={On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization}, + author={Undral Byambadalai and Tomu Hirata and Tatsushi Oka and Shota Yasui}, + booktitle={Proceedings of the 42nd International Conference on Machine Learning}, + year={2025}, + series={ICML'25}, + url={https://arxiv.org/abs/2506.05945} +} + +@misc{hirata2025efficientscalableestimationdistributional, + title={Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks}, + author={Tomu Hirata and Undral Byambadalai and Tatsushi Oka and Shota Yasui and Shunsuke Uto}, + year={2025}, + eprint={2507.07738}, + archivePrefix={arXiv}, + primaryClass={econ.EM}, + url={https://arxiv.org/abs/2507.07738} +} + +@misc{byambadalai2025imperfectcompliance, + title={Beyond the Average: Distributional Causal Inference under Imperfect Compliance}, + author={Undral Byambadalai and Tomu Hirata and Tatsushi Oka and Shota Yasui}, + year={2025}, + eprint={2509.15594}, + archivePrefix={arXiv}, + primaryClass={econ.EM}, + url={https://arxiv.org/abs/2509.15594} +} + +@article{dowhy, + title={DoWhy: An End-to-End Library for Causal Inference}, + author={Sharma, Amit and Kiciman, Emre}, + journal={arXiv preprint arXiv:2011.04216}, + year={2020}, + url={https://arxiv.org/abs/2011.04216} +} + +@inproceedings{econml, + title={Causal Inference and Machine Learning in Practice with {EconML} and {CausalML}: Industrial Use Cases at {Microsoft}, {TripAdvisor}, {Uber}}, + author={Battocchi, Keith and Dillon, Eleanor and Hei, Maggie and Lewis, Greg and Ling, Miruna and Rao, Jing and Shyr, Dennis and Syrgkanis, Vasilis}, + booktitle={Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery \& Data Mining}, + pages={4072--4073}, + year={2021}, + organization={ACM} +} + +@ARTICLE{2020SciPy-NMeth, + author = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and + Haberland, Matt and Reddy, Tyler and Cournapeau, David and + Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and + Bright, Jonathan and van der Walt, St{\'e}fan J. and + Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and + Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and + Kern, Robert and Larson, Eric and Carey, C J and + Polat, {\.I}lhan and Feng, Yu and Moore, Eric W. and + VanderPlas, Jake and Laxalde, Denis and Perktold, Josef and + Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and + Harris, Charles R. and Archibald, Anne M. and + Ribeiro, Ant{\^o}nio H. and Pedregosa, Fabian and + van Mulbregt, Paul and {SciPy 1.0 Contributors}}, + title = {{SciPy} 1.0: Fundamental Algorithms for Scientific Computing in Python}, + journal = {Nature Methods}, + year = {2020}, + volume = {17}, + pages = {261--272}, + doi = {10.1038/s41592-019-0686-2} +} + +@article{kobrosly2020causalcurve, + title={causal-curve: A Python Causal Inference Package to Estimate Causal Dose-Response Curves}, + author={Kobrosly, Roni W.}, + journal={Journal of Open Source Software}, + volume={5}, + number={52}, + pages={2523}, + year={2020}, + doi={10.21105/joss.02523} +} diff --git a/paper.md b/paper.md new file mode 100644 index 0000000..6940ee4 --- /dev/null +++ b/paper.md @@ -0,0 +1,75 @@ +--- +title: 'dte_adj: A Python Package for Estimating Distributional Treatment Effects in Randomized Experiments' +tags: + - Python + - randomized experiments + - causal inference + - distributional treatment effects + - machine learning + - variance reduction +authors: + - name: Tomu Hirata + orcid: 0009-0006-3140-291X + corresponding: true + affiliation: "1, 3" + - name: Undral Byambadalai + affiliation: 1 + - name: Tatsushi Oka + affiliation: "1, 2" + - name: Shota Yasui + affiliation: 1 +affiliations: + - name: CyberAgent, Inc., Japan + index: 1 + - name: Keio University, Japan + index: 2 + - name: Databricks Japan, Japan + index: 3 +date: 24 August 2025 +bibliography: paper.bib +--- + +# Summary + +`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. + +# Statement of Need + +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. + +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. + +# State of the Field + +Several Python packages address causal inference, but none focus on distributional treatment effects with machine learning-based variance reduction: + +- **SciPy** [@2020SciPy-NMeth]: Provides basic empirical cumulative distribution functions but offers no functionality for treatment effect estimation or confidence interval construction in experimental settings. +- **DoWhy** [@dowhy]: Focuses on causal graph-based inference and average treatment effects, without distributional analysis capabilities. +- **EconML** [@econml]: Incorporates machine learning for heterogeneous treatment effect estimation (CATE) but does not address distributional effects. +- **causal-curve** [@kobrosly2020causalcurve]: Estimates dose-response curves but targets continuous treatments rather than distributional outcomes. + +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. + +# Software Design + +`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: + +- **`SimpleDistributionEstimator`** and **`AdjustedDistributionEstimator`**: For simple randomized experiments, implementing methods from @byambadalai2024estimatingdistributionaltreatmenteffects. +- **`SimpleStratifiedDistributionEstimator`** and **`AdjustedStratifiedDistributionEstimator`**: For covariate-adaptive randomization designs, implementing methods from @byambadalai2025efficientestimationdistributionaltreatment. +- **`SimpleLocalDistributionEstimator`** and **`AdjustedLocalDistributionEstimator`**: For settings with imperfect compliance, implementing methods from @byambadalai2025imperfectcompliance. + +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. + +# Research Impact Statement + +The methods implemented in `dte_adj` have been published at top machine learning venues: ICML 2024 [@byambadalai2024estimatingdistributionaltreatmenteffects] and ICML 2025 [@byambadalai2025efficientestimationdistributionaltreatment]. The package has been used internally at CyberAgent, Inc. for analyzing A/B tests where distributional impacts are critical, such as evaluating interventions on user engagement metrics where tail behavior matters more than averages. The documentation includes tutorials demonstrating applications to the Hillstrom email marketing dataset and the Oregon Health Insurance Experiment, facilitating adoption by researchers in economics, marketing, and healthcare. + +# AI Usage Disclosure + +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. + +# Acknowledgements + +We thank CyberAgent, Inc. for supporting this research and the open-source community for valuable feedback during development. + +# References