Feat/update interface

CausalEstimate/estimators/functional/ipw.py

@@ -173,37 +173,70 @@ def compute_ipw_weights(
     ps: np.ndarray,
     weight_type: Literal["ATE", "ATT"] = "ATE",
     clip_percentile: float = 1,
+    eps: float = 1e-9,
 ) -> np.ndarray:
     """
-    Compute IPW weights for ATE or ATT with optional stabilization for ATE.
+    Computes Inverse Propensity Score (IPW) weights with optional clipping.
+
+    This function calculates weights for estimating the Average Treatment Effect (ATE)
+    or the Average Treatment Effect on the Treated (ATT).
+
+    Formulas:
+    - ATE: w = A/ps + (1-A)/(1-ps)
+    - ATT: w = A + (1-A) * ps/(1-ps)
+
+        Args:
+        A: Binary treatment assignment vector (1 for treated, 0 for control).
+        ps: Propensity score vector (estimated probability of treatment).
+        weight_type: The type of estimand, either "ATE" or "ATT".
+        clip_percentile: The upper percentile at which to clip weights to prevent
+                         extreme values. A value of 1.0 (default) applies no
+                         clipping. For example, 0.99 clips the top 1%.
+        eps: A small constant to add to denominators for numerical stability.
+
+    Returns:
+        An array of computed IPW weights.
+
+    Raises:
+        ValueError: If `weight_type` is invalid, shapes mismatch, or
+                    `clip_percentile` is out of bounds.
     """
 
-    if weight_type == "ATE":
-        weight_treated = 1 / ps
-        weight_control = 1 / (1 - ps)
-        if clip_percentile < 1:
-            treated_mask = A == 1
-            control_mask = A == 0
-            threshold = np.percentile(
-                weight_treated[treated_mask], clip_percentile * 100
-            )  # only compute threshold for treated group
-            weight_treated = np.clip(weight_treated, a_min=None, a_max=threshold)
-            threshold = np.percentile(
-                weight_control[control_mask], clip_percentile * 100
-            )  # only compute threshold for control group
-            weight_control = np.clip(weight_control, a_min=None, a_max=threshold)
-        weights = np.where(A == 1, weight_treated, weight_control)
-
-    elif weight_type == "ATT":
-        weight_treated = np.ones_like(A, dtype=float)
-        weight_control = ps / (1 - ps)
-        weights = np.where(A == 1, weight_treated, weight_control)
-        if clip_percentile < 1:
-            # Only clip the weights for the control group
-            control_weights = weights[A == 0]
-            if control_weights.size > 0:
-                threshold = np.percentile(control_weights, clip_percentile * 100)
-                weights[A == 0] = np.clip(control_weights, a_min=None, a_max=threshold)
-    else:
+    # --- 1. Input Validation ---
+    if weight_type not in ["ATE", "ATT"]:
         raise ValueError("weight_type must be 'ATE' or 'ATT'")
+    if not (0 < clip_percentile <= 1.0):
+        raise ValueError("clip_percentile must be in the interval (0, 1.0].")
+    if A.shape != ps.shape:
+        raise ValueError("A and ps must have the same shape.")
+
+    # --- 2. Core Weight Calculation ---
+    if weight_type == "ATE":
+        # Vectorized formula for ATE weights
+        weights = A / (ps + eps) + (1 - A) / (1 - ps + eps)
+    else:  # weight_type == "ATT"
+        # For ATT, treated units have a weight of 1.
+        # Vectorized formula for ATT weights.
+        weights = A + (1 - A) * ps / (1 - ps + eps)
+
+    # --- 3. Unified Weight Clipping ---
+    if clip_percentile < 1.0:
+        q = clip_percentile * 100
+
+        # This logic is now applied to both ATE and ATT.
+        # For ATT, the 'treated_mask' section is a no-op but is still executed.
+        treated_mask = A == 1
+        if np.any(treated_mask):
+            threshold_t = np.percentile(weights[treated_mask], q)
+            weights[treated_mask] = np.clip(
+                weights[treated_mask], a_min=None, a_max=threshold_t
+            )
+
+        control_mask = ~treated_mask
+        if np.any(control_mask):
+            threshold_c = np.percentile(weights[control_mask], q)
+            weights[control_mask] = np.clip(
+                weights[control_mask], a_min=None, a_max=threshold_c
+            )
+
     return weights

CausalEstimate/estimators/functional/tmle.py

@@ -1,13 +1,12 @@
 import warnings
-from typing import Tuple, Optional
+from typing import Tuple
 
 import numpy as np
 from scipy.special import expit, logit
 
-
 from CausalEstimate.estimators.functional.utils import (
-    compute_initial_effect,
     compute_clever_covariate_ate,
+    compute_initial_effect,
     estimate_fluctuation_parameter,
 )
 from CausalEstimate.utils.constants import (
@@ -24,20 +23,20 @@ def compute_tmle_ate(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     Yhat: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> dict:
     """
-    Estimate the ATE using TMLE, with optional weight stabilization.
+    Estimate the ATE using TMLE, with optional weight clipping.
     """
     Q_star_1, Q_star_0 = compute_estimates(
-        A, Y, ps, Y0_hat, Y1_hat, Yhat, stabilized=stabilized
+        A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile
     )
     ate = (Q_star_1 - Q_star_0).mean()
 
     return {
         EFFECT: ate,
-        EFFECT_treated: Q_star_1.mean(),  # Return mean of predictions
-        EFFECT_untreated: Q_star_0.mean(),  # Return mean of predictions
+        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),
     }
 
@@ -49,13 +48,13 @@ def compute_tmle_rr(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     Yhat: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> dict:
     """
-    Estimate the Risk Ratio using TMLE, with optional weight stabilization.
+    Estimate the Risk Ratio using TMLE, with optional weight clipping.
     """
     Q_star_1, Q_star_0 = compute_estimates(
-        A, Y, ps, Y0_hat, Y1_hat, Yhat, stabilized=stabilized
+        A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile
     )
     Q_star_1_m = Q_star_1.mean()
     Q_star_0_m = Q_star_0.mean()
@@ -89,16 +88,14 @@ def compute_estimates(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     Yhat: np.ndarray,
-    stabilized: bool = False,  # New parameter
+    clip_percentile: float = 1,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Compute updated outcome estimates using TMLE targeting step.
     """
-    H = compute_clever_covariate_ate(A, ps, stabilized=stabilized)
+    H = compute_clever_covariate_ate(A, ps, clip_percentile=clip_percentile)
     epsilon = estimate_fluctuation_parameter(H, Y, Yhat)
-
-    pi = A.mean() if stabilized else None
-    Q_star_1, Q_star_0 = update_estimates(ps, Y0_hat, Y1_hat, epsilon, pi=pi)
+    Q_star_1, Q_star_0 = update_estimates(ps, Y0_hat, Y1_hat, epsilon)
 
     return Q_star_1, Q_star_0
 
@@ -108,22 +105,13 @@ def update_estimates(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     epsilon: float,
-    pi: Optional[float] = None,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Update the initial outcome estimates using the fluctuation parameter.
-    If pi is provided, uses stabilized clever covariates.
     """
-    if pi is not None:
-        # Stabilized clever covariates
-        H1 = pi / ps
-        H0 = -(1.0 - pi) / (1.0 - ps)
-    else:
-        # Unstabilized clever covariates
-        H1 = 1.0 / ps
-        H0 = -1.0 / (1.0 - ps)
+    H1 = 1.0 / ps
+    H0 = -1.0 / (1.0 - ps)
 
-    # Update initial estimates with targeting step
     Q_star_1 = expit(logit(Y1_hat) + epsilon * H1)
     Q_star_0 = expit(logit(Y0_hat) + epsilon * H0)
 

CausalEstimate/estimators/functional/tmle_att.py

@@ -24,28 +24,46 @@ def compute_estimates_att(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     Yhat: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> Tuple[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, stabilized=stabilized)
+    H = compute_clever_covariate_att(A, ps, clip_percentile=clip_percentile)
     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
 
-    # The update term for the treated group is always the same
+    # The update term for the potential outcome under treatment, Q(1,W).
+    # This is a scalar value applied to everyone's Y1_hat.
     update_term_1 = epsilon * (1.0 / p_treated)
 
-    # The update term for the control group depends on stabilization
-    if stabilized:
-        # Stabilized update term for controls
-        update_term_0 = -epsilon * (ps * (1 - p_treated) / (p_treated * (1 - ps)))
-    else:
-        # Unstabilized update term for controls
-        update_term_0 = -epsilon * (ps / (p_treated * (1 - ps)))
+    # 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))
+
+    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)
 
@@ -59,14 +77,14 @@ def compute_tmle_att(
     Y0_hat: np.ndarray,
     Y1_hat: np.ndarray,
     Yhat: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> dict:
     """
     Estimate the Average Treatment Effect on the Treated (ATT) using TMLE,
-    with optional weight stabilization for the control group.
+    with optional clipping for the control group.
     """
     Q_star_1, Q_star_0 = compute_estimates_att(
-        A, Y, ps, Y0_hat, Y1_hat, Yhat, stabilized=stabilized
+        A, Y, ps, Y0_hat, Y1_hat, Yhat, clip_percentile=clip_percentile
     )
 
     # The final ATT parameter is the mean difference within the treated population

CausalEstimate/estimators/functional/utils.py

@@ -17,31 +17,49 @@
 def compute_clever_covariate_ate(
     A: np.ndarray,
     ps: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> np.ndarray:
     """
-    Compute the clever covariate H for ATE TMLE.
+    Compute the clever covariate H for ATE TMLE with optional clipping.
 
     Parameters:
     -----------
     A: np.ndarray
         Treatment assignment (0 or 1)
     ps: np.ndarray
         Propensity scores
-    stabilized: bool, optional
-        Whether to use stabilized weights. Default is False.
+    clip_percentile: float, optional
+        Percentile to clip the weights at. Default is 1 (no clipping).
 
     Returns:
     --------
     np.ndarray: The clever covariate H
     """
-    if stabilized:
-        pi = A.mean()
-        # Stabilized clever covariate
-        H = A * pi / ps - (1 - A) * (1 - pi) / (1 - ps)
-    else:
-        # Unstabilized clever covariate
-        H = A / ps - (1 - A) / (1 - ps)
+    # Unstabilized clever covariate components
+    H1_component = 1.0 / ps
+    H0_component = 1.0 / (1 - ps)
+
+    # Apply clipping if requested (similar to IPW clipping logic)
+    if clip_percentile < 1:
+        treated_mask: np.ndarray = A == 1
+        control_mask: np.ndarray = A == 0
+
+        # Clip treated component
+        if treated_mask.sum() > 0:
+            threshold_treated = np.percentile(
+                H1_component[treated_mask], clip_percentile * 100
+            )
+            H1_component = np.clip(H1_component, a_min=None, a_max=threshold_treated)
+
+        # Clip control component
+        if control_mask.sum() > 0:
+            threshold_control = np.percentile(
+                H0_component[control_mask], clip_percentile * 100
+            )
+            H0_component = np.clip(H0_component, a_min=None, a_max=threshold_control)
+
+    # Combine components into clever covariate
+    H = A * H1_component - (1 - A) * H0_component
 
     _validate_clever_covariate(H, "ATE")
     return H
@@ -50,19 +68,19 @@ def compute_clever_covariate_ate(
 def compute_clever_covariate_att(
     A: np.ndarray,
     ps: np.ndarray,
-    stabilized: bool = False,
+    clip_percentile: float = 1,
 ) -> np.ndarray:
     """
-    Compute the clever covariate H for ATT TMLE.
+    Compute the clever covariate H for ATT TMLE with optional clipping.
 
     Parameters:
     -----------
     A: np.ndarray
         Treatment assignment (0 or 1)
     ps: np.ndarray
         Propensity scores
-    stabilized: bool, optional
-        Whether to use stabilized weights. Default is False.
+    clip_percentile: float, optional
+        Percentile to clip the weights at. Default is 1 (no clipping).
 
     Returns:
     --------
@@ -75,18 +93,28 @@ def compute_clever_covariate_att(
         )
         return np.zeros_like(A, dtype=float)
 
-    # Component for treated individuals
+    # Component for treated individuals (no clipping needed, always 1/p_treated)
     H_treated = A / p_treated
 
-    # Component for control individuals
-    if stabilized:
-        # Stabilized clever covariate for controls
-        H_control = (1 - A) * ps * (1 - p_treated) / (p_treated * (1 - ps))
-    else:
-        # Unstabilized clever covariate for controls
-        H_control = (1 - A) * ps / (p_treated * (1 - ps))
+    # Component for control individuals - this contains the weights that need clipping
+    # Using unstabilized weights: ps / (p_treated * (1 - ps))
+    weight_component = ps / (p_treated * (1 - ps))
+
+    # Apply clipping to control weights if requested (similar to ATT IPW clipping)
+    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)
+            weight_component = np.where(
+                control_mask,
+                np.clip(weight_component, a_min=None, a_max=threshold),
+                weight_component,
+            )
 
+    H_control = (1 - A) * weight_component
     H = H_treated - H_control
+
     _validate_clever_covariate(H, "ATT")
     return H