Enhance TMLE estimators with confidence interval calculations - #84
Conversation
…iance 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.
WalkthroughAdds standardized CI/SE constants, a new variance/influence-curve module for CI/std err, integrates CI computation into TMLE ATE/ATT/RR flows (propagating Yhat_star and H), updates TMLE ATT edge cases, centralizes test data generation, and adds simulation-based CI coverage tests. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Caller
participant TMLE as functional.tmle
participant Var as functional.variance
Caller->>TMLE: compute_tmle_ate(A,Y,ps,Y0_hat,Y1_hat,Yhat,...)
TMLE->>TMLE: compute_estimates(...) -> (Q1*, Q0*, Yhat*, H)
TMLE->>Var: compute_ci("ATE", psi, Q1*, Q0*, Y, A, ps, Yhat*, H)
Var-->>TMLE: {STD_ERR, CI95_LOWER, CI95_UPPER}
TMLE-->>Caller: {psi, components, STD_ERR, CI95_LOWER, CI95_UPPER}
note right of Var: RR computed on log-scale then exponentiated
sequenceDiagram
autonumber
participant Caller
participant ATT as functional.tmle_att
participant Var as functional.variance
Caller->>ATT: compute_tmle_att(A,Y,ps,Y0_hat,Y1_hat,Yhat,...)
alt p_treated == 0
ATT-->>Caller: {psi: NaN, STD_ERR: NaN, CI95_LOWER: NaN, CI95_UPPER: NaN}
else
ATT->>ATT: compute_estimates_att(...) -> (Q1*, Q0*, Yhat*, H)
ATT->>Var: compute_ci("ATT", psi, Q1*, Q0*, Y, A, ps, Yhat*, H)
Var-->>ATT: {STD_ERR, CI95_LOWER, CI95_UPPER}
ATT-->>Caller: {psi, EFFECT_treated, EFFECT_untreated, STD_ERR, CI95_LOWER, CI95_UPPER}
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (15)
CausalEstimate/utils/constants.py (1)
1-12: Duplicate OUTCOME_COL definition in fileOUTCOME_COL is defined twice ("outcome" and "Y"). Consider deduplicating to avoid confusion and accidental misuse.
CausalEstimate/core/multi_estimator.py (1)
159-161: Use sample std for bootstrap SE (ddof=1)Bootstrap SE is typically the sample std of replicates. Use ddof=1.
- std_err = float(np.std(effects)) + std_err = float(np.std(effects, ddof=1))tests/helpers/setup.py (1)
2-2: Annotate mutable class attributes with ClassVar (RUF012)alpha and beta are mutable class attributes; annotate with ClassVar. Also import ClassVar.
-from typing import Any, Dict, List +from typing import Any, Dict, List, ClassVar- alpha: list = [0.1, 0.2, -0.3] - beta: list = [0.5, 0.8, -0.6, 0.3] + alpha: ClassVar[List[float]] = [0.1, 0.2, -0.3] + beta: ClassVar[List[float]] = [0.5, 0.8, -0.6, 0.3]Also applies to: 84-85
CausalEstimate/estimators/functional/tmle_att.py (3)
37-41: Avoid estimating epsilon when no treated subjectsIf p_treated == 0, you can skip estimating epsilon to save work.
- H = compute_clever_covariate_att(A, ps, clip_percentile=clip_percentile, eps=eps) - epsilon = estimate_fluctuation_parameter(H, Y, Yhat) - - p_treated = np.mean(A == 1) - if p_treated == 0: - Yhat_star = Yhat.copy() # No update if no treated - return Y1_hat, Y0_hat, Yhat_star, H + H = compute_clever_covariate_att(A, ps, clip_percentile=clip_percentile, eps=eps) + p_treated = np.mean(A == 1) + if p_treated == 0: + Yhat_star = Yhat.copy() # No update if no treated + return Y1_hat, Y0_hat, Yhat_star, H + epsilon = estimate_fluctuation_parameter(H, Y, Yhat)
82-86: Return CI fields when no treated (schema stability)Early return lacks STD_ERR/CI keys; include NaNs to keep a stable interface. Import constants.
-from CausalEstimate.estimators.functional.variance import compute_ci -from CausalEstimate.utils.constants import EFFECT, EFFECT_treated, EFFECT_untreated +from CausalEstimate.estimators.functional.variance import compute_ci +from CausalEstimate.utils.constants import ( + EFFECT, + EFFECT_treated, + EFFECT_untreated, + STD_ERR, + CI95_LOWER, + CI95_UPPER, +)- if not np.any(treated_mask): - # Handle case with no treated subjects - return {EFFECT: np.nan, EFFECT_treated: np.nan, EFFECT_untreated: np.nan} + if not np.any(treated_mask): + # Handle case with no treated subjects + return { + EFFECT: np.nan, + EFFECT_treated: np.nan, + EFFECT_untreated: np.nan, + STD_ERR: np.nan, + CI95_LOWER: np.nan, + CI95_UPPER: np.nan, + }Also applies to: 17-19
46-53: Optional: clip control weights only (align with H implementation)Currently clipping is applied to the full vector. To mirror compute_clever_covariate_att, restrict clipping to controls.
- if control_mask.sum() > 0: - control_weights = weight_component[control_mask] - threshold = np.percentile(control_weights, clip_percentile * 100) - weight_component = np.clip(weight_component, a_min=None, a_max=threshold) + 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, + )tests/test_functional/test_tmle/test_coverage.py (4)
24-26: Annotate mutable class attributes with ClassVar (RUF012)alpha and beta are mutable class attributes; annotate with ClassVar. Add import.
-import numpy as np +import numpy as np +from typing import ClassVar- alpha = [-0.2, 0.5, -0.5] - beta = [0.1, 0.4, 0.6, -2] + alpha: ClassVar[list[float]] = [-0.2, 0.5, -0.5] + beta: ClassVar[list[float]] = [0.1, 0.4, 0.6, -2]- alpha = [-0.2, 0.5, -0.5, 2] - beta = [0.1, 0.4, 0.6, -2, 0.1] + alpha: ClassVar[list[float]] = [-0.2, 0.5, -0.5, 2] + beta: ClassVar[list[float]] = [0.1, 0.4, 0.6, -2, 0.1]Also applies to: 127-129, 3-3
18-20: Docstring mismatchThis class tests correctly specified nuisance models (noise_level=0, matching model form). Update docstring to avoid confusion.
- Tests CI coverage when models are misspecified. - This validates double robustness. + Tests CI coverage with correctly specified nuisance models.
51-53: Optional: check both CI bounds are finite before comparisonMake the finite check symmetric to avoid edge cases.
- if result[CI95_LOWER] is not None and np.isfinite(result[CI95_LOWER]): - if result[CI95_LOWER] <= true_ate <= result[CI95_UPPER]: + if ( + result[CI95_LOWER] is not None + and np.isfinite(result[CI95_LOWER]) + and np.isfinite(result[CI95_UPPER]) + ): + if result[CI95_LOWER] <= true_ate <= result[CI95_UPPER]: coverage_count += 1Also applies to: 154-156
22-28: Note: test runtime may be high300×2000 and 200×2000 runs are heavy. If CI time becomes an issue, consider gating by env var, reducing counts, or marking slow tests.
If needed, I can draft a pattern to honor an env var (e.g., FAST_TESTS) to cut n_simulations during CI.
Also applies to: 125-131
CausalEstimate/estimators/functional/tmle.py (1)
86-91: Avoid infinite-psi CIs for RRWhen rr is set to inf (threshold or zero denominator), compute_ci will produce inf CIs via log(psi). Prefer returning NaN CIs for non-finite psi to avoid misleading intervals. Easiest: guard in compute_ci to early-return NaNs if psi is not finite (applies universally and keeps tmle.py simple).
Apply in variance.py:
@@ - n = len(Y) - if n == 0: + n = len(Y) + if n == 0 or not np.isfinite(psi): return {STD_ERR: np.nan, CI95_LOWER: np.nan, CI95_UPPER: np.nan}Also applies to: 92-103
CausalEstimate/estimators/functional/variance.py (4)
20-22: Handle small samples to avoid ddof=1 warnings and NaNsWith n == 1, np.var(..., ddof=1) yields NaN and a warning. Return NaN CI cleanly for n < 2.
- n = len(Y) - if n == 0: + n = len(Y) + if n < 2: return {STD_ERR: np.nan, CI95_LOWER: np.nan, CI95_UPPER: np.nan} @@ - var_ic = np.var(ic, ddof=1) # Use ddof=1 for sample variance + var_ic = np.var(ic, ddof=1) # sample variance; safe since n >= 2 nowAlso applies to: 40-43
25-26: Remove unused parameter A from ATE IC to silence lint and reduce API surfaceA is not used in _compute_ic_ate; drop it and update the call.
@@ - if effect_type in ["ATE", "ARR"]: - ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, A, Yhat_star, H) + if effect_type in ["ATE", "ARR"]: + ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, Yhat_star, H) @@ -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: +def _compute_ic_ate( + psi: float, + Q_star_1: np.ndarray, + Q_star_0: np.ndarray, + Y: np.ndarray, + Yhat_star: np.ndarray, + H: np.ndarray, +) -> np.ndarray:Based on static analysis hints.
Also applies to: 57-65
45-52: Guard RR CI when psi ≤ 0 or not finiteFor RR, log(psi) requires psi > 0 and finite. You already handle NaN ICs; add a quick guard after computing std_err to return NaN CIs if psi is non‑finite or ≤ 0, preventing inf or -inf bounds.
- if effect_type == "RR": + if effect_type == "RR": + if not np.isfinite(psi) or psi <= 0: + return {STD_ERR: std_err, CI95_LOWER: np.nan, CI95_UPPER: np.nan} # For RR, CIs are calculated on the log scale and then exponentiated log_psi = np.log(psi)
33-35: Optional: shorten error message per TRY003The ValueError message is long. Consider a concise message or a dedicated exception type.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
CausalEstimate/core/multi_estimator.py(2 hunks)CausalEstimate/estimators/functional/tmle.py(6 hunks)CausalEstimate/estimators/functional/tmle_att.py(3 hunks)CausalEstimate/estimators/functional/variance.py(1 hunks)CausalEstimate/estimators/tmle.py(0 hunks)CausalEstimate/utils/constants.py(1 hunks)tests/helpers/setup.py(2 hunks)tests/test_functional/test_tmle/test_coverage.py(1 hunks)
💤 Files with no reviewable changes (1)
- CausalEstimate/estimators/tmle.py
🧰 Additional context used
🧬 Code graph analysis (4)
CausalEstimate/estimators/functional/tmle_att.py (2)
CausalEstimate/estimators/functional/utils.py (3)
compute_initial_effect(153-215)estimate_fluctuation_parameter(218-294)compute_clever_covariate_att(71-125)CausalEstimate/estimators/functional/variance.py (1)
compute_ci(6-54)
tests/test_functional/test_tmle/test_coverage.py (3)
CausalEstimate/estimators/functional/tmle.py (2)
compute_tmle_ate(20-56)compute_tmle_rr(59-110)CausalEstimate/estimators/functional/tmle_att.py (1)
compute_tmle_att(65-107)tests/helpers/setup.py (1)
generate_simulation_data(25-74)
tests/helpers/setup.py (1)
CausalEstimate/simulation/binary_simulation.py (4)
simulate_binary_data(8-116)compute_ATE_theoretical_from_data(136-140)compute_ATT_theoretical_from_data(143-148)compute_RR_theoretical_from_data(151-155)
CausalEstimate/estimators/functional/tmle.py (2)
CausalEstimate/estimators/functional/variance.py (1)
compute_ci(6-54)CausalEstimate/estimators/functional/utils.py (3)
compute_initial_effect(153-215)compute_clever_covariate_ate(17-68)estimate_fluctuation_parameter(218-294)
🪛 Ruff (0.13.3)
tests/test_functional/test_tmle/test_coverage.py
24-24: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
25-25: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
127-127: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
128-128: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
CausalEstimate/estimators/functional/variance.py
33-35: Avoid specifying long messages outside the exception class
(TRY003)
62-62: Unused function argument: A
(ARG001)
tests/helpers/setup.py
84-84: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
85-85: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
🔇 Additional comments (3)
CausalEstimate/utils/constants.py (1)
27-29: Standardized CI/SE keys added — LGTMConsistent naming across codebase. No issues.
CausalEstimate/core/multi_estimator.py (1)
18-20: Good: switched bootstrap summary keys to constantsCentralizing keys via constants improves consistency.
CausalEstimate/estimators/functional/tmle.py (1)
33-36: CI integration and data flow LGTMPassing Yhat_star and H into compute_ci is correct; the unpacking of compute_estimates aligns with updated returns. The merge of ci_results into outputs is clean.
Also applies to: 38-48
…ility 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.
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.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
CausalEstimate/estimators/functional/variance.py (2)
6-54: Extract the 1.96 constant for 95% CI z-score.The z-score multiplier for 95% confidence intervals (1.96) is hardcoded in four places. Consider extracting it to a named constant at the module level for clarity and maintainability.
Add at the top of the file after imports:
Z_SCORE_95 = 1.96Then update the CI calculations:
# 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) + ci_lower = np.exp(log_psi - Z_SCORE_95 * std_err) + ci_upper = np.exp(log_psi + Z_SCORE_95 * std_err) else: # ATE, ATT, ARR - ci_lower = psi - 1.96 * std_err - ci_upper = psi + 1.96 * std_err + ci_lower = psi - Z_SCORE_95 * std_err + ci_upper = psi + Z_SCORE_95 * std_err
87-108: Clarify the dual role ofepsparameter.The
epsparameter (default 1e-9) serves two distinct purposes: numerical stability in divisions (lines 103, 105) and zero-tolerance checks (line 99). Using the same value for both can be confusing and may not be appropriate for all cases.Consider separating these concerns:
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, + eps: float = 1e-12, # for zero tolerance checks + stability_eps: float = 1e-9, # for numerical stability in divisions ) -> 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_mu1 = (A / (ps + stability_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_mu0 = ((1 - A) / (1 - ps + stability_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_rrAlternatively, if you prefer to keep a single parameter, document its dual purpose clearly in the docstring.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
CausalEstimate/estimators/functional/variance.py(1 hunks)tests/test_core/test_bootstrap.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/test_core/test_bootstrap.py (1)
CausalEstimate/core/bootstrap.py (1)
generate_bootstrap_samples(5-29)
🪛 Ruff (0.13.3)
CausalEstimate/estimators/functional/variance.py
33-35: Avoid specifying long messages outside the exception class
(TRY003)
62-62: Unused function argument: A
(ARG001)
🔇 Additional comments (1)
tests/test_core/test_bootstrap.py (1)
40-49: LGTM! More robust test logic.The refactored test now validates the fundamental bootstrap property—that resampled values originate from the original dataset—rather than relying on non-deterministic inequality checks. This is a more reliable and meaningful test.
| 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 |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Remove unused parameter A from function signature.
The parameter A is not used in the influence curve calculation for ATE. This is correctly flagged by static analysis.
Apply this diff:
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:Also update the call site at line 26:
- ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, A, Yhat_star, H)
+ ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, Yhat_star, H)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 | |
| # In CausalEstimate/estimators/functional/variance.py | |
| def _compute_ic_ate( | |
| psi: float, | |
| Q_star_1: np.ndarray, | |
| Q_star_0: np.ndarray, | |
| Y: 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_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 | |
| # At the call site (around line 26) | |
| ic = _compute_ic_ate(psi, Q_star_1, Q_star_0, Y, Yhat_star, H) |
🧰 Tools
🪛 Ruff (0.13.3)
62-62: Unused function argument: A
(ARG001)
1 participant
…iance handling
This commit introduces significant improvements to the TMLE estimators, focusing on the addition of confidence interval calculations and variance estimation.
Key changes include:
variance.pymodule to compute standard errors and 95% confidence intervals for ATE, ATT, and RR using influence curves.compute_tmle_ate,compute_tmle_att, andcompute_tmle_rrto return confidence interval results alongside existing estimates.constants.py.These changes enhance the statistical reliability of the TMLE estimators, providing users with more informative outputs regarding the uncertainty of their estimates.
Summary by CodeRabbit
New Features
Tests
Chores