Overhaul doc site

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:

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:

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 <https://arxiv.org/abs/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:

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 <https://arxiv.org/abs/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:

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 <https://arxiv.org/abs/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 <https://arxiv.org/abs/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 <https://arxiv.org/abs/2507.07738>`_.
+
+Detailed Documentation
+----------------------
+
+.. toctree::
+   :maxdepth: 2
+
+   api/simple
+   api/stratified
+   api/local
+   api/plot

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.
 

docs/source/index.rst

@@ -1,16 +1,49 @@
 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 <https://arxiv.org/abs/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 <https://arxiv.org/abs/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 <https://arxiv.org/abs/2507.07738>`_.
 
 .. toctree::
    :maxdepth: 1
    :caption: Contents:
 
    installation
    get_started
-   modules
+   api_reference
    contributing
 
 Indices and tables

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,18 +381,18 @@ 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]),
                     covariates,
                     treatment_arms,
                     outcomes,
                 )
-                
+
                 # Restore original strata
                 self.strata = original_strata
-                
+
                 if val[0] <= quantile:
                     result = locations[mid]
                     low = mid + 1