Enhance TMLE estimators with confidence interval calculations

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]

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

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,49 +27,39 @@ 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:
         control_mask: np.ndarray = A == 0
         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,
     }

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 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
+
+
+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.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
+    # 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

CausalEstimate/estimators/tmle.py

@@ -1,4 +1,3 @@
-# CausalEstimate/estimators/tmle.py
 import pandas as pd
 
 from CausalEstimate.estimators.base import BaseEstimator

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"