diff --git a/docs/docs/get_started.md b/docs/docs/get_started.md index 0ece1e1..3fde366 100644 --- a/docs/docs/get_started.md +++ b/docs/docs/get_started.md @@ -75,6 +75,9 @@ plot(locations, dte, lower_bound, upper_bound, title="DTE of simple estimator") ![DTE of empirical estimator](assets/dte_empirical.png) +!!! note "Input requirements" + `treatment_arms`, `outcomes` and `strata` must not contain missing values (NaN); `fit` raises a `ValueError` otherwise. They may be 1-D arrays or single columns of shape `(n, 1)`. Missing values in `covariates` are not checked: simple estimators ignore covariates, while for adjusted estimators it depends on whether your base model accepts them (e.g. scikit-learn's `HistGradientBoostingClassifier` does, `LogisticRegression` does not). + 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. @@ -135,6 +138,8 @@ 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. + +The adjustment is meant for randomized experiments: it reduces the variance of the estimates (tighter confidence intervals) and does not correct for confounding. Cross-fitting assigns folds at random, so with very small samples a fold can end up with no training data for a treatment arm; in that case a `ValueError` is raised and you should reduce `folds` (e.g. `folds=2`). For example, the following code use XGBoost classifier to estimate the conditional distribution. ```python diff --git a/docs/docs/tutorials/oregon.md b/docs/docs/tutorials/oregon.md index c9fd6c3..b417cb1 100644 --- a/docs/docs/tutorials/oregon.md +++ b/docs/docs/tutorials/oregon.md @@ -531,7 +531,7 @@ plt.show() - Simple: LDTE ≈ -0.55 at zero costs, converging to zero around $15,000-$20,000 - ML-Adjusted: LDTE ≈ -0.10 to -0.15 at zero costs, stable pattern with improved confidence intervals - Much larger magnitude effects in the Simple estimator, indicating households with multiple members show substantially stronger treatment effects - - ML adjustment provides more conservative estimates, potentially controlling for confounding household characteristics + - ML adjustment provides more conservative estimates and tighter confidence intervals by exploiting household characteristics that predict the outcome (the lottery itself randomizes treatment, so this is variance reduction rather than confounding correction) **2. Heterogeneity Across Strata** @@ -540,7 +540,7 @@ The stratified analysis reveals substantial treatment effect heterogeneity: - **"Signed self up" stratum**: Moderate effects (LDTE ≈ -0.18 to -0.20), suggesting single-person households have more modest increases in ED utilization - **"Signed self up + others" stratum**: Large effects in Simple estimator (LDTE ≈ -0.55), suggesting multi-person households experience much greater increases in ED access when not adjusting for covariates - The 3-4x larger effect in the "signed self up + others" group (Simple estimator) indicates that household composition is a critical moderator of insurance impact -- However, ML adjustment substantially reduces this estimate, suggesting that some of the observed effect may be attributable to observable household characteristics rather than pure treatment effects +- However, ML adjustment substantially reduces this estimate. Because treatment is randomized, the adjusted and unadjusted estimators target the same quantity, so a large gap is more likely sampling noise in the small stratum (reduced by the adjustment) than a difference in what is being estimated **3. Comparison of Estimation Methods** @@ -551,14 +551,14 @@ The stratified analysis reveals substantial treatment effect heterogeneity: - The Simple estimator shows the largest treatment effects across all strata (LDTE ≈ -0.55) - ML adjustment substantially reduces the estimated effect and stabilizes confidence intervals - This divergence suggests that observable covariates (e.g., household size, age composition, baseline health status) explain a significant portion of the treatment effect heterogeneity - - The improved stability of ML-adjusted estimates indicates successful control for confounding factors that may have been correlated with both treatment assignment and outcomes + - The improved stability of ML-adjusted estimates reflects variance reduction from covariates that are predictive of the outcome **4. Practical Implications** - **Household structure matters**: Multi-person households show substantially larger treatment effects in unadjusted analyses, likely because insurance coverage enables care-seeking for multiple family members. -- **The role of covariates**: The difference between Simple and ML-adjusted estimates in the "signed self up + others" stratum highlights the importance of controlling for household characteristics. The unadjusted effect may overstate the pure treatment effect by conflating insurance provision with pre-existing household differences. +- **The role of covariates**: The difference between Simple and ML-adjusted estimates in the "signed self up + others" stratum highlights how much precision household characteristics can add. The unadjusted estimate is noisier in this small stratum and may overstate the effect by chance. - **Stratification reveals hidden heterogeneity**: The overall population estimate masks substantial variation across household types, demonstrating the value of subgroup analysis. -- **Model specification considerations**: ML adjustment improves estimation stability in smaller strata and provides more defensible causal estimates by controlling for observable confounders. The convergence of all estimates to zero at higher cost levels confirms that the treatment primarily affects the lower tail of the cost distribution. +- **Model specification considerations**: ML adjustment improves estimation stability in smaller strata and provides more precise estimates by exploiting predictive covariates (it is not a correction for confounding, since treatment is randomized). The convergence of all estimates to zero at higher cost levels confirms that the treatment primarily affects the lower tail of the cost distribution. ### Conclusion diff --git a/dte_adj/base.py b/dte_adj/base.py index dba2477..e4abefb 100644 --- a/dte_adj/base.py +++ b/dte_adj/base.py @@ -418,30 +418,31 @@ def _compute_qtes( locations = np.sort(outcomes) def find_quantile(quantile, arm): + # Smallest location y with F(y) >= quantile, i.e. the generalized inverse + # of the estimated CDF. Returns the largest location if no such y exists. low, high = 0, locations.shape[0] - 1 - result = -1 - while low <= high: - mid = (low + high) // 2 - # 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 + result = locations[-1] + # Temporarily use the provided strata when evaluating the CDF + original_strata = self.strata + self.strata = strata + try: + while low <= high: + mid = (low + high) // 2 + val, _, _ = self._compute_cumulative_distribution( + arm, + np.full((1), locations[mid]), + covariates, + treatment_arms, + outcomes, + ) + + if val[0] >= quantile - 1e-12: + result = locations[mid] + high = mid - 1 + else: + low = mid + 1 + finally: self.strata = original_strata - - if val[0] <= quantile: - result = locations[mid] - low = mid + 1 - else: - high = mid - 1 return result result = np.zeros(quantiles.shape) diff --git a/dte_adj/local.py b/dte_adj/local.py index 6e65917..f1a8b59 100644 --- a/dte_adj/local.py +++ b/dte_adj/local.py @@ -10,7 +10,8 @@ ArrayLike, compute_ldte, compute_lpte, - _convert_to_ndarray, + _to_1d, + _check_no_missing, _infer_default_locations, ) @@ -55,8 +56,13 @@ def fit( Returns: SimpleLocalDistributionEstimator: The fitted estimator. """ - treatment_indicator = _convert_to_ndarray(treatment_indicator) + treatment_indicator = _to_1d("treatment_indicator", treatment_indicator) super().fit(covariates, treatment_arms, outcomes, strata) + if treatment_indicator.shape[0] != self.covariates.shape[0]: + raise ValueError( + "The shape of covariates and treatment_indicator should be same" + ) + _check_no_missing("treatment_indicator", treatment_indicator) self.treatment_indicator = treatment_indicator return self @@ -219,9 +225,10 @@ class AdjustedLocalDistributionEstimator(AdjustedStratifiedDistributionEstimator A class for computing Local Distribution Treatment Effects (LDTE) and Local Probability Treatment Effects (LPTE) using machine learning adjustment. - This estimator combines the benefits of ML adjustment with local treatment effect estimation, - providing precise estimates of treatment effects that are weighted by treatment propensity - within each stratum. It uses cross-fitting to avoid overfitting issues. + This estimator combines ML adjustment with local treatment effect estimation, providing + variance-reduced estimates in randomized experiments with partial compliance. The ML + adjustment targets efficiency (tighter confidence intervals), not correction for + confounding. It uses cross-fitting to avoid overfitting issues. """ def fit( @@ -245,8 +252,13 @@ def fit( Returns: AdjustedLocalDistributionEstimator: The fitted estimator. """ - treatment_indicator = _convert_to_ndarray(treatment_indicator) + treatment_indicator = _to_1d("treatment_indicator", treatment_indicator) super().fit(covariates, treatment_arms, outcomes, strata) + if treatment_indicator.shape[0] != self.covariates.shape[0]: + raise ValueError( + "The shape of covariates and treatment_indicator should be same" + ) + _check_no_missing("treatment_indicator", treatment_indicator) self.treatment_indicator = treatment_indicator return self @@ -288,13 +300,12 @@ def predict_ldte( from sklearn.ensemble import RandomForestClassifier from dte_adj import AdjustedLocalDistributionEstimator - # Generate confounded data with strata + # Generate data with strata from a randomized experiment np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1], size=1000) - # Treatment assignment depends on covariates - Z_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1] + strata))) - Z = np.random.binomial(1, Z_prob, 1000) + # Treatment assignment is random (independent of covariates) + Z = np.random.binomial(1, 0.5, 1000) D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) @@ -366,13 +377,12 @@ def predict_lpte( from sklearn.linear_model import LogisticRegression from dte_adj import AdjustedLocalDistributionEstimator - # Generate confounded data with strata + # Generate data with strata from a randomized experiment np.random.seed(42) X = np.random.randn(1000, 5) strata = np.random.choice([0, 1], size=1000) - # Treatment assignment depends on covariates - Z_prob = 1 / (1 + np.exp(-(X[:, 0] + strata))) - Z = np.random.binomial(1, Z_prob, 1000) + # Treatment assignment is random (independent of covariates) + Z = np.random.binomial(1, 0.5, 1000) D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) diff --git a/dte_adj/simple.py b/dte_adj/simple.py index 1e1a6ca..a4f53a0 100644 --- a/dte_adj/simple.py +++ b/dte_adj/simple.py @@ -5,7 +5,7 @@ SimpleStratifiedDistributionEstimator, AdjustedStratifiedDistributionEstimator, ) -from dte_adj.util import ArrayLike, _convert_to_ndarray +from dte_adj.util import ArrayLike, _prepare_fit_inputs class SimpleDistributionEstimator(SimpleStratifiedDistributionEstimator): @@ -15,8 +15,8 @@ class SimpleDistributionEstimator(SimpleStratifiedDistributionEstimator): 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. + It provides a baseline approach for randomized experiments where covariate adjustment + is not needed. Example: ```python @@ -61,20 +61,14 @@ def fit( Returns: SimpleDistributionEstimator: The fitted estimator. """ - covariates = _convert_to_ndarray(covariates) - treatment_arms = _convert_to_ndarray(treatment_arms) - outcomes = _convert_to_ndarray(outcomes) - - if covariates.shape[0] != treatment_arms.shape[0]: - raise ValueError("The shape of covariates and treatment_arm should be same") - - if covariates.shape[0] != outcomes.shape[0]: - raise ValueError("The shape of covariates and outcome should be same") + covariates, treatment_arms, outcomes, strata = _prepare_fit_inputs( + covariates, treatment_arms, outcomes + ) self.covariates = covariates self.treatment_arms = treatment_arms self.outcomes = outcomes - self.strata = np.zeros(len(self.covariates)) + self.strata = strata return self @@ -83,10 +77,18 @@ class AdjustedDistributionEstimator(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. + This estimator uses cross-fitting with ML models of the conditional distribution given + covariates to compute Distribution Treatment Effects (DTE), Probability Treatment Effects + (PTE), and Quantile Treatment Effects (QTE) with reduced variance. It is designed for + randomized experiments, where treatment is assigned independently of the covariates: in + that setting the estimator stays consistent regardless of how well the ML model fits, and + a good fit yields tighter confidence intervals than ``SimpleDistributionEstimator``. + + Note: + The method targets efficiency, not bias correction. It does not correct for + confounding in observational data; if treatment assignment depends on covariates, the + estimates are valid only under selection on observables, and ``fit`` assumes the + assignment probability is constant across observations (no propensity weighting). Example: ```python @@ -94,10 +96,10 @@ class AdjustedDistributionEstimator(AdjustedStratifiedDistributionEstimator): from sklearn.ensemble import RandomForestClassifier from dte_adj import AdjustedDistributionEstimator - # Generate confounded data + # Generate data from a randomized experiment: treatment is independent of X, + # while the outcome depends on X (this is what the ML adjustment exploits) X = np.random.randn(1000, 5) - treatment_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1]))) - D = np.random.binomial(1, treatment_prob, 1000) + D = np.random.binomial(1, 0.5, 1000) Y = X.sum(axis=1) + 2 * D + np.random.randn(1000) # Fit adjusted estimator @@ -125,19 +127,13 @@ def fit( Returns: AdjustedDistributionEstimator: The fitted estimator. """ - covariates = _convert_to_ndarray(covariates) - treatment_arms = _convert_to_ndarray(treatment_arms) - outcomes = _convert_to_ndarray(outcomes) - - if covariates.shape[0] != treatment_arms.shape[0]: - raise ValueError("The shape of covariates and treatment_arm should be same") - - if covariates.shape[0] != outcomes.shape[0]: - raise ValueError("The shape of covariates and outcome should be same") + covariates, treatment_arms, outcomes, strata = _prepare_fit_inputs( + covariates, treatment_arms, outcomes + ) self.covariates = covariates self.treatment_arms = treatment_arms self.outcomes = outcomes - self.strata = np.zeros(len(self.covariates)) + self.strata = strata return self diff --git a/dte_adj/stratified.py b/dte_adj/stratified.py index a4038e0..646170f 100644 --- a/dte_adj/stratified.py +++ b/dte_adj/stratified.py @@ -5,7 +5,11 @@ from copy import deepcopy from tqdm.auto import tqdm from dte_adj.base import DistributionEstimatorBase -from dte_adj.util import ArrayLike, _convert_to_ndarray +from dte_adj.util import ( + ArrayLike, + _prepare_fit_inputs, + _check_folds_have_training_data, +) class SimpleStratifiedDistributionEstimator(DistributionEstimatorBase): @@ -30,16 +34,9 @@ def fit( Returns: DistributionEstimatorBase: The fitted estimator. """ - covariates = _convert_to_ndarray(covariates) - treatment_arms = _convert_to_ndarray(treatment_arms) - outcomes = _convert_to_ndarray(outcomes) - strata = _convert_to_ndarray(strata) - - if covariates.shape[0] != treatment_arms.shape[0]: - raise ValueError("The shape of covariates and treatment_arm should be same") - - if covariates.shape[0] != outcomes.shape[0]: - raise ValueError("The shape of covariates and outcome should be same") + covariates, treatment_arms, outcomes, strata = _prepare_fit_inputs( + covariates, treatment_arms, outcomes, strata + ) self.covariates = covariates self.treatment_arms = treatment_arms @@ -155,7 +152,12 @@ def _compute_interval_probability( class AdjustedStratifiedDistributionEstimator(DistributionEstimatorBase): - """A class is for estimating the adjusted distribution function and computing the Distributional parameters for CAR.""" + """A class is for estimating the adjusted distribution function and computing the Distributional parameters for CAR. + + The ML adjustment (cross-fitted conditional distribution models) is intended to reduce + the variance of estimates in randomized experiments, where treatment assignment is + independent of covariates within strata. It does not correct for confounding. + """ def __init__(self, base_model: Any, folds=3, is_multi_task=False): """ @@ -163,7 +165,9 @@ def __init__(self, base_model: Any, folds=3, is_multi_task=False): Args: base_model (scikit-learn estimator): The base model implementing used for conditional distribution function estimators. The model should implement fit(data, targets) and predict_proba(data). - folds (int): The number of folds for cross-fitting. + folds (int): The number of folds for cross-fitting. Folds are assigned at random, so + with small samples a fold may by chance leave no training observations for the + target arm; a ValueError suggesting fewer folds is raised in that case. is_multi_task(bool): Whether to use multi-task learning. If True, your base model needs to support multi-task prediction (n_samples, n_features) -> (n_samples, n_targets). Returns: @@ -199,16 +203,9 @@ def fit( Returns: DistributionEstimatorBase: The fitted estimator. """ - covariates = _convert_to_ndarray(covariates) - treatment_arms = _convert_to_ndarray(treatment_arms) - outcomes = _convert_to_ndarray(outcomes) - strata = _convert_to_ndarray(strata) - - if covariates.shape[0] != treatment_arms.shape[0]: - raise ValueError("The shape of covariates and treatment_arm should be same") - - if covariates.shape[0] != outcomes.shape[0]: - raise ValueError("The shape of covariates and outcome should be same") + covariates, treatment_arms, outcomes, strata = _prepare_fit_inputs( + covariates, treatment_arms, outcomes, strata + ) self.covariates = covariates self.treatment_arms = treatment_arms @@ -249,6 +246,7 @@ def _compute_cumulative_distribution( prediction = np.zeros((n_records, n_loc)) treatment_mask = treatment_arms == target_treatment_arm folds = np.random.randint(self.folds, size=n_records) + _check_folds_have_training_data(folds, self.folds, treatment_mask) strata = self.strata s_list = np.unique(strata) if self.is_multi_task: @@ -360,6 +358,7 @@ def _compute_interval_probability( prediction = np.zeros((n_records, n_loc - 1)) treatment_mask = treatment_arms == target_treatment_arm folds = np.random.randint(self.folds, size=n_records) + _check_folds_have_training_data(folds, self.folds, treatment_mask) strata = self.strata s_list = np.unique(strata) binominals = (outcomes[:, np.newaxis] <= locations) * 1 # (n_records, n_loc) diff --git a/dte_adj/util.py b/dte_adj/util.py index 4a68d03..0b518c4 100644 --- a/dte_adj/util.py +++ b/dte_adj/util.py @@ -32,6 +32,79 @@ def _convert_to_ndarray(data: ArrayLike) -> np.ndarray: return np.asarray(data) +def _to_1d(name: str, data: ArrayLike) -> np.ndarray: + """Convert to a 1-D ndarray, accepting a single-column 2-D input as well.""" + arr = _convert_to_ndarray(data) + if arr.ndim == 2 and arr.shape[1] == 1: + arr = arr[:, 0] + if arr.ndim != 1: + raise ValueError( + f"{name} must be a 1-D array (or a single column), got shape {arr.shape}" + ) + return arr + + +def _check_no_missing(name: str, arr: np.ndarray) -> None: + """Raise if a float array contains missing values (NaN).""" + if arr.dtype.kind in "fc" and np.isnan(arr).any(): + raise ValueError( + f"{name} must not contain missing values (NaN). " + "Drop or impute the affected observations before calling fit." + ) + + +def _prepare_fit_inputs( + covariates: ArrayLike, + treatment_arms: ArrayLike, + outcomes: ArrayLike, + strata: ArrayLike = None, +): + """Convert and validate the inputs shared by every ``fit`` method. + + ``treatment_arms``, ``outcomes`` and ``strata`` are flattened to 1-D (a single + column such as shape ``(n, 1)`` is accepted) and must not contain NaN. + ``covariates`` are passed through unchanged; whether missing values in them are + acceptable depends on the base model used by adjusted estimators. If ``strata`` + is None, all observations are placed in a single stratum. + """ + covariates = _convert_to_ndarray(covariates) + treatment_arms = _to_1d("treatment_arms", treatment_arms) + outcomes = _to_1d("outcomes", outcomes) + + if covariates.shape[0] != treatment_arms.shape[0]: + raise ValueError("The shape of covariates and treatment_arm should be same") + + if covariates.shape[0] != outcomes.shape[0]: + raise ValueError("The shape of covariates and outcome should be same") + + if strata is None: + strata = np.zeros(covariates.shape[0]) + else: + strata = _to_1d("strata", strata) + if covariates.shape[0] != strata.shape[0]: + raise ValueError("The shape of covariates and strata should be same") + + _check_no_missing("treatment_arms", treatment_arms) + _check_no_missing("outcomes", outcomes) + _check_no_missing("strata", strata) + return covariates, treatment_arms, outcomes, strata + + +def _check_folds_have_training_data( + folds: np.ndarray, n_folds: int, treatment_mask: np.ndarray +) -> None: + """Raise an informative error if cross-fitting would train on no data.""" + for fold in range(n_folds): + if not ((folds != fold) & treatment_mask).any(): + raise ValueError( + f"Cross-fitting produced a fold ({fold} of {n_folds}) whose " + "complementary training set contains no observations of the target " + "treatment arm. This can happen by chance when the sample (or the " + "treatment arm) is small relative to the number of folds. " + "Reduce `folds` (e.g. folds=2) or use more data." + ) + + def _infer_default_locations( outcomes: np.ndarray, for_intervals: bool = False, diff --git a/tests/test_distribution_estimator_base.py b/tests/test_distribution_estimator_base.py index 50a6b23..fc7d2e7 100644 --- a/tests/test_distribution_estimator_base.py +++ b/tests/test_distribution_estimator_base.py @@ -161,7 +161,15 @@ def test_predict_qte(self): target_treatment_arm = 1 control_treatment_arm = 0 quantiles = np.array([0.1 * i for i in range(1, 10)]) - expected_qte = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + expected_qte = np.ones(9) + + # CDF of outcomes 0..19: arm 0 is (y + 1) / 20, arm 1 is y / 20 (shifted by one) + def shifted_cdf(arm, locations, covariates, treatment_arms, outcomes): + n = outcomes.shape[0] + cdf = np.clip((locations + (1 - arm)) / 20, 0, 1) + return cdf, np.zeros((n, len(locations))), np.zeros((n, len(locations))) + + self.estimator.compute_cumulative_distribution.side_effect = shifted_cdf # Act qte, lower_bound, upper_bound = self.estimator.predict_qte( diff --git a/tests/test_simple_estimator.py b/tests/test_simple_estimator.py index abcd862..1531b4e 100644 --- a/tests/test_simple_estimator.py +++ b/tests/test_simple_estimator.py @@ -2,7 +2,11 @@ import numpy as np from unittest.mock import patch, MagicMock from sklearn.linear_model import LogisticRegression -from dte_adj import SimpleDistributionEstimator, AdjustedDistributionEstimator +from dte_adj import ( + SimpleDistributionEstimator, + SimpleStratifiedDistributionEstimator, + AdjustedDistributionEstimator, +) np.random.seed(123) @@ -285,3 +289,64 @@ def test_e2e(self): ), "Adjusted estimator does not have narrower intervals", ) + + +class TestInputShapes(unittest.TestCase): + def setUp(self): + self.X = np.zeros((4, 1)) + self.D = np.array([0, 0, 1, 1]) + self.Y = np.array([0.0, 1.0, 2.0, 3.0]) + + def test_single_column_outcomes_match_1d(self): + results = [] + for outcomes in (self.Y, self.Y[:, None]): + est = SimpleDistributionEstimator().fit(self.X, self.D, outcomes) + effect, lower, upper = est.predict_dte( + 1, 0, np.array([1.0]), display_progress=False + ) + results.append((effect, lower, upper)) + for a, b in zip(*results): + self.assertEqual(a.shape, (1,)) + np.testing.assert_allclose(a, b) + np.testing.assert_allclose(results[1][0], [-1.0]) + + def test_single_column_treatment_arms_and_strata(self): + est = SimpleDistributionEstimator().fit(self.X, self.D[:, None], self.Y) + self.assertEqual(est.treatment_arms.shape, (4,)) + est = SimpleStratifiedDistributionEstimator().fit( + self.X, self.D, self.Y, np.zeros((4, 1)) + ) + self.assertEqual(est.strata.shape, (4,)) + + def test_multi_column_outcomes_rejected(self): + with self.assertRaises(ValueError): + SimpleDistributionEstimator().fit(self.X, self.D, np.zeros((4, 2))) + + def test_missing_values_rejected(self): + Y = self.Y.copy() + Y[0] = np.nan + with self.assertRaises(ValueError) as cm: + SimpleDistributionEstimator().fit(self.X, self.D, Y) + self.assertIn("outcomes", str(cm.exception)) + with self.assertRaises(ValueError): + SimpleDistributionEstimator().fit( + self.X, np.array([0, 1, np.nan, 1]), self.Y + ) + + def test_covariates_with_nan_allowed_for_simple(self): + X = self.X.copy() + X[0, 0] = np.nan + SimpleDistributionEstimator().fit(X, self.D, self.Y) + + +class TestFoldValidation(unittest.TestCase): + def test_empty_training_fold_raises_informative_error(self): + X = np.zeros((6, 2)) + D = np.array([0, 0, 0, 1, 1, 1]) + Y = np.arange(6.0) + est = AdjustedDistributionEstimator(MagicMock(), folds=2).fit(X, D, Y) + # All treated units fall in fold 0 -> training set for fold 0 is empty + with patch("numpy.random.randint", return_value=np.array([1, 1, 1, 0, 0, 0])): + with self.assertRaises(ValueError) as cm: + est.predict(1, np.array([2.0]), display_progress=False) + self.assertIn("Reduce `folds`", str(cm.exception)) diff --git a/tests/test_statistical_baselines.py b/tests/test_statistical_baselines.py new file mode 100644 index 0000000..0a414c7 --- /dev/null +++ b/tests/test_statistical_baselines.py @@ -0,0 +1,129 @@ +"""Statistical-correctness tests against hand-computed and simulated baselines.""" + +import unittest +import numpy as np +from sklearn.linear_model import LogisticRegression + +from dte_adj import ( + SimpleDistributionEstimator, + SimpleStratifiedDistributionEstimator, + AdjustedDistributionEstimator, +) + + +class TestKnownBaselines(unittest.TestCase): + def test_qte_location_shift(self): + # Y = 5 * D: every quantile shifts by exactly 5 + D = np.repeat([0, 1], 20) + X = np.zeros((40, 1)) + Y = 5.0 * D + est = SimpleDistributionEstimator().fit(X, D, Y) + qte, _, _ = est.predict_qte( + 1, + 0, + quantiles=np.array([0.1, 0.5, 0.9]), + n_bootstrap=5, + display_progress=False, + ) + np.testing.assert_allclose(qte, [5.0, 5.0, 5.0], rtol=0, atol=1e-10) + + def test_qte_hand_computed(self): + # control sorted: 1..10, treated sorted: 2,4,...,20 (a scale-by-2 effect) + D = np.repeat([0, 1], 10) + X = np.zeros((20, 1)) + Y = np.concatenate([np.arange(1, 11), 2 * np.arange(1, 11)]).astype(float) + est = SimpleDistributionEstimator().fit(X, D, Y) + qte, _, _ = est.predict_qte( + 1, + 0, + quantiles=np.array([0.1, 0.5, 0.9]), + n_bootstrap=5, + display_progress=False, + ) + # q-th quantile is the smallest y with F(y) >= q: control 1,5,9; treated 2,10,18 + np.testing.assert_allclose(qte, [1.0, 5.0, 9.0]) + + def test_dte_pte_hand_computed(self): + D = np.array([0, 0, 0, 0, 1, 1, 1, 1]) + X = np.zeros((8, 1)) + Y = np.array([1.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0, 6.0]) + est = SimpleDistributionEstimator().fit(X, D, Y) + locations = np.array([0.0, 2.0, 4.0, 6.0]) + dte, _, _ = est.predict_dte(1, 0, locations, display_progress=False) + # F_1 = [0, 0, .5, 1], F_0 = [0, .5, 1, 1] + np.testing.assert_allclose(dte, [0.0, -0.5, -0.5, 0.0]) + pte, _, _ = est.predict_pte(1, 0, locations, display_progress=False) + np.testing.assert_allclose(pte, [-0.5, 0.0, 0.5]) + + def test_stratified_equals_simple_for_single_stratum(self): + rng = np.random.default_rng(0) + X = rng.normal(size=(200, 2)) + D = rng.binomial(1, 0.5, 200) + Y = rng.normal(size=200) + D + locations = np.linspace(-2, 3, 8) + a = SimpleDistributionEstimator().fit(X, D, Y) + b = SimpleStratifiedDistributionEstimator().fit(X, D, Y, np.zeros(200)) + np.testing.assert_allclose( + a.predict_dte(1, 0, locations, display_progress=False)[0], + b.predict_dte(1, 0, locations, display_progress=False)[0], + ) + + +class TestSimulatedDGP(unittest.TestCase): + """Normal location-shift DGP with true DTE(y) = F_1(y) - F_0(y) = Phi((y - tau) / sd) - Phi(y / sd).""" + + @classmethod + def setUpClass(cls): + from scipy.stats import norm + + rng = np.random.default_rng(1) + n, cls.tau = 4000, 1.0 + cls.X = rng.normal(size=(n, 3)) + cls.D = rng.binomial(1, 0.5, n) + cls.Y = cls.X[:, 0] + cls.tau * cls.D + rng.normal(size=n) + # Y | D=d ~ N(d * tau, 2) + sd = np.sqrt(2.0) + cls.locations = np.array([-1.0, 0.0, 1.0, 2.0]) + cls.true_dte = norm.cdf((cls.locations - cls.tau) / sd) - norm.cdf( + cls.locations / sd + ) + cls.norm = norm + + def _assert_covers(self, dte, lo, hi): + # Pointwise 95% bands miss at each location with prob. ~5%, so check the truth + # lies within a band widened to ~3 standard errors, and the estimate is accurate. + half = (hi - lo) / 2 + widened = half * self.norm.ppf(0.9987) / self.norm.ppf(0.975) + self.assertTrue(np.all(np.abs(dte - self.true_dte) <= widened)) + + def test_simple_dte_covers_truth(self): + est = SimpleDistributionEstimator().fit(self.X, self.D, self.Y) + dte, lo, hi = est.predict_dte(1, 0, self.locations, display_progress=False) + self._assert_covers(dte, lo, hi) + + def test_simple_qte_close_to_truth(self): + est = SimpleDistributionEstimator().fit(self.X, self.D, self.Y) + qte, _, _ = est.predict_qte( + 1, + 0, + quantiles=np.array([0.25, 0.5, 0.75]), + n_bootstrap=20, + display_progress=False, + ) + np.testing.assert_allclose(qte, [self.tau] * 3, atol=0.2) + + def test_adjusted_dte_covers_truth_with_tighter_interval(self): + np.random.seed(0) + simple = SimpleDistributionEstimator().fit(self.X, self.D, self.Y) + adj = AdjustedDistributionEstimator(LogisticRegression(), folds=3).fit( + self.X, self.D, self.Y + ) + _, s_lo, s_hi = simple.predict_dte(1, 0, self.locations, display_progress=False) + dte, a_lo, a_hi = adj.predict_dte(1, 0, self.locations, display_progress=False) + self._assert_covers(dte, a_lo, a_hi) + # Y depends on X[:, 0], so adjustment should reduce the CI width + self.assertLess((a_hi - a_lo).mean(), (s_hi - s_lo).mean()) + + +if __name__ == "__main__": + unittest.main()