From dcbf22433bd9b88a49a764a1808e3287ed865159 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Sat, 11 Oct 2025 14:44:22 +0200 Subject: [PATCH 1/3] Enhance TMLE estimators with confidence interval calculations and variance handling This commit introduces significant improvements to the TMLE estimators, focusing on the addition of confidence interval calculations and variance estimation. Key changes include: - **New Variance Module**: Added a new `variance.py` module to compute standard errors and 95% confidence intervals for ATE, ATT, and RR using influence curves. - **Updated TMLE Functions**: Modified `compute_tmle_ate`, `compute_tmle_att`, and `compute_tmle_rr` to return confidence interval results alongside existing estimates. - **Refactoring**: Improved the structure of the TMLE functions to enhance clarity and maintainability. - **Constants Addition**: Introduced new constants for standard error and confidence interval bounds in `constants.py`. - **Testing Enhancements**: Added comprehensive tests for confidence interval coverage under various scenarios, ensuring robustness against model misspecification. These changes enhance the statistical reliability of the TMLE estimators, providing users with more informative outputs regarding the uncertainty of their estimates. --- CausalEstimate/core/multi_estimator.py | 9 +- CausalEstimate/estimators/functional/tmle.py | 54 ++++- .../estimators/functional/tmle_att.py | 70 +++--- .../estimators/functional/variance.py | 108 +++++++++ CausalEstimate/estimators/tmle.py | 1 - CausalEstimate/utils/constants.py | 4 + tests/helpers/setup.py | 170 ++++++------- .../test_tmle/test_coverage.py | 223 ++++++++++++++++++ 8 files changed, 512 insertions(+), 127 deletions(-) create mode 100644 CausalEstimate/estimators/functional/variance.py create mode 100644 tests/test_functional/test_tmle/test_coverage.py diff --git a/CausalEstimate/core/multi_estimator.py b/CausalEstimate/core/multi_estimator.py index 3be4404..d35201d 100644 --- a/CausalEstimate/core/multi_estimator.py +++ b/CausalEstimate/core/multi_estimator.py @@ -15,6 +15,9 @@ INITIAL_EFFECT_untreated, ADJUSTMENT_treated, ADJUSTMENT_untreated, + STD_ERR, + CI95_LOWER, + CI95_UPPER, ) @@ -158,9 +161,9 @@ def _compute_bootstrap( summary: Dict[str, Any] = { EFFECT: mean_effect, - "std_err": std_err, - "CI95_lower": ci95_lower, - "CI95_upper": ci95_upper, + STD_ERR: std_err, + CI95_LOWER: ci95_lower, + CI95_UPPER: ci95_upper, } other_keys = [key for key in result_keys if key != EFFECT] diff --git a/CausalEstimate/estimators/functional/tmle.py b/CausalEstimate/estimators/functional/tmle.py index 364c6d2..8502675 100644 --- a/CausalEstimate/estimators/functional/tmle.py +++ b/CausalEstimate/estimators/functional/tmle.py @@ -9,6 +9,7 @@ compute_initial_effect, estimate_fluctuation_parameter, ) +from CausalEstimate.estimators.functional.variance import compute_ci from CausalEstimate.utils.constants import ( EFFECT, EFFECT_treated, @@ -29,16 +30,29 @@ def compute_tmle_ate( """ Estimate the ATE using TMLE, with optional weight clipping. """ - Q_star_1, Q_star_0 = compute_estimates( + Q_star_1, Q_star_0, Yhat_star, H = compute_estimates( A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile, eps=eps ) ate = (Q_star_1 - Q_star_0).mean() + ci_results = compute_ci( + effect_type="ATE", + psi=ate, + Q_star_1=Q_star_1, + Q_star_0=Q_star_0, + Y=Y, + A=A, + ps=ps, + Yhat_star=Yhat_star, + H=H, + ) + return { EFFECT: ate, EFFECT_treated: Q_star_1.mean(), EFFECT_untreated: Q_star_0.mean(), **compute_initial_effect(Y1_hat, Y0_hat, Q_star_1, Q_star_0), + **ci_results, } @@ -55,7 +69,7 @@ def compute_tmle_rr( """ Estimate the Risk Ratio using TMLE, with optional weight clipping. """ - Q_star_1, Q_star_0 = compute_estimates( + Q_star_1, Q_star_0, Yhat_star, H = compute_estimates( A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile, eps=eps ) Q_star_1_m = Q_star_1.mean() @@ -75,11 +89,24 @@ def compute_tmle_rr( ) rr = np.inf + ci_results = compute_ci( + effect_type="RR", + psi=rr, + Q_star_1=Q_star_1, + Q_star_0=Q_star_0, + Y=Y, + A=A, + ps=ps, + Yhat_star=Yhat_star, + H=H, + ) + return { EFFECT: rr, EFFECT_treated: Q_star_1_m, EFFECT_untreated: Q_star_0_m, **compute_initial_effect(Y1_hat, Y0_hat, Q_star_1, Q_star_0, rr=True), + **ci_results, } @@ -92,15 +119,23 @@ def compute_estimates( Yhat: np.ndarray, clip_percentile: float = 1, eps: float = 1e-9, -) -> Tuple[np.ndarray, np.ndarray]: +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Compute updated outcome estimates using TMLE targeting step. + Returns: + Q_star_1: Updated outcome estimates under treatment + Q_star_0: Updated outcome estimates under control + Yhat_star: Targeted predictions Q*(A,W) + H: Clever covariate """ H = compute_clever_covariate_ate(A, ps, clip_percentile=clip_percentile, eps=eps) epsilon = estimate_fluctuation_parameter(H, Y, Yhat) - Q_star_1, Q_star_0 = update_estimates(ps, Y0_hat, Y1_hat, epsilon) + Q_star_1, Q_star_0 = update_estimates(ps, Y0_hat, Y1_hat, epsilon, eps=eps) - return Q_star_1, Q_star_0 + Yhat_clipped = np.clip(Yhat, eps, 1 - eps) + Yhat_star = expit(logit(Yhat_clipped) + epsilon * H) + + return Q_star_1, Q_star_0, Yhat_star, H def update_estimates( @@ -112,13 +147,14 @@ def update_estimates( ) -> Tuple[np.ndarray, np.ndarray]: """ Update the initial outcome estimates using the fluctuation parameter. - eps: float = 1e-9, - Guard against division by zero + Returns: + Q_star_1: Updated outcome estimates under treatment + Q_star_0: Updated outcome estimates under control """ H1 = 1.0 / (ps + eps) H0 = -1.0 / (1.0 - ps + eps) - Q_star_1 = expit(logit(Y1_hat) + epsilon * H1) - Q_star_0 = expit(logit(Y0_hat) + epsilon * H0) + Q_star_1 = expit(logit(np.clip(Y1_hat, eps, 1 - eps)) + epsilon * H1) + Q_star_0 = expit(logit(np.clip(Y0_hat, eps, 1 - eps)) + epsilon * H0) return Q_star_1, Q_star_0 diff --git a/CausalEstimate/estimators/functional/tmle_att.py b/CausalEstimate/estimators/functional/tmle_att.py index 18032bf..923fd80 100644 --- a/CausalEstimate/estimators/functional/tmle_att.py +++ b/CausalEstimate/estimators/functional/tmle_att.py @@ -10,10 +10,11 @@ from scipy.special import expit, logit from CausalEstimate.estimators.functional.utils import ( - compute_initial_effect, compute_clever_covariate_att, + compute_initial_effect, estimate_fluctuation_parameter, ) +from CausalEstimate.estimators.functional.variance import compute_ci from CausalEstimate.utils.constants import EFFECT, EFFECT_treated, EFFECT_untreated @@ -26,32 +27,20 @@ def compute_estimates_att( Yhat: np.ndarray, clip_percentile: float = 1, eps: float = 1e-9, -) -> Tuple[np.ndarray, np.ndarray]: +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Compute updated outcome estimates for ATT using a one-step TMLE targeting step. """ - # Estimate the fluctuation parameter epsilon using a logistic regression: H = compute_clever_covariate_att(A, ps, clip_percentile=clip_percentile, eps=eps) epsilon = estimate_fluctuation_parameter(H, Y, Yhat) - # --- Step 2: Define the CORRECT, separate update terms --- - # This is the part that was incorrect in your new code. We revert to the logic - # from your old implementation. p_treated = np.mean(A == 1) - if ( - p_treated == 0 - ): # Should be caught by compute_clever_covariate_att but good practice - return Y1_hat, Y0_hat + if p_treated == 0: + Yhat_star = Yhat.copy() # No update if no treated + return Y1_hat, Y0_hat, Yhat_star, H - # The update term for the potential outcome under treatment, Q(1,W). - # This is a scalar value applied to everyone's Y1_hat. + # Update terms update_term_1 = epsilon * (1.0 / (p_treated + eps)) - - # The update term for the potential outcome under control, Q(0,W). - # This is a vector of values applied to everyone's Y0_hat. - # We must re-calculate the weight component here. - # For theoretical consistency, if ps were clipped to find H, they should be clipped here too. - weight_component = ps / (p_treated * (1 - ps) + eps) if clip_percentile < 1: @@ -59,16 +48,18 @@ def compute_estimates_att( if control_mask.sum() > 0: control_weights = weight_component[control_mask] threshold = np.percentile(control_weights, clip_percentile * 100) - # Clip the component for ALL subjects based on the threshold from controls weight_component = np.clip(weight_component, a_min=None, a_max=threshold) update_term_0 = -epsilon * weight_component - # --- Step 3: Apply the separate updates to the potential outcome models --- - Q_star_1 = expit(logit(Y1_hat) + update_term_1) - Q_star_0 = expit(logit(Y0_hat) + update_term_0) + # Apply updates + Q_star_1 = expit(logit(np.clip(Y1_hat, eps, 1 - eps)) + update_term_1) + Q_star_0 = expit(logit(np.clip(Y0_hat, eps, 1 - eps)) + update_term_0) - return Q_star_1, Q_star_0 + Yhat_clipped = np.clip(Yhat, eps, 1 - eps) + Yhat_star = expit(logit(Yhat_clipped) + epsilon * H) + + return Q_star_1, Q_star_0, Yhat_star, H def compute_tmle_att( @@ -82,22 +73,35 @@ def compute_tmle_att( eps: float = 1e-9, ) -> dict: """ - Estimate the Average Treatment Effect on the Treated (ATT) using TMLE, - with optional clipping for the control group. - eps: float = 1e-9, - Guard against division by zero + Estimate the Average Treatment Effect on the Treated (ATT) using TMLE. """ - Q_star_1, Q_star_0 = compute_estimates_att( + Q_star_1, Q_star_0, Yhat_star, H = compute_estimates_att( A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile, eps=eps ) - # The final ATT parameter is the mean difference within the treated population - psi = np.mean(Q_star_1[A == 1] - Q_star_0[A == 1]) + treated_mask = A == 1 + if not np.any(treated_mask): + # Handle case with no treated subjects + return {EFFECT: np.nan, EFFECT_treated: np.nan, EFFECT_untreated: np.nan} + + psi = np.mean(Q_star_1[treated_mask] - Q_star_0[treated_mask]) + + ci_results = compute_ci( + effect_type="ATT", + psi=psi, + Q_star_1=Q_star_1, + Q_star_0=Q_star_0, + Y=Y, + A=A, + ps=ps, + Yhat_star=Yhat_star, + H=H, + ) return { EFFECT: psi, - # For clarity, return the mean of the updated predictions - EFFECT_treated: np.mean(Q_star_1[A == 1]), - EFFECT_untreated: np.mean(Q_star_0[A == 1]), + EFFECT_treated: np.mean(Q_star_1[treated_mask]), + EFFECT_untreated: np.mean(Q_star_0[treated_mask]), **compute_initial_effect(Y1_hat, Y0_hat, Q_star_1, Q_star_0), + **ci_results, } diff --git a/CausalEstimate/estimators/functional/variance.py b/CausalEstimate/estimators/functional/variance.py new file mode 100644 index 0000000..273dc89 --- /dev/null +++ b/CausalEstimate/estimators/functional/variance.py @@ -0,0 +1,108 @@ +import numpy as np + +from CausalEstimate.utils.constants import CI95_LOWER, CI95_UPPER, STD_ERR + + +def compute_ci( + effect_type: str, + psi: float, + Q_star_1: np.ndarray, + Q_star_0: np.ndarray, + Y: np.ndarray, + A: np.ndarray, + ps: np.ndarray, + Yhat_star: np.ndarray, + H: np.ndarray = None, +) -> dict: + """ + Compute the standard deviation and 95% confidence interval using the influence curve. + """ + n = len(Y) + if n == 0: + return {STD_ERR: np.nan, CI95_LOWER: np.nan, CI95_UPPER: np.nan} + + # Select the appropriate influence curve based on the effect type + if effect_type in ["ATE", "ARR"]: + ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, A, Yhat_star, H) + elif effect_type == "ATT": + p_treated = np.mean(A) + ic = _compute_ic_att(psi, Q_star_1, Q_star_0, Y, A, Yhat_star, H, p_treated) + elif effect_type == "RR": + ic = _compute_ic_rr(Q_star_1, Q_star_0, Y, A, ps) + else: + raise ValueError( + f"CI calculation for effect type '{effect_type}' is not supported." + ) + + if np.any(np.isnan(ic)): + return {STD_ERR: np.nan, CI95_LOWER: np.nan, CI95_UPPER: np.nan} + + # Compute variance and standard error + var_ic = np.var(ic, ddof=1) # Use ddof=1 for sample variance + std_err = np.sqrt(var_ic / n) + + # Compute confidence interval + if effect_type == "RR": + # For RR, CIs are calculated on the log scale and then exponentiated + log_psi = np.log(psi) + ci_lower = np.exp(log_psi - 1.96 * std_err) + ci_upper = np.exp(log_psi + 1.96 * std_err) + else: # ATE, ATT, ARR + ci_lower = psi - 1.96 * std_err + ci_upper = psi + 1.96 * std_err + + return {STD_ERR: std_err, CI95_LOWER: ci_lower, CI95_UPPER: ci_upper} + + +def _compute_ic_ate( + psi: float, + Q_star_1: np.ndarray, + Q_star_0: np.ndarray, + Y: np.ndarray, + A: np.ndarray, + Yhat_star: np.ndarray, + H: np.ndarray, +) -> np.ndarray: + """Influence curve for ATE.""" + return H * (Y - Yhat_star) + (Q_star_1 - Q_star_0) - psi + + +def _compute_ic_att( + psi: float, + Q_star_1: np.ndarray, + Q_star_0: np.ndarray, + Y: np.ndarray, + A: np.ndarray, + Yhat_star: np.ndarray, + H: np.ndarray, + p_treated: float, +) -> np.ndarray: + """Influence curve for ATT.""" + if p_treated == 0: + return np.full_like(Y, np.nan) + ic = H * (Y - Yhat_star) + (A / p_treated) * (Q_star_1 - Q_star_0 - psi) + return ic + + +def _compute_ic_rr( + Q_star_1: np.ndarray, + Q_star_0: np.ndarray, + Y: np.ndarray, + A: np.ndarray, + ps: np.ndarray, + eps: float = 1e-9, +) -> np.ndarray: + """Influence curve for log(Risk Ratio).""" + mu1_star = np.mean(Q_star_1) + mu0_star = np.mean(Q_star_0) + + if np.isclose(mu0_star, 0) or np.isclose(mu1_star, 0): + return np.full_like(Y, np.nan) + + # IC for mu1 + ic_mu1 = (A / (ps + eps)) * (Y - Q_star_1) + Q_star_1 - mu1_star + # IC for mu0 + ic_mu0 = ((1 - A) / (1 - ps + eps)) * (Y - Q_star_0) + Q_star_0 - mu0_star + + ic_log_rr = (1 / mu1_star) * ic_mu1 - (1 / mu0_star) * ic_mu0 + return ic_log_rr diff --git a/CausalEstimate/estimators/tmle.py b/CausalEstimate/estimators/tmle.py index cce7d7c..d50e5b4 100644 --- a/CausalEstimate/estimators/tmle.py +++ b/CausalEstimate/estimators/tmle.py @@ -1,4 +1,3 @@ -# CausalEstimate/estimators/tmle.py import pandas as pd from CausalEstimate.estimators.base import BaseEstimator diff --git a/CausalEstimate/utils/constants.py b/CausalEstimate/utils/constants.py index 45ff013..afbb5b5 100644 --- a/CausalEstimate/utils/constants.py +++ b/CausalEstimate/utils/constants.py @@ -23,3 +23,7 @@ INITIAL_EFFECT_untreated = "initial_effect_0" ADJUSTMENT_treated = "adjustment_1" ADJUSTMENT_untreated = "adjustment_0" + +STD_ERR = "std_err" +CI95_LOWER = "CI95_lower" +CI95_UPPER = "CI95_upper" diff --git a/tests/helpers/setup.py b/tests/helpers/setup.py index 2484403..223e54e 100644 --- a/tests/helpers/setup.py +++ b/tests/helpers/setup.py @@ -1,7 +1,8 @@ import unittest -from typing import List +from typing import Any, Dict, List import numpy as np +import pandas as pd from scipy.special import expit from CausalEstimate.simulation.binary_simulation import ( @@ -21,96 +22,103 @@ ) -class TestEffectBase(unittest.TestCase): +def generate_simulation_data( + n: int, alpha: List[float], beta: List[float], noise_level: float, seed: int +) -> Dict[str, Any]: """ - Base class for testing causal effect estimators. - - The TRUE data generating process (DGP) is controlled by the full `alpha` and - `beta` vectors. Child classes can override these to create different DGPs. + Generates a single, fresh dataset for a simulation run. + This is the logic extracted from TestEffectBase.setUpClass. + """ + rng = np.random.default_rng(seed) + + # 1. Simulate data using the true DGP + data = simulate_binary_data(n, alpha=alpha, beta=beta, seed=seed) + X_raw = data[["X1", "X2"]].values + A = data[TREATMENT_COL].values + Y = data[OUTCOME_COL].values + + # 2. Generate nuisance predictions using the assumed (misspecified) model + ps_model_coeffs = np.array(alpha[:3]) + X_ps_design = np.column_stack([np.ones(n), X_raw]) + ps_logit = X_ps_design @ ps_model_coeffs + noise_level * rng.normal(size=n) + ps = expit(ps_logit) + + outcome_model_coeffs = np.array(beta[:4]) + X_y1_design = np.column_stack([np.ones(n), np.ones(n), X_raw]) + X_y0_design = np.column_stack([np.ones(n), np.zeros(n), X_raw]) + X_y_obs_design = np.column_stack([np.ones(n), A, X_raw]) + + Y1_hat = expit( + X_y1_design @ outcome_model_coeffs + noise_level * rng.normal(size=n) + ) + Y0_hat = expit( + X_y0_design @ outcome_model_coeffs + noise_level * rng.normal(size=n) + ) + Yhat = expit( + X_y_obs_design @ outcome_model_coeffs + noise_level * rng.normal(size=n) + ) + + # 3. Finalize data and compute true values + eps = 1e-7 + + return { + "A": A, + "Y": Y, + "ps": np.clip(ps, eps, 1 - eps), + "Y1_hat": np.clip(Y1_hat, eps, 1 - eps), + "Y0_hat": np.clip(Y0_hat, eps, 1 - eps), + "Yhat": np.clip(Yhat, eps, 1 - eps), + "true_ate": compute_ATE_theoretical_from_data(data, beta=beta), + "true_att": compute_ATT_theoretical_from_data(data, beta=beta), + "true_rr": compute_RR_theoretical_from_data(data, beta=beta), + } - - alpha: [intercept, X1, X2, X1*X2 interaction] - - beta: [intercept, A, X1, X2, X1*X2 interaction] - The predictions for ps, Y1_hat, and Y0_hat are generated from an *assumed model* - that only uses the main effects (the first 3 or 4 coefficients). This is done - intentionally, so that if a child class sets a non-zero interaction term in - `alpha` or `beta`, the model used for predictions becomes misspecified. - The default values are set to a correctly specified DGP. +class TestEffectBase(unittest.TestCase): + """ + Base class for single-run tests. It now uses the centralized + generate_simulation_data function to create its fixture. """ n: int = 30_000 - alpha: List[float] = [0.1, 0.2, -0.3] - beta: List[float] = [0.5, 0.8, -0.6, 0.3] - noise_level: float = 0 # logit - cutoff_epsilon: float = 1e-7 + alpha: list = [0.1, 0.2, -0.3] + beta: list = [0.5, 0.8, -0.6, 0.3] + noise_level: float = 0 seed: int = 41 @classmethod def setUpClass(cls): - # Simulate realistic data for testing - rng = np.random.default_rng(cls.seed) - - # 1. Simulate data using the full coefficient vectors (the TRUE DGP). - # Child classes override cls.alpha/cls.beta to change this DGP. - data = simulate_binary_data( - cls.n, alpha=cls.alpha, beta=cls.beta, seed=cls.seed + """ + Generate a single, large dataset to be used as a fixture + for all the simple, single-run estimator tests. + """ + sim_data = generate_simulation_data( + n=cls.n, + alpha=cls.alpha, + beta=cls.beta, + noise_level=cls.noise_level, + seed=cls.seed, ) - X_raw = data[["X1", "X2"]].values - A = data[TREATMENT_COL].values - Y = data[OUTCOME_COL].values - - # 2. Generate predictions using an ASSUMED model that only considers - # main effects. This is where the misspecification is introduced. - - # --- Propensity Score Model --- - # The assumed model for PS uses only the first 3 coefficients (intercept, X1, X2). - # If `cls.alpha` has a non-zero 4th element, this model is MISSPECIFIED. - ps_model_coeffs = np.array(cls.alpha[:3]) - X_ps_design = np.column_stack( - [np.ones(cls.n), X_raw] - ) # Design matrix for main effects - ps = expit(X_ps_design @ ps_model_coeffs) + cls.noise_level * rng.normal( - size=cls.n + # Unpack the dictionary into class attributes for the tests to use + cls.A = sim_data["A"] + cls.Y = sim_data["Y"] + cls.ps = sim_data["ps"] + cls.Y1_hat = sim_data["Y1_hat"] + cls.Y0_hat = sim_data["Y0_hat"] + cls.Yhat = sim_data["Yhat"] + cls.true_ate = sim_data["true_ate"] + cls.true_att = sim_data["true_att"] + cls.true_rr = sim_data["true_rr"] + + cls.data = pd.DataFrame( + { + TREATMENT_COL: cls.A, + OUTCOME_COL: cls.Y, + PS_COL: cls.ps, + PROBAS_T1_COL: cls.Y1_hat, + PROBAS_T0_COL: cls.Y0_hat, + PROBAS_COL: cls.Yhat, + PID_COL: np.arange(len(cls.A)), + } ) - - outcome_model_coeffs = np.array(cls.beta[:4]) - - # Design matrices for main effects outcome model - X_y1_design = np.column_stack([np.ones(cls.n), np.ones(cls.n), X_raw]) # A=1 - X_y0_design = np.column_stack([np.ones(cls.n), np.zeros(cls.n), X_raw]) # A=0 - X_y_obs_design = np.column_stack([np.ones(cls.n), A, X_raw]) # A=observed - - # Generate predictions for Y1_hat, Y0_hat, and Yhat - Y1_hat = expit( - X_y1_design @ outcome_model_coeffs - ) + cls.noise_level * rng.normal(size=cls.n) - - Y0_hat = expit( - X_y0_design @ outcome_model_coeffs - ) + cls.noise_level * rng.normal(size=cls.n) - Yhat = expit( - X_y_obs_design @ outcome_model_coeffs - ) + cls.noise_level * rng.normal(size=cls.n) - - # 3. Finalize data preparation - eps = cls.cutoff_epsilon - cls.A, cls.Y = A, Y - cls.ps = np.clip(ps, eps, 1 - eps) - cls.Y1_hat = np.clip(Y1_hat, eps, 1 - eps) - cls.Y0_hat = np.clip(Y0_hat, eps, 1 - eps) - cls.Yhat = np.clip(Yhat, eps, 1 - eps) - - cls.true_ate = compute_ATE_theoretical_from_data(data, beta=cls.beta) - cls.true_att = compute_ATT_theoretical_from_data(data, beta=cls.beta) - cls.true_rr = compute_RR_theoretical_from_data(data, beta=cls.beta) - - # for classes that take dataframe as input - cls.data = data - cls.data[PID_COL] = np.arange(len(data)) - cls.data[TREATMENT_COL] = A - cls.data[OUTCOME_COL] = Y - cls.data[PS_COL] = ps - cls.data[PROBAS_T1_COL] = Y1_hat - cls.data[PROBAS_T0_COL] = Y0_hat - cls.data[PROBAS_COL] = Yhat diff --git a/tests/test_functional/test_tmle/test_coverage.py b/tests/test_functional/test_tmle/test_coverage.py new file mode 100644 index 0000000..8150845 --- /dev/null +++ b/tests/test_functional/test_tmle/test_coverage.py @@ -0,0 +1,223 @@ +import unittest + +import numpy as np + +from CausalEstimate.estimators.functional.tmle import ( + compute_tmle_ate, + compute_tmle_rr, +) +from CausalEstimate.estimators.functional.tmle_att import ( + compute_tmle_att, +) +from CausalEstimate.utils.constants import CI95_LOWER, CI95_UPPER +from tests.helpers.setup import generate_simulation_data + + +class TestTMLECoverage(unittest.TestCase): + """ + Tests CI coverage when models are misspecified. + This validates double robustness. + """ + + n_simulations = 300 + n_samples = 2000 + alpha = [-0.2, 0.5, -0.5] + beta = [0.1, 0.4, 0.6, -2] + noise_level = 0.0 + + def test_ate_coverage(self): + """Test if the ATE 95% CI covers the true value ~95% of the time.""" + coverage_count = 0 + + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, # Use loop index as the seed for reproducibility + ) + + result = compute_tmle_ate( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + + true_ate = sim_data["true_ate"] + if result[CI95_LOWER] is not None and np.isfinite(result[CI95_LOWER]): + if result[CI95_LOWER] <= true_ate <= result[CI95_UPPER]: + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"\nATE Coverage: {coverage_probability:.3f}") + + self.assertGreaterEqual(coverage_probability, 0.94) + + def test_att_coverage(self): + """Test ATT coverage with misspecified nuisance models.""" + coverage_count = 0 + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, + ) + result = compute_tmle_att( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + true_att = sim_data["true_att"] + if np.isfinite(result[CI95_LOWER]) and ( + result[CI95_LOWER] <= true_att <= result[CI95_UPPER] + ): + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"ATT Coverage: {coverage_probability:.3f}") + self.assertGreaterEqual(coverage_probability, 0.94) + + def test_rr_coverage(self): + """Test RR coverage with misspecified nuisance models.""" + coverage_count = 0 + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, + ) + result = compute_tmle_rr( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + true_rr = sim_data["true_rr"] + if np.isfinite(result[CI95_LOWER]) and ( + result[CI95_LOWER] <= true_rr <= result[CI95_UPPER] + ): + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"RR Coverage: {coverage_probability:.3f}") + self.assertGreaterEqual(coverage_probability, 0.94) + + +class TestTMLECoverageMisspecified(unittest.TestCase): + """ + Tests CI coverage when models are misspecified. + This validates double robustness. + """ + + n_simulations = 200 + n_samples = 2000 + alpha = [-0.2, 0.5, -0.5, 2] + beta = [0.1, 0.4, 0.6, -2, 0.1] + noise_level = 0.01 + + def test_ate_coverage(self): + """Test if the ATE 95% CI covers the true value ~95% of the time.""" + coverage_count = 0 + + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, # Use loop index as the seed for reproducibility + ) + + result = compute_tmle_ate( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + + true_ate = sim_data["true_ate"] + if result[CI95_LOWER] is not None and np.isfinite(result[CI95_LOWER]): + if result[CI95_LOWER] <= true_ate <= result[CI95_UPPER]: + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"\nATE Coverage (Misspecified): {coverage_probability:.3f}") + + self.assertGreaterEqual(coverage_probability, 0.92) + + def test_att_coverage(self): + """Test ATT coverage with misspecified nuisance models.""" + coverage_count = 0 + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, + ) + result = compute_tmle_att( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + true_att = sim_data["true_att"] + if np.isfinite(result[CI95_LOWER]) and ( + result[CI95_LOWER] <= true_att <= result[CI95_UPPER] + ): + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"ATT Coverage (Misspecified): {coverage_probability:.3f}") + self.assertGreaterEqual(coverage_probability, 0.90) + + def test_rr_coverage(self): + """Test RR coverage with misspecified nuisance models.""" + coverage_count = 0 + for i in range(self.n_simulations): + sim_data = generate_simulation_data( + n=self.n_samples, + alpha=self.alpha, + beta=self.beta, + noise_level=self.noise_level, + seed=i, + ) + result = compute_tmle_rr( + sim_data["A"], + sim_data["Y"], + sim_data["ps"], + sim_data["Y0_hat"], + sim_data["Y1_hat"], + sim_data["Yhat"], + ) + true_rr = sim_data["true_rr"] + if np.isfinite(result[CI95_LOWER]) and ( + result[CI95_LOWER] <= true_rr <= result[CI95_UPPER] + ): + coverage_count += 1 + + coverage_probability = coverage_count / self.n_simulations + print(f"RR Coverage (Misspecified): {coverage_probability:.3f}") + self.assertGreaterEqual(coverage_probability, 0.90) + + +if __name__ == "__main__": + unittest.main() From b4971262709dd7c4e2ff9a8c5ae4273dbd0278c7 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Sat, 11 Oct 2025 15:07:33 +0200 Subject: [PATCH 2/3] Refine influence curve calculations in variance.py for numerical stability This commit enhances the handling of edge cases in the influence curve calculations for ATT and RR estimators. Key changes include: - Updated the condition for returning NaN values when `p_treated` and the means of control and treatment groups are close to zero, using `np.isclose` with a specified tolerance. - Ensured that NaN arrays are created with the correct shape and data type for consistency. These modifications improve the robustness of the estimators against numerical instability and edge cases. --- CausalEstimate/estimators/functional/variance.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/CausalEstimate/estimators/functional/variance.py b/CausalEstimate/estimators/functional/variance.py index 273dc89..cd77a00 100644 --- a/CausalEstimate/estimators/functional/variance.py +++ b/CausalEstimate/estimators/functional/variance.py @@ -78,8 +78,8 @@ def _compute_ic_att( p_treated: float, ) -> np.ndarray: """Influence curve for ATT.""" - if p_treated == 0: - return np.full_like(Y, np.nan) + if np.isclose(p_treated, 0.0, atol=1e-12): + return np.full(Y.shape, np.nan, dtype=float) ic = H * (Y - Yhat_star) + (A / p_treated) * (Q_star_1 - Q_star_0 - psi) return ic @@ -96,8 +96,8 @@ def _compute_ic_rr( mu1_star = np.mean(Q_star_1) mu0_star = np.mean(Q_star_0) - if np.isclose(mu0_star, 0) or np.isclose(mu1_star, 0): - return np.full_like(Y, np.nan) + if np.isclose(mu0_star, 0.0, atol=eps) or np.isclose(mu1_star, 0.0, atol=eps): + return np.full(Y.shape, np.nan, dtype=float) # IC for mu1 ic_mu1 = (A / (ps + eps)) * (Y - Q_star_1) + Q_star_1 - mu1_star From 24e40287ed52e924f94e47f339be846eb6faf55b Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Sat, 11 Oct 2025 15:07:42 +0200 Subject: [PATCH 3/3] Update bootstrap sample tests to verify original dataset values This commit modifies the test for generating bootstrap samples to ensure that all values in the samples originate from the original dataset. The previous test for sample uniqueness has been replaced with assertions that confirm each sample contains only values present in the original DataFrame. This change enhances the reliability of the bootstrap sampling tests. --- tests/test_core/test_bootstrap.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_core/test_bootstrap.py b/tests/test_core/test_bootstrap.py index f60497b..f2063e8 100644 --- a/tests/test_core/test_bootstrap.py +++ b/tests/test_core/test_bootstrap.py @@ -37,13 +37,16 @@ def test_generate_bootstrap_samples_with_seed(self): # Check if the samples are identical when using the same seed pd.testing.assert_frame_equal(samples1, samples2) - def test_generate_bootstrap_samples_different_samples(self): - # Test that bootstrap samples are actually different + def test_generate_bootstrap_samples_values_from_original(self): + # Test that bootstrap samples contain values from the original dataset samples = generate_bootstrap_samples(self.test_df, 2) - # Check that the two samples are different (they should be, with very high probability) - with self.assertRaises(AssertionError): - pd.testing.assert_frame_equal(samples[0], samples[1]) + # Each bootstrap sample should only contain values from the original + for sample in samples: + # All values in column A should be from the original [1,2,3,4,5] + self.assertTrue(sample["A"].isin(self.test_df["A"]).all()) + # All values in column B should be from the original [10,20,30,40,50] + self.assertTrue(sample["B"].isin(self.test_df["B"]).all()) def test_generate_bootstrap_samples_with_empty_df(self): # Test behavior with empty DataFrame