CyberAgentAILab · TomeHirata · Aug 20, 2026 · Aug 20, 2026
diff --git a/.github/workflows/build-doc.yml b/.github/workflows/build-doc.yml
@@ -29,7 +29,7 @@ jobs:
     - name: Build documentation
       run: |
         cd docs
-        uv run make html
+        uv run mkdocs build --strict
 
     - name: Deploy to GitHub Pages
       uses: peaceiris/actions-gh-pages@84c30a85c19949d7eee79c4ff27748b70285e453 # v4.1.0

diff --git a/.gitignore b/.gitignore
@@ -68,7 +68,8 @@ instance/
 # Scrapy stuff:
 .scrapy
 
-# Sphinx documentation
+# MkDocs documentation
+docs/build/
 docs/_build/
 
 # PyBuilder

diff --git a/_typos.toml b/_typos.toml
@@ -2,7 +2,12 @@
 # Allow dataset-specific column names from Hillstrom dataset
 womens = "womens"
 mens = "mens"
+# Domain term: subjects who comply with treatment assignment
+compliers = "compliers"
+complier = "complier"
 
 [default.extend-identifiers]
 # Author name in SciPy citation
 VanderPlas = "VanderPlas"
+# numpy function
+arange = "arange"
diff --git a/docs/Makefile b/docs/Makefile
diff --git a/docs/docs/api/index.md b/docs/docs/api/index.md
@@ -0,0 +1,29 @@
+# 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, Byambadalai et al. (2025)[^car2025] for covariate-adaptive randomization, and Byambadalai et al. (2025)[^compliance2025] for imperfect compliance scenarios.
+
+For multi-task learning approaches that train models for all locations simultaneously (using `is_multi_task=True`), see the neural network framework in [^multitask2025].
+
+## Detailed Documentation
+
+- [Simple Randomization Estimators](simple.md)
+- [Covariate Adaptive Randomization Estimators](stratified.md)
+- [Local Distribution Estimators](local.md)
+- [Plotting Utilities](plot.md)
+
+[^simple2024]: Byambadalai, U., Oka, T., & Yasui, S. (2024). Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction. In Proceedings of the 41st International Conference on Machine Learning (ICML'24). [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. In Proceedings of the 42nd International Conference on Machine Learning (ICML'25). [arXiv:2506.05945](https://arxiv.org/abs/2506.05945).
+[^multitask2025]: Hirata, T., Byambadalai, U., Oka, T., Yasui, S., & Uto, S. (2025). Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks. arXiv preprint [arXiv:2507.07738](https://arxiv.org/abs/2507.07738).
+[^compliance2025]: Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). Beyond the Average: Distributional Causal Inference under Imperfect Compliance. arXiv preprint [arXiv:2509.15594](https://arxiv.org/abs/2509.15594).
diff --git a/docs/source/api/local.rst → docs/docs/api/local.md b/docs/source/api/local.rst → docs/docs/api/local.md
@@ -1,26 +1,15 @@
-Local Distribution Estimators
-==============================
+# 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 and scenarios with imperfect compliance.
 
 Local distribution treatment effects (LDTE) and local probability treatment effects (LPTE) provide methods for causal inference that account for treatment assignment vs. treatment receipt differences. For theoretical foundations on imperfect compliance scenarios, see:
 
-* Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). *Beyond the Average: Distributional Causal Inference under Imperfect Compliance*. `arXiv:2509.15594 <https://arxiv.org/abs/2509.15594>`_.
+- Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). *Beyond the Average: Distributional Causal Inference under Imperfect Compliance*. [arXiv:2509.15594](https://arxiv.org/abs/2509.15594).
 
-SimpleLocalDistributionEstimator
---------------------------------
+## SimpleLocalDistributionEstimator
 
-.. autoclass:: dte_adj.SimpleLocalDistributionEstimator
-   :members:
-   :undoc-members:
-   :show-inheritance:
-   :no-index:
+::: dte_adj.SimpleLocalDistributionEstimator
 
-AdjustedLocalDistributionEstimator
-----------------------------------
+## AdjustedLocalDistributionEstimator
 
-.. autoclass:: dte_adj.AdjustedLocalDistributionEstimator
-   :members:
-   :undoc-members:
-   :show-inheritance:
-   :no-index:
+::: dte_adj.AdjustedLocalDistributionEstimator
diff --git a/docs/docs/api/plot.md b/docs/docs/api/plot.md
@@ -0,0 +1,5 @@
+# Plotting Utilities
+
+This page documents plotting utilities for visualizing distribution treatment effects, confidence intervals, and other results.
+
+::: dte_adj.plot
diff --git a/docs/docs/api/simple.md b/docs/docs/api/simple.md
@@ -0,0 +1,23 @@
+# 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.
+
+## SimpleDistributionEstimator
+
+::: dte_adj.SimpleDistributionEstimator
+    options:
+      inherited_members: true
+      show_inheritance_diagram: true
+
+## AdjustedDistributionEstimator
+
+::: dte_adj.AdjustedDistributionEstimator
+    options:
+      inherited_members: true
+      show_inheritance_diagram: true
+
+[^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).
diff --git a/docs/docs/api/stratified.md b/docs/docs/api/stratified.md
@@ -0,0 +1,21 @@
+# 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.
+
+## SimpleStratifiedDistributionEstimator
+
+::: dte_adj.SimpleStratifiedDistributionEstimator
+    options:
+      inherited_members: true
+
+## AdjustedStratifiedDistributionEstimator
+
+::: dte_adj.AdjustedStratifiedDistributionEstimator
+    options:
+      inherited_members: true
+
+[^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).
diff --git a/.../source/_static/distributional_effect.png → docs/docs/assets/distributional_effect.png b/.../source/_static/distributional_effect.png → docs/docs/assets/distributional_effect.png
diff --git a/docs/source/_static/dte_empirical.png → docs/docs/assets/dte_empirical.png b/docs/source/_static/dte_empirical.png → docs/docs/assets/dte_empirical.png
diff --git a/docs/source/_static/dte_moment.png → docs/docs/assets/dte_moment.png b/docs/source/_static/dte_moment.png → docs/docs/assets/dte_moment.png
diff --git a/docs/source/_static/dte_simple.png → docs/docs/assets/dte_simple.png b/docs/source/_static/dte_simple.png → docs/docs/assets/dte_simple.png
diff --git a/docs/source/_static/dte_uniform.png → docs/docs/assets/dte_uniform.png b/docs/source/_static/dte_uniform.png → docs/docs/assets/dte_uniform.png
diff --git a/docs/source/_static/hillstorm_dte.png → docs/docs/assets/hillstorm_dte.png b/docs/source/_static/hillstorm_dte.png → docs/docs/assets/hillstorm_dte.png
diff --git a/.../source/_static/hillstorm_dte_control.png → docs/docs/assets/hillstorm_dte_control.png b/.../source/_static/hillstorm_dte_control.png → docs/docs/assets/hillstorm_dte_control.png
diff --git a/docs/source/_static/hillstorm_pte.png → docs/docs/assets/hillstorm_pte.png b/docs/source/_static/hillstorm_pte.png → docs/docs/assets/hillstorm_pte.png
diff --git a/.../source/_static/hillstorm_pte_control.png → docs/docs/assets/hillstorm_pte_control.png b/.../source/_static/hillstorm_pte_control.png → docs/docs/assets/hillstorm_pte_control.png
diff --git a/...ce/_static/hillstorm_subgroup_men_dte.png → ...ocs/assets/hillstorm_subgroup_men_dte.png b/...ce/_static/hillstorm_subgroup_men_dte.png → ...ocs/assets/hillstorm_subgroup_men_dte.png
diff --git a/...ce/_static/hillstorm_subgroup_men_pte.png → ...ocs/assets/hillstorm_subgroup_men_pte.png b/...ce/_static/hillstorm_subgroup_men_pte.png → ...ocs/assets/hillstorm_subgroup_men_pte.png
diff --git a/.../_static/hillstorm_subgroup_women_dte.png → ...s/assets/hillstorm_subgroup_women_dte.png b/.../_static/hillstorm_subgroup_women_dte.png → ...s/assets/hillstorm_subgroup_women_dte.png
diff --git a/.../_static/hillstorm_subgroup_women_pte.png → ...s/assets/hillstorm_subgroup_women_pte.png b/.../_static/hillstorm_subgroup_women_pte.png → ...s/assets/hillstorm_subgroup_women_pte.png
diff --git a/.../_static/oregon_ldte_costs_comparison.png → ...s/assets/oregon_ldte_costs_comparison.png b/.../_static/oregon_ldte_costs_comparison.png → ...s/assets/oregon_ldte_costs_comparison.png
diff --git a/...urce/_static/oregon_ldte_costs_strata.png → .../docs/assets/oregon_ldte_costs_strata.png b/...urce/_static/oregon_ldte_costs_strata.png → .../docs/assets/oregon_ldte_costs_strata.png
diff --git a/..._static/oregon_ldte_visits_comparison.png → .../assets/oregon_ldte_visits_comparison.png b/..._static/oregon_ldte_visits_comparison.png → .../assets/oregon_ldte_visits_comparison.png
diff --git a/.../_static/oregon_lpte_costs_comparison.png → ...s/assets/oregon_lpte_costs_comparison.png b/.../_static/oregon_lpte_costs_comparison.png → ...s/assets/oregon_lpte_costs_comparison.png
diff --git a/..._static/oregon_lpte_visits_comparison.png → .../assets/oregon_lpte_visits_comparison.png b/..._static/oregon_lpte_visits_comparison.png → .../assets/oregon_lpte_visits_comparison.png
diff --git a/docs/source/_static/pte_empirical.png → docs/docs/assets/pte_empirical.png b/docs/source/_static/pte_empirical.png → docs/docs/assets/pte_empirical.png
diff --git a/docs/source/_static/qte.png → docs/docs/assets/qte.png b/docs/source/_static/qte.png → docs/docs/assets/qte.png
diff --git a/docs/docs/contributing.md b/docs/docs/contributing.md
@@ -0,0 +1,3 @@
+# Contribution Guide
+
+Regarding how to contribute to this package, please refer to [CONTRIBUTING.md](https://github.com/CyberAgentAILab/python-dte-adjustment/blob/main/CONTRIBUTING.md) for more details.
diff --git a/docs/docs/get_started.md b/docs/docs/get_started.md
@@ -0,0 +1,156 @@
+# Get Started
+
+This page contains basic usage of dte_adj library.
+
+Generate data for training cumulative distribution function:
+
+```python
+import numpy as np
+
+def generate_data(n, d_x=100, rho=0.5):
+    """
+    Generate data according to the described data generating process (DGP).
+
+    Args:
+    n (int): Number of samples.
+    d_x (int): Number of covariates. Default is 100.
+    rho (float): Success probability for the Bernoulli distribution. Default is 0.5.
+
+    Returns:
+    X (np.ndarray): Covariates matrix of shape (n, d_x).
+    D (np.ndarray): Treatment variable array of shape (n,).
+    Y (np.ndarray): Outcome variable array of shape (n,).
+    """
+    # Generate covariates X from a uniform distribution on (0, 1)
+    X = np.random.uniform(0, 1, (n, d_x))
+
+    # Generate treatment variable D from a Bernoulli distribution with success probability rho
+    D = np.random.binomial(1, rho, n)
+
+    # Define beta_j and gamma_j according to the problem statement
+    beta = np.zeros(d_x)
+    gamma = np.zeros(d_x)
+
+    # Set the first 50 values of beta and gamma to 1
+    beta[:50] = 1
+    gamma[:50] = 1
+
+    # Compute the outcome Y
+    U = np.random.normal(0, 1, n)  # Error term
+    linear_term = np.dot(X, beta)
+    quadratic_term = np.dot(X**2, gamma)
+
+    # Outcome equation
+    Y = 5 * D + linear_term + quadratic_term + U
+
+    return X, D, Y
+
+n = 1000  # Sample size
+X, D, Y = generate_data(n)
+```
+
+Then, let's build an empirical cumulative distribution function (CDF).
+
+```python
+import dte_adj
+from dte_adj.plot import plot
+
+estimator = dte_adj.SimpleDistributionEstimator()
+estimator.fit(X, D, Y)
+locations = np.linspace(Y.min(), Y.max(), 20)
+cdf = estimator.predict(1, locations)
+```
+
+Distributional treatment effect (DTE) can be computed easily in the following code.
+
+```python
+dte, lower_bound, upper_bound = estimator.predict_dte(target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="simple")
+```
+
+A convenience function is available to visualize distribution effects. This method can be used for other distribution parameters including Probability Treatment Effect (PTE) and Quantile Treatment Effect (QTE).
+
+```python
+plot(locations, dte, lower_bound, upper_bound, title="DTE of simple estimator")
+```
+
+![DTE of empirical estimator](assets/dte_empirical.png)
+
+To initialize the adjusted distribution function, the base model for conditional distribution function needs to be passed.
+In the following example, Logistic Regression is used. Please make sure that your base model implements `fit` and `predict_proba` methods.
+
+```python
+from sklearn.linear_model import LogisticRegression
+logit = LogisticRegression()
+estimator = dte_adj.AdjustedDistributionEstimator(logit, folds=3)
+estimator.fit(X, D, Y)
+cdf = estimator.predict(1, locations)
+```
+
+DTE can be computed and visualized in the following code.
+
+```python
+dte, lower_bound, upper_bound = estimator.predict_dte(target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="simple")
+plot(locations, dte, lower_bound, upper_bound, title="DTE of adjusted estimator with simple confidence band")
+```
+
+![DTE of adjusted estimator with simple confidence band](assets/dte_simple.png)
+
+Confidence bands can be computed in different ways. In the following code, moment condition is used to calculate the confidence band.
+
+```python
+dte, lower_bound, upper_bound = estimator.predict_dte(target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="moment")
+plot(locations, dte, lower_bound, upper_bound, title="DTE of adjusted estimator with moment confidence band")
+```
+
+![DTE of adjusted estimator with moment confidence band](assets/dte_moment.png)
+
+Also, an uniform confidence band is used when `uniform` is specified for the `variance_type` argument.
+
+```python
+dte, lower_bound, upper_bound = estimator.predict_dte(target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="uniform")
+plot(locations, dte, lower_bound, upper_bound, title="DTE of adjusted estimator with uniform confidence band")
+```
+
+![DTE of adjusted estimator with uniform confidence band](assets/dte_uniform.png)
+
+To compute PTE, you can use `predict_pte` method. The `locations` parameter defines interval boundaries, and the method returns probability treatment effects for each interval.
+For each interval, the starting point is not included but the ending point is included. For example, if the `locations` is [0, 1, 2], PTE is computed for `(0, 1]` and `(1, 2]`.
+
+```python
+pte, lower_bound, upper_bound = estimator.predict_pte(target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="simple")
+# Note: pte will have shape (len(locations)-1,) since it computes intervals between locations
+plot(locations[:-1], pte, lower_bound, upper_bound, chart_type="bar", title="PTE of adjusted estimator with simple confidence band")
+```
+
+![PTE of adjusted estimator with simple confidence band](assets/pte_empirical.png)
+
+To compute QTE, you can use `predict_qte` method. The confidence band is computed by bootstrap method.
+
+```python
+quantiles = np.array([0.1 * i for i in range(1, 10)], dtype=np.float32)
+qte, lower_bound, upper_bound = estimator.predict_qte(target_treatment_arm=1, control_treatment_arm=0, quantiles=quantiles, n_bootstrap=30)
+plot(quantiles, qte, lower_bound, upper_bound, title="QTE of adjusted estimator")
+```
+
+![QTE of adjusted estimator](assets/qte.png)
+
+You can use any model with `predict_proba` or `predict` method to adjust the distribution function estimation.
+For example, the following code use XGBoost classifier to estimate the conditional distribution.
+
+```python
+import xgboost as xgb
+estimator = dte_adj.AdjustedDistributionEstimator(xgb.XGBClassifier(), folds=3)
+estimator.fit(X, D, Y)
+cdf = estimator.predict(1, locations)
+```
+
+`predict_dte` and `predict_pte` methods provide an option to train a model for multiple locations simultaneously.
+To enable the feature, pass `is_multi_task=True`.
+
+```python
+from sklearn.linear_model import LinearRegression
+model = LinearRegression()
+estimator = dte_adj.AdjustedDistributionEstimator(model, folds=3)
+estimator.fit(X, D, Y)
+dte, lower_bound, upper_bound = estimator.predict_dte(target_treatment_arm=1, control_treatment_arm=0, is_multi_task=True, locations=locations, variance_type="moment")
+```
diff --git a/docs/docs/index.md b/docs/docs/index.md
@@ -0,0 +1,36 @@
+# dte_adj
+
+## 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.
+
+<p align="center">
+  <img src="assets/distributional_effect.png" alt="distributional effect" width="700">
+</p>
+
+## 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**: Hirata et al. (2025)[^multitask2025]
+- **Imperfect compliance**: Byambadalai et al. (2025)[^compliance2025]
+
+## License
+
+MIT License
+
+[^simple2024]: Byambadalai, U., Oka, T., & Yasui, S. (2024). Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction. In Proceedings of the 41st International Conference on Machine Learning (ICML'24). [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. In Proceedings of the 42nd International Conference on Machine Learning (ICML'25). [arXiv:2506.05945](https://arxiv.org/abs/2506.05945).
+[^multitask2025]: Hirata, T., Byambadalai, U., Oka, T., Yasui, S., & Uto, S. (2025). Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks. arXiv preprint [arXiv:2507.07738](https://arxiv.org/abs/2507.07738).
+[^compliance2025]: Byambadalai, U., Hirata, T., Oka, T., & Yasui, S. (2025). Beyond the Average: Distributional Causal Inference under Imperfect Compliance. arXiv preprint [arXiv:2509.15594](https://arxiv.org/abs/2509.15594).
diff --git a/docs/docs/installation.md b/docs/docs/installation.md
@@ -0,0 +1,25 @@
+# Installation Guide
+
+This package can be installed either through PyPI or source code.
+
+## Requirement
+
+You need to use Python version 3.10 or higher to use this package.
+
+## Install from PyPI
+
+For installing the package from PyPI, please use the following command.
+
+```bash
+pip install dte_adj
+```
+
+## Install from source code
+
+For installing the package from the source code, please use the following commands.
+
+```bash
+git clone https://github.com/CyberAgentAILab/python-dte-adjustment
+cd python-dte-adjustment
+pip install -e .
+```
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# Contribution Guide

		Regarding how to contribute to this package, please refer to [CONTRIBUTING.md](https://github.com/CyberAgentAILab/python-dte-adjustment/blob/main/CONTRIBUTING.md) for more details.