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) [1] 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¶
- class dte_adj.SimpleDistributionEstimator[source]
Bases:
SimpleStratifiedDistributionEstimator
A class for computing the empirical distribution function and distributional treatment effects using simple (unadjusted) estimation methods.
This estimator computes Distribution Treatment Effects (DTE), Probability Treatment Effects (PTE), and Quantile Treatment Effects (QTE) without using machine learning models for adjustment. It provides a baseline approach suitable when treatment assignment is random or when covariate adjustment is not needed.
Example
import numpy as np from dte_adj import SimpleDistributionEstimator # Generate sample data X = np.random.randn(1000, 5) D = np.random.binomial(1, 0.5, 1000) # Random treatment Y = X[:, 0] + 2 * D + np.random.randn(1000) # Fit simple estimator estimator = SimpleDistributionEstimator() estimator.fit(X, D, Y) # Compute treatment effects locations = np.linspace(Y.min(), Y.max(), 20) dte, lower, upper = estimator.predict_dte(1, 0, locations) pte, pte_lower, pte_upper = estimator.predict_pte(1, 0, locations)
- fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) SimpleDistributionEstimator [source]
Set parameters.
- Parameters:
covariates (np.ndarray) – Pre-treatment covariates.
treatment_arms (np.ndarray) – The index of the treatment arm.
outcomes (np.ndarray) – Scalar-valued observed outcome.
- Returns:
The fitted estimator.
- Return type:
SimpleDistributionEstimator
- predict(treatment_arm: int, locations: ndarray) ndarray
Compute cumulative distribution values.
- Parameters:
treatment_arm (int) – The index of the treatment arm.
outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.
- Returns:
Estimated cumulative distribution values for the input.
- Return type:
np.ndarray
- predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.
The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
locations (np.ndarray) – Scalar values to be used for computing the cumulative distribution.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are “moment”, “simple”, and “uniform”. Defaults to “moment”.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected DTEs (np.ndarray): Treatment effect estimates at each location
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
import numpy as np from dte_adj import SimpleDistributionEstimator # Generate sample data X = np.random.randn(1000, 5) D = np.random.binomial(1, 0.5, 1000) Y = X[:, 0] + 2 * D + np.random.randn(1000) # Fit estimator estimator = SimpleDistributionEstimator() estimator.fit(X, D, Y) # Compute DTE locations = np.linspace(Y.min(), Y.max(), 20) dte, lower, upper = estimator.predict_dte( target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="moment" ) print(f"DTE shape: {dte.shape}") # Should match locations.shape print(f"Average DTE: {dte.mean():.3f}")
- predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.
The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
locations (np.ndarray) – Scalar values defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are “moment”, “simple”, and “uniform”. Defaults to “moment”.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
import numpy as np from dte_adj import SimpleDistributionEstimator # Generate sample data X = np.random.randn(1000, 5) D = np.random.binomial(1, 0.5, 1000) Y = X[:, 0] + 2 * D + np.random.randn(1000) # Fit estimator estimator = SimpleDistributionEstimator() estimator.fit(X, D, Y) # Define interval boundaries locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] # Compute PTE pte, lower, upper = estimator.predict_pte( target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="moment" ) print(f"PTE shape: {pte.shape}") # Should be (4,) for 4 intervals print(f"Interval effects: {pte}")
- predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.
The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
quantiles (np.ndarray, optional) – Quantiles used for QTE. Defaults to [0.1, 0.2, …, 0.9].
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected QTEs (np.ndarray): Treatment effect estimates at each quantile
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
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( target_treatment_arm=1, control_treatment_arm=0, quantiles=quantiles, n_bootstrap=100 ) print(f"QTE at quantiles {quantiles}: {qte}") print(f"Median effect (50th percentile): {qte[1]:.3f}")
AdjustedDistributionEstimator¶
- class dte_adj.AdjustedDistributionEstimator(base_model: Any, folds=3, is_multi_task=False)[source]
Bases:
AdjustedStratifiedDistributionEstimator
A class for computing distribution treatment effects using machine learning adjustment.
This estimator uses cross-fitting with ML models to adjust for confounding when computing Distribution Treatment Effects (DTE), Probability Treatment Effects (PTE), and Quantile Treatment Effects (QTE). It provides more precise estimates when treatment assignment depends on observed covariates.
Example
import numpy as np from sklearn.ensemble import RandomForestClassifier from dte_adj import AdjustedDistributionEstimator # Generate confounded data X = np.random.randn(1000, 5) treatment_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1]))) D = np.random.binomial(1, treatment_prob, 1000) Y = X.sum(axis=1) + 2 * D + np.random.randn(1000) # Fit adjusted estimator base_model = RandomForestClassifier(n_estimators=100) estimator = AdjustedDistributionEstimator(base_model, folds=3) estimator.fit(X, D, Y) # Compute adjusted treatment effects locations = np.linspace(Y.min(), Y.max(), 20) dte, lower, upper = estimator.predict_dte(1, 0, locations, variance_type="moment")
- fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) AdjustedDistributionEstimator [source]
Set parameters.
- Parameters:
covariates (np.ndarray) – Pre-treatment covariates.
treatment_arms (np.ndarray) – The index of the treatment arm.
outcomes (np.ndarray) – Scalar-valued observed outcome.
- Returns:
The fitted estimator.
- Return type:
AdjustedDistributionEstimator
- predict(treatment_arm: int, locations: ndarray) ndarray
Compute cumulative distribution values.
- Parameters:
treatment_arm (int) – The index of the treatment arm.
outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.
- Returns:
Estimated cumulative distribution values for the input.
- Return type:
np.ndarray
- predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.
The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
locations (np.ndarray) – Scalar values to be used for computing the cumulative distribution.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are “moment”, “simple”, and “uniform”. Defaults to “moment”.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected DTEs (np.ndarray): Treatment effect estimates at each location
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
import numpy as np from dte_adj import SimpleDistributionEstimator # Generate sample data X = np.random.randn(1000, 5) D = np.random.binomial(1, 0.5, 1000) Y = X[:, 0] + 2 * D + np.random.randn(1000) # Fit estimator estimator = SimpleDistributionEstimator() estimator.fit(X, D, Y) # Compute DTE locations = np.linspace(Y.min(), Y.max(), 20) dte, lower, upper = estimator.predict_dte( target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="moment" ) print(f"DTE shape: {dte.shape}") # Should match locations.shape print(f"Average DTE: {dte.mean():.3f}")
- predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.
The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
locations (np.ndarray) – Scalar values defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are “moment”, “simple”, and “uniform”. Defaults to “moment”.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
import numpy as np from dte_adj import SimpleDistributionEstimator # Generate sample data X = np.random.randn(1000, 5) D = np.random.binomial(1, 0.5, 1000) Y = X[:, 0] + 2 * D + np.random.randn(1000) # Fit estimator estimator = SimpleDistributionEstimator() estimator.fit(X, D, Y) # Define interval boundaries locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] # Compute PTE pte, lower, upper = estimator.predict_pte( target_treatment_arm=1, control_treatment_arm=0, locations=locations, variance_type="moment" ) print(f"PTE shape: {pte.shape}") # Should be (4,) for 4 intervals print(f"Interval effects: {pte}")
- predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) Tuple[ndarray, ndarray, ndarray]
Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.
The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.
- Parameters:
target_treatment_arm (int) – The index of the treatment arm of the treatment group.
control_treatment_arm (int) – The index of the treatment arm of the control group.
quantiles (np.ndarray, optional) – Quantiles used for QTE. Defaults to [0.1, 0.2, …, 0.9].
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.
- Returns:
- A tuple containing:
Expected QTEs (np.ndarray): Treatment effect estimates at each quantile
Lower bounds (np.ndarray): Lower confidence interval bounds
Upper bounds (np.ndarray): Upper confidence interval bounds
- Return type:
Tuple[np.ndarray, np.ndarray, np.ndarray]
Example
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( target_treatment_arm=1, control_treatment_arm=0, quantiles=quantiles, n_bootstrap=100 ) print(f"QTE at quantiles {quantiles}: {qte}") print(f"Median effect (50th percentile): {qte[1]:.3f}")