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¶
- class dte_adj.SimpleLocalDistributionEstimator[source]
Bases:
SimpleStratifiedDistributionEstimator
A class for computing Local Distribution Treatment Effects (LDTE) and Local Probability Treatment Effects (LPTE) using simple empirical estimation.
This estimator computes treatment effects that are weighted by treatment propensity within each stratum, providing estimates that are locally robust to treatment assignment heterogeneity across strata. It uses empirical methods without ML adjustment.
- fit(covariates: ndarray, treatment_arms: ndarray, treatment_indicator: ndarray, outcomes: ndarray, strata: ndarray) SimpleLocalDistributionEstimator [source]
Train the SimpleLocalDistributionEstimator.
- Parameters:
covariates (np.ndarray) – Pre-treatment covariates.
treatment_arms (np.ndarray) – Treatment assignment variable (Z).
treatment_indicator (np.ndarray) – Treatment indicator variable (D).
outcomes (np.ndarray) – Scalar-valued observed outcome.
strata (np.ndarray) – Stratum indicators.
- Returns:
The fitted estimator.
- Return type:
SimpleLocalDistributionEstimator
- predict_ldte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) Tuple[ndarray, ndarray, ndarray] [source]
Compute Local Distribution Treatment Effects (LDTE).
LDTE measures the difference in cumulative distribution functions between treatment groups weighted by treatment propensity within each stratum. This provides estimates that are locally robust to treatment assignment heterogeneity across 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.
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.
- Returns:
- A tuple containing:
Expected LDTEs (np.ndarray): Local 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 sklearn.linear_model import LogisticRegression from dte_adj import AdjustedLocalDistributionEstimator # Generate sample data with strata np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1], size=1000) # Binary strata D = np.random.binomial(1, 0.3 + 0.4 * strata, 1000) # Treatment depends on strata Y = X[:, 0] + 2 * D + strata + np.random.randn(1000) # Fit local estimator base_model = LogisticRegression() estimator = AdjustedLocalDistributionEstimator(base_model) estimator.fit(X, D, D, Y, strata) # treatment_arms = treatment_indicator for binary case # Compute LDTE locations = np.linspace(Y.min(), Y.max(), 20) ldte, lower, upper = estimator.predict_ldte( target_treatment_arm=1, control_treatment_arm=0, locations=locations ) print(f"LDTE shape: {ldte.shape}") # Should match locations.shape print(f"Average LDTE: {ldte.mean():.3f}")
- predict_lpte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) Tuple[ndarray, ndarray, ndarray] [source]
Compute Local Probability Treatment Effects (LPTE).
LPTE measures the difference in probability mass between treatment groups for intervals weighted by treatment propensity within each stratum. This provides interval-based treatment effect estimates that are locally robust to treatment assignment heterogeneity.
- 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 LPTE is computed.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
- Returns:
- A tuple containing:
Expected LPTEs (np.ndarray): Local 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 sklearn.linear_model import LogisticRegression from dte_adj import SimpleLocalDistributionEstimator # Generate sample data with strata np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1, 2], size=1000) # Multiple strata D = np.random.binomial(1, 0.2 + 0.3 * (strata == 1) + 0.4 * (strata == 2), 1000) Y = X[:, 0] + 1.5 * D + 0.5 * strata + np.random.randn(1000) # Fit simple local estimator estimator = SimpleLocalDistributionEstimator() estimator.fit(X, D, D, Y, strata) # Define interval boundaries locations = np.array([-2, -1, 0, 1, 2]) # Creates 4 intervals # Compute LPTE lpte, lower, upper = estimator.predict_lpte( target_treatment_arm=1, control_treatment_arm=0, locations=locations ) print(f"LPTE shape: {lpte.shape}") # Should be (4,) for 4 intervals print(f"Interval effects: {lpte}")
AdjustedLocalDistributionEstimator¶
- class dte_adj.AdjustedLocalDistributionEstimator(base_model: Any, folds=3, is_multi_task=False)[source]
Bases:
AdjustedStratifiedDistributionEstimator
A class for computing Local Distribution Treatment Effects (LDTE) and Local Probability Treatment Effects (LPTE) using ML-adjusted estimation.
This estimator combines local treatment effect estimation with machine learning adjustment, providing treatment effects that are both locally robust to treatment assignment heterogeneity and adjusted for confounding through observed covariates. It uses cross-fitting for more precise estimates in complex treatment assignment scenarios.
- fit(covariates: ndarray, treatment_arms: ndarray, treatment_indicator: ndarray, outcomes: ndarray, strata: ndarray) AdjustedLocalDistributionEstimator [source]
Train the AdjustedLocalDistributionEstimator.
- Parameters:
covariates (np.ndarray) – Pre-treatment covariates.
treatment_arms (np.ndarray) – Treatment assignment variable (Z).
treatment_indicator (np.ndarray) – Treatment indicator variable (D).
outcomes (np.ndarray) – Scalar-valued observed outcome.
strata (np.ndarray) – Stratum indicators.
- Returns:
The fitted estimator.
- Return type:
AdjustedLocalDistributionEstimator
- predict_ldte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) Tuple[ndarray, ndarray, ndarray] [source]
Compute Local Distribution Treatment Effects (LDTE) with ML-adjusted estimation. LDTE measures the difference in cumulative distribution functions between treatment groups weighted by treatment propensity within each stratum, using ML models for adjustment. Currently, this API only supports analytical confidence interval.
- 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.
- Returns:
- A tuple containing:
Expected LDTEs (np.ndarray): ML-adjusted local treatment effect estimates
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 sklearn.ensemble import RandomForestClassifier from dte_adj import AdjustedLocalDistributionEstimator # Generate sample data with complex treatment assignment np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1], size=1000) # Treatment depends on covariates and strata treatment_prob = 0.2 + 0.3 * (X[:, 0] > 0) + 0.2 * strata D = np.random.binomial(1, treatment_prob, 1000) Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) # Fit adjusted local estimator base_model = RandomForestClassifier(n_estimators=50, random_state=42) estimator = AdjustedLocalDistributionEstimator(base_model, folds=3) estimator.fit(X, D, D, Y, strata) # Compute LDTE locations = np.linspace(Y.min(), Y.max(), 15) ldte, lower, upper = estimator.predict_ldte( target_treatment_arm=1, control_treatment_arm=0, locations=locations ) print(f"ML-adjusted LDTE shape: {ldte.shape}") print(f"Average LDTE: {ldte.mean():.3f}")
- predict_lpte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) Tuple[ndarray, ndarray, ndarray] [source]
Compute Local Probability Treatment Effects (LPTE) with ML-adjusted estimation. LPTE measures the difference in probability mass between treatment groups for intervals weighted by treatment propensity within each stratum, using ML models for adjustment. Currently, this API only supports analytical confidence interval.
- 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 LPTE is computed.
alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.
- Returns:
- A tuple containing:
Expected LPTEs (np.ndarray): ML-adjusted local treatment effect estimates, 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 sklearn.ensemble import GradientBoostingClassifier from dte_adj import AdjustedLocalDistributionEstimator # Generate sample data with confounding np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1, 2], size=1000) # Complex treatment assignment mechanism logit_score = X[:, 0] + 0.5 * X[:, 1] + strata treatment_prob = 1 / (1 + np.exp(-logit_score)) D = np.random.binomial(1, treatment_prob, 1000) Y = X.sum(axis=1) + 1.5 * D + 0.3 * strata + np.random.randn(1000) # Fit adjusted estimator with gradient boosting base_model = GradientBoostingClassifier(n_estimators=100, random_state=42) estimator = AdjustedLocalDistributionEstimator(base_model, folds=5) estimator.fit(X, D, D, Y, strata) # Define intervals and compute LPTE locations = np.array([-3, -1, 0, 1, 3]) # 4 intervals lpte, lower, upper = estimator.predict_lpte( target_treatment_arm=1, control_treatment_arm=0, locations=locations ) print(f"ML-adjusted LPTE shape: {lpte.shape}") # Should be (4,) print(f"Interval effects: {lpte}")