diff --git a/docs/source/_static/distributional_effect.png b/docs/source/_static/distributional_effect.png new file mode 100644 index 0000000..6be0e86 Binary files /dev/null and b/docs/source/_static/distributional_effect.png differ diff --git a/docs/source/api/local.rst b/docs/source/api/local.rst new file mode 100644 index 0000000..a8bbc00 --- /dev/null +++ b/docs/source/api/local.rst @@ -0,0 +1,22 @@ +Local Distribution Estimators +============================== + +This page documents local distribution treatment effect estimators that compute treatment effects weighted by treatment propensity within each stratum. These estimators are particularly useful for handling treatment assignment heterogeneity across strata. + +SimpleLocalDistributionEstimator +-------------------------------- + +.. autoclass:: dte_adj.SimpleLocalDistributionEstimator + :members: + :undoc-members: + :show-inheritance: + :no-index: + +AdjustedLocalDistributionEstimator +---------------------------------- + +.. autoclass:: dte_adj.AdjustedLocalDistributionEstimator + :members: + :undoc-members: + :show-inheritance: + :no-index: \ No newline at end of file diff --git a/docs/source/api/plot.rst b/docs/source/api/plot.rst new file mode 100644 index 0000000..a4390e2 --- /dev/null +++ b/docs/source/api/plot.rst @@ -0,0 +1,10 @@ +Plotting Utilities +================== + +This page documents plotting utilities for visualizing distribution treatment effects, confidence intervals, and other results. + +.. automodule:: dte_adj.plot + :members: + :undoc-members: + :show-inheritance: + :no-index: \ No newline at end of file diff --git a/docs/source/api/simple.rst b/docs/source/api/simple.rst new file mode 100644 index 0000000..77b111a --- /dev/null +++ b/docs/source/api/simple.rst @@ -0,0 +1,30 @@ +Simple Randomization Estimators +=============================== + +This page documents estimators that work with simple randomized experimental designs where treatment assignment is completely randomized. + +These estimators leverage pre-treatment covariates through distributional regression frameworks to improve the precision of distributional treatment effect estimates. The key methodological contribution is using machine learning techniques for variance reduction while maintaining validity as long as nuisance components are reasonably well estimated. + +Byambadalai et al. (2024) [#simple2024]_ propose a regression adjustment method that incorporates covariates into distributional regression, enabling deeper insights beyond average treatment effects by estimating full distributional treatment effects in randomized experiments. + +.. [#simple2024] Byambadalai, U., Oka, T., & Yasui, S. (2024). Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction. arXiv preprint `arXiv:2407.16037 `_. + +SimpleDistributionEstimator +--------------------------- + +.. autoclass:: dte_adj.SimpleDistributionEstimator + :members: + :inherited-members: + :undoc-members: + :show-inheritance: + :no-index: + +AdjustedDistributionEstimator +----------------------------- + +.. autoclass:: dte_adj.AdjustedDistributionEstimator + :members: + :inherited-members: + :undoc-members: + :show-inheritance: + :no-index: \ No newline at end of file diff --git a/docs/source/api/stratified.rst b/docs/source/api/stratified.rst new file mode 100644 index 0000000..3341f14 --- /dev/null +++ b/docs/source/api/stratified.rst @@ -0,0 +1,30 @@ +Covariate Adaptive Randomization Estimators +=========================================== + +This page documents estimators that work with stratified experimental designs, particularly for covariate-adaptive randomization (CAR) within strata. + +These estimators are designed to handle stratified block randomization where participants are grouped into strata based on baseline covariates before treatment assignment. The key methodological contribution is leveraging additional covariates beyond strata indicators using machine learning methods to enhance the precision of distributional treatment effect estimates. + +Byambadalai et al. (2025) [#car2025]_ propose a flexible distribution regression framework that achieves the semiparametric efficiency bound for distributional treatment effects under CAR, demonstrating that regression-adjusted estimators can optimally utilize covariate information in stratified designs. + +.. [#car2025] Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization. arXiv preprint `arXiv:2506.05945 `_. + +SimpleStratifiedDistributionEstimator +------------------------------------- + +.. autoclass:: dte_adj.SimpleStratifiedDistributionEstimator + :members: + :inherited-members: + :undoc-members: + :show-inheritance: + :no-index: + +AdjustedStratifiedDistributionEstimator +--------------------------------------- + +.. autoclass:: dte_adj.AdjustedStratifiedDistributionEstimator + :members: + :inherited-members: + :undoc-members: + :show-inheritance: + :no-index: \ No newline at end of file diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst new file mode 100644 index 0000000..7c0853c --- /dev/null +++ b/docs/source/api_reference.rst @@ -0,0 +1,36 @@ +API Reference +============= + +This section provides comprehensive documentation for all classes and functions in the dte_adj package. The API is organized into logical groups based on functionality and use cases. + +Overview +-------- + +The dte_adj package provides several types of estimators for computing distribution treatment effects: + +* **Simple Randomization Estimators**: For estimating distributional effects in simple randomized experiments where treatment assignment is independent of all covariates +* **Covariate Adaptive Randomization Estimators**: For estimating distributional effects under covariate-adaptive randomization (CAR) designs, including stratified block randomization and other adaptive schemes +* **Local Distribution Estimators**: For estimating local distribution treatment effects weighted by treatment propensity within strata +* **Utility Functions**: Helper functions for confidence intervals and statistical computations +* **Plotting Utilities**: Visualization tools for treatment effects and distributions + +For theoretical foundations, see Byambadalai et al. (2024) [#simple2024]_ for simple randomization and Byambadalai et al. (2025) [#car2025]_ for covariate-adaptive randomization. + +For multi-task learning approaches that train models for all locations simultaneously (using ``is_multi_task=True``), see the neural network framework in [#multitask2024]_. + +.. [#simple2024] Byambadalai, U., Oka, T., & Yasui, S. (2024). Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction. arXiv preprint `arXiv:2407.16037 `_. + +.. [#car2025] Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization. arXiv preprint `arXiv:2506.05945 `_. + +.. [#multitask2024] Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2024). Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks. arXiv preprint `arXiv:2507.07738 `_. + +Detailed Documentation +---------------------- + +.. toctree:: + :maxdepth: 2 + + api/simple + api/stratified + api/local + api/plot \ No newline at end of file diff --git a/docs/source/get_started.rst b/docs/source/get_started.rst index 38da5f8..4fba511 100644 --- a/docs/source/get_started.rst +++ b/docs/source/get_started.rst @@ -1,5 +1,5 @@ -Get Started with dte_adj -======================== +Get Started +=========== This page contains basic usage of dte_adj library. diff --git a/docs/source/index.rst b/docs/source/index.rst index da29d27..2c9bc87 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,8 +1,41 @@ dte_adj ======= -`dte_adj` is a Python package for estimating distribution treatment effects. -It provides APIs for conducting regression adjustment to estimate precise distribution functions as well as convenient utils. +A Python Package for Estimating Distribution Treatment Effects +-------------------------------------------------------------- + +`dte_adj` is a Python package for estimating distribution treatment effects in randomized experiments. +It provides APIs for conducting regression adjustment to estimate precise distribution functions, enabling deeper insights beyond average treatment effects through machine learning-enhanced estimation methods. + +.. image:: _static/distributional_effect.png + :alt: distributional effect + :height: 400px + :width: 700px + :align: center + +Estimator Types +--------------- + +The package provides several types of estimators for computing distribution treatment effects: + +* **Simple Randomization Estimators**: For estimating distributional effects in simple randomized experiments where treatment assignment is independent of all covariates +* **Covariate Adaptive Randomization Estimators**: For estimating distributional effects under covariate-adaptive randomization (CAR) designs, including stratified block randomization and other adaptive schemes +* **Local Distribution Estimators**: For estimating local distribution treatment effects weighted by treatment propensity within strata + +Theoretical Foundations +----------------------- + +For theoretical foundations, see: + +* **Simple randomization**: Byambadalai et al. (2024) [#simple2024]_ +* **Covariate-adaptive randomization**: Byambadalai et al. (2025) [#car2025]_ +* **Multi-task learning**: Byambadalai et al. (2024) [#multitask2024]_ + +.. [#simple2024] Byambadalai, U., Oka, T., & Yasui, S. (2024). Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction. arXiv preprint `arXiv:2407.16037 `_. + +.. [#car2025] Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization. arXiv preprint `arXiv:2506.05945 `_. + +.. [#multitask2024] Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2024). Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks. arXiv preprint `arXiv:2507.07738 `_. .. toctree:: :maxdepth: 1 @@ -10,7 +43,7 @@ It provides APIs for conducting regression adjustment to estimate precise distri installation get_started - modules + api_reference contributing Indices and tables diff --git a/docs/source/modules.rst b/docs/source/modules.rst deleted file mode 100644 index 941d8e9..0000000 --- a/docs/source/modules.rst +++ /dev/null @@ -1,16 +0,0 @@ -Modules -======= - -.. toctree:: - :maxdepth: 2 - -.. automodule:: dte_adj - :members: - :inherited-members: - :undoc-members: - :show-inheritance: - -.. automodule:: dte_adj.plot - :members: - :undoc-members: - :show-inheritance: \ No newline at end of file diff --git a/dte_adj/__init__.py b/dte_adj/__init__.py index ee4887b..527b92d 100644 --- a/dte_adj/__init__.py +++ b/dte_adj/__init__.py @@ -200,17 +200,17 @@ def predict_qte( import numpy as np from dte_adj import SimpleStratifiedDistributionEstimator - + # Generate stratified sample data X = np.random.randn(1000, 5) strata = np.random.choice([0, 1, 2], size=1000) D = np.random.binomial(1, 0.5, 1000) Y = X[:, 0] + 2 * D + 0.5 * strata + np.random.randn(1000) - + # Fit stratified estimator estimator = SimpleStratifiedDistributionEstimator() estimator.fit(X, D, Y, strata) - + # Compute QTE at specific quantiles quantiles = np.array([0.25, 0.5, 0.75]) # 25th, 50th, 75th percentiles qte, lower, upper = estimator.predict_qte( @@ -219,7 +219,7 @@ def predict_qte( quantiles=quantiles, n_bootstrap=100 ) - + print(f"QTE at quantiles {quantiles}: {qte}") print(f"Median effect (50th percentile): {qte[1]:.3f}") """ @@ -238,7 +238,7 @@ def predict_qte( qtes = np.zeros((n_bootstrap, qte.shape[0])) for b in range(n_bootstrap): bootstrap_indexes = np.random.choice(indexes, size=n_obs, replace=True) - + qtes[b] = self._compute_qtes( target_treatment_arm, control_treatment_arm, @@ -381,7 +381,7 @@ def find_quantile(quantile, arm): # Temporarily store original strata and use the provided strata original_strata = self.strata self.strata = strata - + val, _, _ = self._compute_cumulative_distribution( arm, np.full((1), locations[mid]), @@ -389,10 +389,10 @@ def find_quantile(quantile, arm): treatment_arms, outcomes, ) - + # Restore original strata self.strata = original_strata - + if val[0] <= quantile: result = locations[mid] low = mid + 1