From a35c87e9f826ae07373de5146f7131ec76dcfb94 Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Fri, 9 Aug 2024 15:28:57 +0900 Subject: [PATCH 1/5] docs: draft for JOSS --- paper.bib | 49 ++++++++++++++++++++++++++++++++++++++ paper.md | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 paper.bib create mode 100644 paper.md diff --git a/paper.bib b/paper.bib new file mode 100644 index 0000000..b58caf7 --- /dev/null +++ b/paper.bib @@ -0,0 +1,49 @@ +@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} +} \ No newline at end of file diff --git a/paper.md b/paper.md new file mode 100644 index 0000000..dfc041b --- /dev/null +++ b/paper.md @@ -0,0 +1,71 @@ +--- +title: 'dte_adj: A Python package for Distributional Treatment Effects' +tags: + - Python + - Distributional Treatment Effects + - Variance Reduction +authors: + - name: Tomu Hirata + orcid: 0009-0006-3140-291X + equal-contrib: true + affiliation: "1, 3" + - name: Undral Byambadalai + corresponding: true + affiliation: 1 + - name: Tatsushi Oka + corresponding: true + affiliation: "1, 2" + - name: Shota Yasui + corresponding: true + affiliation: 1 +affiliations: + - name: Cyber Agent, Inc, Japan + index: 1 + - name: Keio University, Japan + index: 2 + - name: Indeed Technologies Japan, Japan + index: 3 +date: 9 August 2024 +bibliography: paper.bib + +# Optional fields if submitting to a AAS journal too, see this blog post: +# https://blog.joss.theoj.org/2018/12/a-new-collaboration-with-aas-publishing +aas-doi: 10.3847/xxxxx +aas-journal: International Conference on Machine Learning +--- + +# Summary + +`dte_adj` is a Python package for computing empirical cumulative distribution function (CDF) and distributional treatment effect (DTE) from data obtained by Randomized control tests. This package also contains a novel method to reduce variance of DTE using pre-treatment covariates introduced in `@Undral:2024`. + +# Statement of need + +Since the groundbreaking work by `@Fisher:1935`, randomized experiments have been essential in understanding the impact of interventions and shaping policy decisions. A widely used metric in this context is the Average Treatment Effect (ATE). However, exploring the distributional treatment effects often offers a more nuanced understanding than focusing solely on the average effects. +Python is widely used in the research community recently with its flexibility and ease-of-use in the user-interface. However, there is no popular Python library for computing Distributional Treatment Effect from data obtained from randomized experiments. While scipy provides a method for computing the empirical cumulative distribution function, it lacks convenient functions for calculating DTE or for estimating the variance of the distribution. +`dte_adj` was developed to fill the gap by offering the functionalities for 1) computing CDF from data, 2) calculating DTE and its confidence band based on CDF and 3) visualizing DTE. This library uses `numpy` as input and output of methods, which is widely used for matrix computation in Python. The main classes of this library also follows the interface of popular library `scikit-learn`, which makes it easy for the users with Machine Learning development experieneces. + +# Functionalities + +The high level functionalities of `dte_adj` are as follows: +1. Computing CDF and its variance based on number arrays +2. Calculating distributional parameters and their confidence bands +3. Visualiving distributional parameters and the confidence bands + +It currently offers two classes to compute CDF and its variance. +- `SimpleDistributionEstimator`: this class offers a standard way to compute empirical CDF +- `AdjustedDistributionEstimator`: this class offers a way to compute CDF with smaller variance adjusted by pre-treatment covariates introduced in `@Undral:2024` + +Both classes implement following methods to calculate distributional parameters. +- `predict_dte`: method for computing Distributional Treatment Effect $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $y$ is an outcome variable, $w$ is treatment type , and $F_{Y(w)}(y)$ is cumulative likelihood for treatment type $w$ and outcome $y$. +- `predict_pte`: method for computing Probability Treatment Effect (PTE) $PTE_{w, w'}(y, h) := \left( F_{Y(w)}(y+h) - F_{Y(w)}(y) \right) - \left( F_{Y(w')}(y+h) - F_{Y(w')}(y) \right)$, where $h > 0$ is an interval of each evaluation window. +- `predict_qte`: method for computing Quantile Treatment Effect (QTE) $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, where $\tau$ is quantile. + +Lastly, `dte_adj.plot` module can be used for visualiting the distribution parameters. The examples of the visualization are available in the figures below. + +![DTE](docs/source/_static/dte_moment.png) +![PTE](docs/source/_static/pte_simple.png) +![QTE](docs/source/_static/qte.png) + +# Acknowledgements + +# References \ No newline at end of file From f0c4af66bba272af7158c82d0cc9da893d018536 Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Sun, 24 Aug 2025 21:21:22 +0900 Subject: [PATCH 2/5] yodate paper.md --- paper.bib | 26 ++++++++++++++-- paper.md | 90 +++++++++++++++++++++++++++++++++++-------------------- 2 files changed, 80 insertions(+), 36 deletions(-) diff --git a/paper.bib b/paper.bib index b58caf7..851c0b5 100644 --- a/paper.bib +++ b/paper.bib @@ -1,11 +1,11 @@ @misc{byambadalai2024estimatingdistributionaltreatmenteffects, - title={Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction}, + 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}, + url={https://arxiv.org/abs/2407.16037}, } @book{fisher1935design, @@ -46,4 +46,24 @@ @article{scikit-learn volume={12}, pages={2825--2830}, year={2011} -} \ No newline at end of file +} + +@misc{byambadalai2025efficientestimationdistributionaltreatment, + title={On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization}, + author={Undral Byambadalai and Tatsushi Oka and Shota Yasui}, + year={2025}, + eprint={2506.05945}, + archivePrefix={arXiv}, + primaryClass={econ.EM}, + 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}, + year={2025}, + eprint={2507.07738}, + archivePrefix={arXiv}, + primaryClass={econ.EM}, + url={https://arxiv.org/abs/2507.07738} +} diff --git a/paper.md b/paper.md index dfc041b..d2885e0 100644 --- a/paper.md +++ b/paper.md @@ -1,71 +1,95 @@ --- -title: 'dte_adj: A Python package for Distributional Treatment Effects' +title: 'dte_adj: A Python Package for Estimating Distributional Treatment Effects in Randomized Experiments' tags: - Python - - Distributional Treatment Effects - - Variance Reduction + - randomized experiments + - causal inference + - distributional treatment effects + - machine learning + - variance reduction authors: - name: Tomu Hirata orcid: 0009-0006-3140-291X equal-contrib: true affiliation: "1, 3" - - name: Undral Byambadalai + - name: Undral Byambadalai corresponding: true affiliation: 1 - name: Tatsushi Oka corresponding: true affiliation: "1, 2" - - name: Shota Yasui + - name: Shota Yasui corresponding: true affiliation: 1 affiliations: - - name: Cyber Agent, Inc, Japan + - name: CyberAgent, Inc., Japan index: 1 - name: Keio University, Japan index: 2 - - name: Indeed Technologies Japan, Japan + - name: Databricks Japan, Japan index: 3 -date: 9 August 2024 +date: 24 August 2025 bibliography: paper.bib - -# Optional fields if submitting to a AAS journal too, see this blog post: -# https://blog.joss.theoj.org/2018/12/a-new-collaboration-with-aas-publishing -aas-doi: 10.3847/xxxxx -aas-journal: International Conference on Machine Learning --- # Summary -`dte_adj` is a Python package for computing empirical cumulative distribution function (CDF) and distributional treatment effect (DTE) from data obtained by Randomized control tests. This package also contains a novel method to reduce variance of DTE using pre-treatment covariates introduced in `@Undral:2024`. +`dte_adj` is a Python package designed for estimating distributional treatment effects (DTEs) in randomized experiments. 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 to achieve variance reduction, making distributional effect estimation more precise and computationally efficient. It supports multiple experimental designs including simple randomization, covariate-adaptive randomization (CAR), and local distributional treatment effect (LDTE) estimation. The package provides a scikit-learn compatible API and 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 the pioneering work of @Fisher:1935, 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, not just the mean. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask. + +Despite the growing importance of distributional analysis in fields ranging from economics to medicine, the Python ecosystem lacks comprehensive tools for DTE estimation. While SciPy provides basic empirical cumulative distribution functions, it offers no specialized functionality for treatment effect estimation, variance reduction, or confidence interval construction in experimental settings. Existing R packages like `RDDtools` focus on regression discontinuity rather than randomized experiments, and lack modern machine learning integration. + +`dte_adj` addresses this gap by providing a comprehensive Python framework for distributional treatment effect analysis. The package implements state-of-the-art variance reduction techniques using machine learning models for regression adjustment [@byambadalai2024estimatingdistributionaltreatmenteffects], enabling more precise DTE estimates with smaller sample sizes. It supports multiple experimental designs including covariate-adaptive randomization [@byambadalai2025efficientestimationdistributionaltreatment] and local treatment effects, with a scikit-learn [@scikit-learn] compatible API that integrates seamlessly into existing machine learning workflows. This makes advanced distributional analysis accessible to the broader Python research community, supporting more nuanced causal inference in experimental studies. + +# Features + +`dte_adj` provides a comprehensive suite of tools for distributional treatment effect analysis: + +## Estimator Classes -# Statement of need +The package implements multiple estimator classes following a hierarchical design pattern: -Since the groundbreaking work by `@Fisher:1935`, randomized experiments have been essential in understanding the impact of interventions and shaping policy decisions. A widely used metric in this context is the Average Treatment Effect (ATE). However, exploring the distributional treatment effects often offers a more nuanced understanding than focusing solely on the average effects. -Python is widely used in the research community recently with its flexibility and ease-of-use in the user-interface. However, there is no popular Python library for computing Distributional Treatment Effect from data obtained from randomized experiments. While scipy provides a method for computing the empirical cumulative distribution function, it lacks convenient functions for calculating DTE or for estimating the variance of the distribution. -`dte_adj` was developed to fill the gap by offering the functionalities for 1) computing CDF from data, 2) calculating DTE and its confidence band based on CDF and 3) visualizing DTE. This library uses `numpy` as input and output of methods, which is widely used for matrix computation in Python. The main classes of this library also follows the interface of popular library `scikit-learn`, which makes it easy for the users with Machine Learning development experieneces. +**Simple Randomization Estimators:** +- `SimpleDistributionEstimator`: Basic empirical distribution function estimator for simple randomized experiments +- `AdjustedDistributionEstimator`: Machine learning-enhanced estimator with regression adjustment for variance reduction -# Functionalities +**Stratified Estimators (for Covariate-Adaptive Randomization):** +- `SimpleStratifiedDistributionEstimator`: Handles stratified block randomization designs +- `AdjustedStratifiedDistributionEstimator`: Combines stratification with ML-based variance reduction -The high level functionalities of `dte_adj` are as follows: -1. Computing CDF and its variance based on number arrays -2. Calculating distributional parameters and their confidence bands -3. Visualiving distributional parameters and the confidence bands +**Local Distribution Estimators:** +- `SimpleLocalDistributionEstimator`: Estimates local distributional treatment effects (LDTE) +- `AdjustedLocalDistributionEstimator`: LDTE estimation with ML adjustment for improved precision -It currently offers two classes to compute CDF and its variance. -- `SimpleDistributionEstimator`: this class offers a standard way to compute empirical CDF -- `AdjustedDistributionEstimator`: this class offers a way to compute CDF with smaller variance adjusted by pre-treatment covariates introduced in `@Undral:2024` +## Core Methods -Both classes implement following methods to calculate distributional parameters. -- `predict_dte`: method for computing Distributional Treatment Effect $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $y$ is an outcome variable, $w$ is treatment type , and $F_{Y(w)}(y)$ is cumulative likelihood for treatment type $w$ and outcome $y$. -- `predict_pte`: method for computing Probability Treatment Effect (PTE) $PTE_{w, w'}(y, h) := \left( F_{Y(w)}(y+h) - F_{Y(w)}(y) \right) - \left( F_{Y(w')}(y+h) - F_{Y(w')}(y) \right)$, where $h > 0$ is an interval of each evaluation window. -- `predict_qte`: method for computing Quantile Treatment Effect (QTE) $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, where $\tau$ is quantile. +All estimators implement a consistent API with three primary methods: -Lastly, `dte_adj.plot` module can be used for visualiting the distribution parameters. The examples of the visualization are available in the figures below. +- `predict_dte()`: Computes Distributional Treatment Effects $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $F_{Y(w)}(y)$ represents the cumulative distribution function for treatment $w$ at outcome level $y$. + +- `predict_pte()`: Computes Probability Treatment Effects over specified intervals, measuring differences in probability mass between treatment groups. + +- `predict_qte()`: Computes Quantile Treatment Effects $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, comparing quantiles across treatments. + +## Advanced Features + +**Multi-task Learning:** The package supports multi-task neural networks (`is_multi_task=True`) for computational efficiency when analyzing many outcome locations simultaneously [@hirata2025efficientscalableestimationdistributional]. + +**Cross-fitting:** Adjusted estimators use K-fold cross-fitting to prevent overfitting in machine learning models, ensuring robust treatment effect estimates. + +**Confidence Intervals:** Built-in bootstrap methods provide confidence intervals with multiple variance estimation approaches (`moment`, `simple`, `uniform`). + +**Visualization:** The `dte_adj.plot` module enables easy plotting of treatment effects and confidence bands. ![DTE](docs/source/_static/dte_moment.png) -![PTE](docs/source/_static/pte_simple.png) +![PTE](docs/source/_static/pte_empirical.png) ![QTE](docs/source/_static/qte.png) # Acknowledgements -# References \ No newline at end of file +We thank CyberAgent, Inc. for supporting this research and the open-source community for valuable feedback during development. + +# References From b312938cf43948ebe6bc6bdf2051fc4ac7156aa0 Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Tue, 31 Mar 2026 20:07:02 +0900 Subject: [PATCH 3/5] Address PR review comments and update for latest codebase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- paper.bib | 38 ++++++++++++++++++++++++++++++++------ paper.md | 10 +++++----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/paper.bib b/paper.bib index 851c0b5..e6c0c5c 100644 --- a/paper.bib +++ b/paper.bib @@ -48,22 +48,48 @@ @article{scikit-learn year={2011} } -@misc{byambadalai2025efficientestimationdistributionaltreatment, +@inproceedings{byambadalai2025efficientestimationdistributionaltreatment, title={On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization}, - author={Undral Byambadalai and Tatsushi Oka and Shota Yasui}, + author={Undral Byambadalai and Tomu Hirata and Tatsushi Oka and Shota Yasui}, + booktitle={Proceedings of the 42nd International Conference on Machine Learning}, year={2025}, - eprint={2506.05945}, - archivePrefix={arXiv}, - primaryClass={econ.EM}, + 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}, + 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} +} diff --git a/paper.md b/paper.md index d2885e0..8d35b21 100644 --- a/paper.md +++ b/paper.md @@ -38,9 +38,9 @@ bibliography: paper.bib # Statement of Need -Randomized experiments have been fundamental to scientific inquiry since the pioneering work of @Fisher:1935, 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, not just the mean. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask. +Randomized experiments (RCTs, also known as A/B tests) have been fundamental to scientific inquiry since the pioneering work of @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, not just the mean. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask. -Despite the growing importance of distributional analysis in fields ranging from economics to medicine, the Python ecosystem lacks comprehensive tools for DTE estimation. While SciPy provides basic empirical cumulative distribution functions, it offers no specialized functionality for treatment effect estimation, variance reduction, or confidence interval construction in experimental settings. Existing R packages like `RDDtools` focus on regression discontinuity rather than randomized experiments, and lack modern machine learning integration. +Despite the growing importance of distributional analysis in fields ranging from economics to medicine, the Python ecosystem lacks comprehensive tools for DTE estimation. While SciPy provides basic empirical cumulative distribution functions, it offers no specialized functionality for treatment effect estimation, variance reduction, or confidence interval construction in experimental settings. Existing Python packages for causal inference, such as DoWhy [@dowhy] and EconML [@econml], focus primarily on average treatment effects and conditional average treatment effects (CATE), with EconML incorporating machine learning for heterogeneous treatment effect estimation. However, these packages do not address distributional treatment effects, which capture how treatments affect the entire outcome distribution rather than just the mean. Existing R packages like `RDDtools` focus on regression discontinuity rather than randomized experiments, and lack modern machine learning integration. `dte_adj` addresses this gap by providing a comprehensive Python framework for distributional treatment effect analysis. The package implements state-of-the-art variance reduction techniques using machine learning models for regression adjustment [@byambadalai2024estimatingdistributionaltreatmenteffects], enabling more precise DTE estimates with smaller sample sizes. It supports multiple experimental designs including covariate-adaptive randomization [@byambadalai2025efficientestimationdistributionaltreatment] and local treatment effects, with a scikit-learn [@scikit-learn] compatible API that integrates seamlessly into existing machine learning workflows. This makes advanced distributional analysis accessible to the broader Python research community, supporting more nuanced causal inference in experimental studies. @@ -60,9 +60,9 @@ The package implements multiple estimator classes following a hierarchical desig - `SimpleStratifiedDistributionEstimator`: Handles stratified block randomization designs - `AdjustedStratifiedDistributionEstimator`: Combines stratification with ML-based variance reduction -**Local Distribution Estimators:** -- `SimpleLocalDistributionEstimator`: Estimates local distributional treatment effects (LDTE) -- `AdjustedLocalDistributionEstimator`: LDTE estimation with ML adjustment for improved precision +**Local Distribution Estimators (for Imperfect Compliance):** +- `SimpleLocalDistributionEstimator`: Estimates local distributional treatment effects (LDTE) for settings with imperfect compliance +- `AdjustedLocalDistributionEstimator`: LDTE estimation with ML adjustment for improved precision [@byambadalai2025imperfectcompliance] ## Core Methods From 29eb364f9f9dab1a85a107eded6880673927b403 Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Tue, 31 Mar 2026 20:42:36 +0900 Subject: [PATCH 4/5] Restructure paper to comply with JOSS submission guidelines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- _typos.toml | 4 ++++ paper.bib | 33 +++++++++++++++++++++++++++++++ paper.md | 57 +++++++++++++++++++---------------------------------- 3 files changed, 57 insertions(+), 37 deletions(-) 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 index e6c0c5c..becfdf9 100644 --- a/paper.bib +++ b/paper.bib @@ -93,3 +93,36 @@ @inproceedings{econml 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 index 8d35b21..84af10d 100644 --- a/paper.md +++ b/paper.md @@ -34,59 +34,42 @@ bibliography: paper.bib # Summary -`dte_adj` is a Python package designed for estimating distributional treatment effects (DTEs) in randomized experiments. 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 to achieve variance reduction, making distributional effect estimation more precise and computationally efficient. It supports multiple experimental designs including simple randomization, covariate-adaptive randomization (CAR), and local distributional treatment effect (LDTE) estimation. The package provides a scikit-learn compatible API and comprehensive functionality for computing distribution functions, probability treatment effects, and quantile treatment effects with confidence intervals. +`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 (RCTs, also known as A/B tests) have been fundamental to scientific inquiry since the pioneering work of @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, not just the mean. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask. +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 fields ranging from economics to medicine, the Python ecosystem lacks comprehensive tools for DTE estimation. While SciPy provides basic empirical cumulative distribution functions, it offers no specialized functionality for treatment effect estimation, variance reduction, or confidence interval construction in experimental settings. Existing Python packages for causal inference, such as DoWhy [@dowhy] and EconML [@econml], focus primarily on average treatment effects and conditional average treatment effects (CATE), with EconML incorporating machine learning for heterogeneous treatment effect estimation. However, these packages do not address distributional treatment effects, which capture how treatments affect the entire outcome distribution rather than just the mean. Existing R packages like `RDDtools` focus on regression discontinuity rather than randomized experiments, and lack modern machine learning integration. +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. -`dte_adj` addresses this gap by providing a comprehensive Python framework for distributional treatment effect analysis. The package implements state-of-the-art variance reduction techniques using machine learning models for regression adjustment [@byambadalai2024estimatingdistributionaltreatmenteffects], enabling more precise DTE estimates with smaller sample sizes. It supports multiple experimental designs including covariate-adaptive randomization [@byambadalai2025efficientestimationdistributionaltreatment] and local treatment effects, with a scikit-learn [@scikit-learn] compatible API that integrates seamlessly into existing machine learning workflows. This makes advanced distributional analysis accessible to the broader Python research community, supporting more nuanced causal inference in experimental studies. +# State of the Field -# Features +Several Python packages address causal inference, but none focus on distributional treatment effects with machine learning-based variance reduction: -`dte_adj` provides a comprehensive suite of tools for distributional treatment effect analysis: +- **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. -## Estimator Classes +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. -The package implements multiple estimator classes following a hierarchical design pattern: +# Software Design -**Simple Randomization Estimators:** -- `SimpleDistributionEstimator`: Basic empirical distribution function estimator for simple randomized experiments -- `AdjustedDistributionEstimator`: Machine learning-enhanced estimator with regression adjustment for variance reduction +`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: -**Stratified Estimators (for Covariate-Adaptive Randomization):** -- `SimpleStratifiedDistributionEstimator`: Handles stratified block randomization designs -- `AdjustedStratifiedDistributionEstimator`: Combines stratification with ML-based variance reduction +- **`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. -**Local Distribution Estimators (for Imperfect Compliance):** -- `SimpleLocalDistributionEstimator`: Estimates local distributional treatment effects (LDTE) for settings with imperfect compliance -- `AdjustedLocalDistributionEstimator`: LDTE estimation with ML adjustment for improved precision [@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. -## Core Methods +# Research Impact Statement -All estimators implement a consistent API with three primary methods: +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. -- `predict_dte()`: Computes Distributional Treatment Effects $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $F_{Y(w)}(y)$ represents the cumulative distribution function for treatment $w$ at outcome level $y$. +# AI Usage Disclosure -- `predict_pte()`: Computes Probability Treatment Effects over specified intervals, measuring differences in probability mass between treatment groups. - -- `predict_qte()`: Computes Quantile Treatment Effects $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, comparing quantiles across treatments. - -## Advanced Features - -**Multi-task Learning:** The package supports multi-task neural networks (`is_multi_task=True`) for computational efficiency when analyzing many outcome locations simultaneously [@hirata2025efficientscalableestimationdistributional]. - -**Cross-fitting:** Adjusted estimators use K-fold cross-fitting to prevent overfitting in machine learning models, ensuring robust treatment effect estimates. - -**Confidence Intervals:** Built-in bootstrap methods provide confidence intervals with multiple variance estimation approaches (`moment`, `simple`, `uniform`). - -**Visualization:** The `dte_adj.plot` module enables easy plotting of treatment effects and confidence bands. - -![DTE](docs/source/_static/dte_moment.png) -![PTE](docs/source/_static/pte_empirical.png) -![QTE](docs/source/_static/qte.png) +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 From 38ea69b5ed767795040b0d588e460224f5d3d465 Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Tue, 31 Mar 2026 20:48:21 +0900 Subject: [PATCH 5/5] Address Copilot review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- paper.bib | 2 +- paper.md | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/paper.bib b/paper.bib index becfdf9..f6897d4 100644 --- a/paper.bib +++ b/paper.bib @@ -32,7 +32,7 @@ @ARTICLE{2020NumPy-Array journal = {Nature}, year = {2020}, volume = {585}, - pages = {357–362}, + pages = {357--362}, doi = {10.1038/s41586-020-2649-2} } diff --git a/paper.md b/paper.md index 84af10d..6940ee4 100644 --- a/paper.md +++ b/paper.md @@ -10,16 +10,13 @@ tags: authors: - name: Tomu Hirata orcid: 0009-0006-3140-291X - equal-contrib: true + corresponding: true affiliation: "1, 3" - name: Undral Byambadalai - corresponding: true affiliation: 1 - name: Tatsushi Oka - corresponding: true affiliation: "1, 2" - name: Shota Yasui - corresponding: true affiliation: 1 affiliations: - name: CyberAgent, Inc., Japan