Feat/update interface - #83
Conversation
…nverse-propensity weights. Updated functions to replace the stabilized flag with a clip_percentile parameter, enhancing weight calculation and improving handling of extreme values. Documentation updated to reflect these changes.
…otential outcomes. Reintroduced separate update terms for treatment and control groups, ensuring proper handling of weights and clipping based on the specified percentile. This change enhances the accuracy of the estimates while maintaining theoretical consistency.
…his cleanup eliminates redundant test cases that are no longer applicable following recent refactors and enhancements to the TMLE framework.
…ameter and replace it with clip_percentile Updated documentation to reflect these changes and ensure clarity in the usage of the new parameter.
Introduced a new test suite for the clip_percentile feature in TMLE estimators, validating the behavior of clever covariates under extreme propensity scores. The tests ensure that clipping effectively reduces extreme values and maintains the integrity of ATE and ATT estimates. Additionally, updated existing tests to incorporate clipping scenarios, enhancing coverage and robustness of the testing framework.
WalkthroughReplaces stabilization flags with percentile-based clipping across IPW and TMLE codepaths; introduces vectorized IPW weight formulas with an eps guard; unifies clever-covariate construction to clipping-based forms; updates estimator constructors and public function signatures; removes stabilized tests and adds comprehensive clipping tests. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant Estimator as TMLE / IPW Estimator
participant Functional as Functional API (tmle/ipw)
participant Utils as Utilities (covariates/weights)
User->>Estimator: compute(..., clip_percentile, eps)
Estimator->>Functional: compute_* (A, Y, ps, Yhat..., clip_percentile, eps)
Note right of Functional: stabilized branches removed — single clipping flow
Functional->>Utils: compute_clever_covariate_* (A, ps, clip_percentile, eps)
Utils-->>Functional: H (components optionally clipped per-group)
Functional->>Functional: vectorized weight/covariate computations (uses eps)
Functional-->>Estimator: updated estimates / Q*
Estimator-->>User: effect result
sequenceDiagram
rect rgba(240,248,255,0.45)
participant Data as (A, ps)
participant IPWFunc as functional/ipw.py
Data->>IPWFunc: inputs
IPWFunc->>IPWFunc: compute w:
IPWFunc->>IPWFunc: ATE: w = A/(ps+eps) + (1-A)/(1-ps+eps)
IPWFunc->>IPWFunc: ATT: w = A + (1-A)*ps/(1-ps+eps)
alt clip_percentile < 1
IPWFunc->>IPWFunc: compute per-group percentiles and clip treated/control weights
end
IPWFunc-->>Data: weights
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60–90 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
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 |
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
CausalEstimate/estimators/functional/utils.py (1)
69-117: ATT clever covariate: add eps, validate inputs, and clip-safe denominator.Same stability/validation concerns apply here; bound (1-ps) and document parameters.
-def compute_clever_covariate_att( - A: np.ndarray, - ps: np.ndarray, - clip_percentile: float = 1, -) -> np.ndarray: +def compute_clever_covariate_att( + A: np.ndarray, + ps: np.ndarray, + clip_percentile: float = 1.0, + eps: float = 1e-9, +) -> np.ndarray: @@ - clip_percentile: float, optional - Percentile to clip the weights at. Default is 1 (no clipping). + clip_percentile: float, optional + Upper percentile for clipping, in (0, 1]. Default 1 (no clipping). + eps: float, optional + Small constant for numerical stability in denominators. Default 1e-9. @@ - p_treated = np.mean(A == 1) + if A.shape != ps.shape: + raise ValueError("A and ps must have the same shape.") + if not (0 < clip_percentile <= 1.0): + raise ValueError("clip_percentile must be in (0, 1].") + p_treated = np.mean(A == 1) @@ - # Using unstabilized weights: ps / (p_treated * (1 - ps)) - weight_component = ps / (p_treated * (1 - ps)) + # Using unstabilized weights with numeric protection: ps / (p_treated * (1 - ps)) + denom0 = np.clip(1 - ps, eps, 1 - eps) + weight_component = ps / (p_treated * denom0)
🧹 Nitpick comments (14)
CausalEstimate/estimators/functional/utils.py (2)
133-144: Revise misleading “small H” warning.|H| ≪ 1 does not indicate ps near {0,1}; it risks false positives. Drop it or warn on ps-extremes upstream.
- if np.any(np.abs(H) < 1e-6): - warnings.warn( - f"Extremely small values < 1e-6 detected in clever covariate H for {estimand}. " - "This may indicate issues with propensity scores near 0 or 1.", - RuntimeWarning, - ) + # Consider logging ps-extremes during construction rather than inspecting |H| << 1 here.
178-186: Remove redundant local import; use module-level warnings.Minor cleanup.
- import warnings - warnings.warn( "Initial effect for untreated group is 0 or very close to 0, risk ratio undefined", RuntimeWarning, )CausalEstimate/estimators/functional/ipw.py (3)
20-24: Header comment references stabilized weights but code/API removed them.Update or drop this block to prevent confusion.
-We also provide an option to use stabilized weights as described in: -Miguel A Hernán 1, James M Robins -Estimating causal effects from epidemiological data -https://pubmed.ncbi.nlm.nih.gov/16790829/ +Note: This module uses percentile clipping; stabilized-weight options were removed.
96-101: Docstring: this is Hajek (ratio) estimator, not “simple Horvitz–Thompson”.Rename in docstring to match implementation (ratio of weighted sums).
- Computes E[Y(1)] and E[Y(0)] for the ATE using the simple Horvitz-Thompson estimator, + Computes E[Y(1)] and E[Y(0)] for the ATE using the Hajek (ratio) estimator,
176-241: Harden input errors and align docs with implementation; keep Ruff happy.
- Use shorter error messages (TRY003).
- Mention eps in formulas.
- Optional: warn when group denominators ~ 0 (poor overlap).
- Formulas: - - ATE: w = A/ps + (1-A)/(1-ps) - - ATT: w = A + (1-A) * ps/(1-ps) + Formulas (with numeric stabilizer eps): + - ATE: w = A/(ps+eps) + (1-A)/(1-ps+eps) + - ATT: w = A + (1-A) * ps/(1-ps+eps) @@ - if weight_type not in ["ATE", "ATT"]: - raise ValueError("weight_type must be 'ATE' or 'ATT'") + if weight_type not in ["ATE", "ATT"]: + raise ValueError("invalid weight_type") @@ - if not (0 < clip_percentile <= 1.0): - raise ValueError("clip_percentile must be in the interval (0, 1.0].") + if not (0 < clip_percentile <= 1.0): + raise ValueError("clip_percentile out of bounds") @@ - if A.shape != ps.shape: - raise ValueError("A and ps must have the same shape.") + if A.shape != ps.shape: + raise ValueError("shape mismatch: A vs ps")If you want, I can add overlap warnings where denominators approach 0 in the ratio estimators to improve debuggability. Shall I push a patch?
tests/test_functional/test_ipw_clipping.py (3)
44-59: Use tolerant equality for percentile-threshold max checks.Percentiles can vary by 1–2 ulp; prefer isclose to strict equality.
- # The new maximum value should be the threshold itself - self.assertAlmostEqual( - weights_clipped[treated_mask].max(), threshold_treated - ) + # The new maximum should equal the threshold within numerical tolerance + self.assertTrue(np.isclose(weights_clipped[treated_mask].max(), threshold_treated, rtol=1e-9, atol=1e-12)) @@ - self.assertAlmostEqual( - weights_clipped[control_mask].max(), threshold_control - ) + self.assertTrue(np.isclose(weights_clipped[control_mask].max(), threshold_control, rtol=1e-9, atol=1e-12))
183-192: Tighten “no clipping” equality check to match eps-stabilized implementation.You already use atol=1e-5. Consider adding a short comment about tiny eps in denominators to explain the tolerance.
249-252: Add negative tests for invalid inputs (bounds, shapes, weight_type).Covers new validations and prevents silent misuses.
+ def test_invalid_inputs_raise(self): + A = np.array([1, 0, 1]) + ps = np.array([0.8, 0.2]) # shape mismatch + with self.assertRaises(ValueError): + compute_ipw_weights(A, ps, weight_type="ATE", clip_percentile=1.0) + ps2 = np.array([0.8, 0.2, 0.5]) + with self.assertRaises(ValueError): + compute_ipw_weights(A, ps2, weight_type="BAD", clip_percentile=1.0) + with self.assertRaises(ValueError): + compute_ipw_weights(A, ps2, weight_type="ATE", clip_percentile=0.0) + with self.assertRaises(ValueError): + compute_ipw_weights(A, ps2, weight_type="ATE", clip_percentile=1.1)tests/test_functional/test_utils.py (2)
33-38: ATE clever covariate: good baseline check. Consider noting eps in comments.To preempt confusion once eps is added to utils, add a note that exact equality holds up to tiny eps.
119-130: Add invalid-parameter tests for clever covariates once validation is added.Guard against misuse (clip_percentile bounds, shape mismatch).
def test_clipping_percentile_bounds(self): """Test that clipping percentile parameter is properly bounded.""" # Test with valid clip_percentile values for clip_pct in [0.1, 0.5, 0.9, 1.0]: H = compute_clever_covariate_ate(self.A, self.ps, clip_percentile=clip_pct) self.assertTrue(np.all(np.isfinite(H))) # Test ATT as well for clip_pct in [0.1, 0.5, 0.9, 1.0]: H = compute_clever_covariate_att(self.A, self.ps, clip_percentile=clip_pct) self.assertTrue(np.all(np.isfinite(H))) + + def test_clever_covariate_invalid_params(self): + # clip_percentile out of bounds + with self.assertRaises(ValueError): + compute_clever_covariate_ate(self.A, self.ps, clip_percentile=0.0) + with self.assertRaises(ValueError): + compute_clever_covariate_att(self.A, self.ps, clip_percentile=1.1) + # shape mismatch + with self.assertRaises(ValueError): + compute_clever_covariate_ate(self.A, self.ps[:-1], clip_percentile=1.0)tests/test_functional/test_tmle_clipping.py (1)
108-171: Good end-to-end checks; consider adding an overlap-warning assertion later.Once overlap warnings are added in IPW/TMLE paths, assert they’re emitted for aggressive clipping/edge ps.
CausalEstimate/estimators/tmle.py (1)
11-21: Add deprecation shim for legacystabilizedkwarg in TMLE.initREADME references
stabilized=True(README.md:69) — accept legacystabilizedvia **kwargs and emit a DeprecationWarning to avoid breaking external users.@@ - def __init__( + def __init__( self, effect_type: str = "ATE", @@ - clip_percentile: float = 1, + clip_percentile: float = 1, + **kwargs, ):Then, right after super().init(...):
# Handle legacy keyword to avoid hard breakage if "stabilized" in kwargs: import warnings as _warnings _warnings.warn( "Parameter 'stabilized' has been removed. Use 'clip_percentile' instead.", DeprecationWarning, stacklevel=2, )Location: CausalEstimate/estimators/tmle.py (TMLE.init); README.md:69 mentions
stabilized=True.CausalEstimate/estimators/functional/tmle_att.py (1)
39-44: Guard against degenerate or near-degenerate treated shareDivision by p_treated and weight construction can become explosive when p_treated ≈ 0. Add a small-ε guard.
Apply:
- 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 + p_treated = np.mean(A == 1) + if p_treated <= 1e-8: + import warnings + warnings.warn("No (or virtually no) treated subjects; returning original predictions.", RuntimeWarning) + return Y1_hat, Y0_hatCausalEstimate/estimators/functional/tmle.py (1)
70-75: Nit: make RR “unrealistically large” threshold configurableHardcoding 1e5 may be surprising. Consider a module-level constant or a parameter with a default.
Apply:
- if rr > 1e5: + RR_MAX = 1e5 + if rr > RR_MAX: warnings.warn( - "Risk ratio is unrealistically large, returning inf.", RuntimeWarning + f"Risk ratio is unrealistically large (>{RR_MAX}), returning inf.", RuntimeWarning )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
CausalEstimate/estimators/functional/ipw.py(1 hunks)CausalEstimate/estimators/functional/tmle.py(5 hunks)CausalEstimate/estimators/functional/tmle_att.py(2 hunks)CausalEstimate/estimators/functional/utils.py(3 hunks)CausalEstimate/estimators/ipw.py(0 hunks)CausalEstimate/estimators/tmle.py(4 hunks)tests/test_functional/test_ipw_clipping.py(4 hunks)tests/test_functional/test_tmle/test_stabilization.py(0 hunks)tests/test_functional/test_tmle_clipping.py(1 hunks)tests/test_functional/test_utils.py(2 hunks)
💤 Files with no reviewable changes (2)
- CausalEstimate/estimators/ipw.py
- tests/test_functional/test_tmle/test_stabilization.py
🧰 Additional context used
🧬 Code graph analysis (6)
CausalEstimate/estimators/functional/tmle.py (1)
CausalEstimate/estimators/functional/utils.py (3)
compute_clever_covariate_ate(17-65)compute_initial_effect(147-209)estimate_fluctuation_parameter(212-288)
CausalEstimate/estimators/functional/tmle_att.py (1)
CausalEstimate/estimators/functional/utils.py (2)
compute_clever_covariate_att(68-119)estimate_fluctuation_parameter(212-288)
tests/test_functional/test_tmle_clipping.py (3)
CausalEstimate/estimators/functional/utils.py (2)
compute_clever_covariate_ate(17-65)compute_clever_covariate_att(68-119)CausalEstimate/estimators/functional/tmle.py (2)
compute_tmle_ate(19-41)compute_tmle_rr(44-81)CausalEstimate/estimators/functional/tmle_att.py (1)
compute_tmle_att(73-99)
CausalEstimate/estimators/tmle.py (2)
CausalEstimate/estimators/functional/tmle_att.py (1)
compute_tmle_att(73-99)CausalEstimate/estimators/functional/tmle.py (1)
compute_tmle_rr(44-81)
tests/test_functional/test_ipw_clipping.py (3)
CausalEstimate/estimators/functional/ipw.py (1)
compute_ipw_weights(171-242)tests/helpers/setup.py (1)
setUpClass(49-116)tests/test_functional/test_ipw.py (3)
setUpClass(23-30)setUpClass(58-64)setUpClass(124-127)
tests/test_functional/test_utils.py (1)
CausalEstimate/estimators/functional/utils.py (2)
compute_clever_covariate_ate(17-65)compute_clever_covariate_att(68-119)
🪛 Ruff (0.12.2)
CausalEstimate/estimators/functional/ipw.py
207-207: Avoid specifying long messages outside the exception class
(TRY003)
209-209: Avoid specifying long messages outside the exception class
(TRY003)
211-211: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (4)
tests/test_functional/test_tmle_clipping.py (1)
24-74: Solid, deterministic setup; covers extreme-ps regimes well.LGTM. No action needed.
CausalEstimate/estimators/tmle.py (1)
71-82: LGTM: clean pass-through of clip_percentile to functional APIsATE/ARR aliasing and ATT/RR branches look consistent.
CausalEstimate/estimators/functional/tmle_att.py (1)
56-63: Clipping scope for update_term_0: confirm intended threshold basisYou compute the threshold using controls only, then clip weight_component for all subjects. That attenuates treated updates too. If the intent is “clip only the control component,” restrict clipping to control indices; if the intent is “uniform cap,” compute the percentile over all subjects.
Option A (clip controls only — mirrors clever covariate construction):
- 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) + 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, + )Option B (uniform cap):
- 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) + threshold = np.percentile(weight_component, clip_percentile * 100) + weight_component = np.clip(weight_component, a_min=None, a_max=threshold)CausalEstimate/estimators/functional/tmle.py (1)
26-33: LGTM: ATE path cleanly delegates and returns consistent diagnosticsReturn payload includes EFFECT and initial/adjustment fields via compute_initial_effect; looks good.
| # --- 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) |
There was a problem hiding this comment.
Avoid logit(0) / logit(1) -> inf; clip Y1_hat/Y0_hat before logit
This will otherwise produce inf/NaN updates.
Apply:
- Q_star_1 = expit(logit(Y1_hat) + update_term_1)
- Q_star_0 = expit(logit(Y0_hat) + update_term_0)
+ eps = 1e-6
+ Y1_c = np.clip(Y1_hat, eps, 1 - eps)
+ Y0_c = np.clip(Y0_hat, eps, 1 - eps)
+ Q_star_1 = expit(logit(Y1_c) + update_term_1)
+ Q_star_0 = expit(logit(Y0_c) + update_term_0)📝 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.
| # --- 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) | |
| # --- Step 3: Apply the separate updates to the potential outcome models --- | |
| eps = 1e-6 | |
| Y1_c = np.clip(Y1_hat, eps, 1 - eps) | |
| Y0_c = np.clip(Y0_hat, eps, 1 - eps) | |
| Q_star_1 = expit(logit(Y1_c) + update_term_1) | |
| Q_star_0 = expit(logit(Y0_c) + update_term_0) |
🤖 Prompt for AI Agents
In CausalEstimate/estimators/functional/tmle_att.py around lines 66 to 68, the
code calls logit(Y1_hat)/logit(Y0_hat) directly which can produce inf/NaN when
predictions are exactly 0 or 1; before taking logit, clip Y1_hat and Y0_hat into
a safe range (e.g., [eps, 1-eps]) using a small epsilon like 1e-6 (or a
configurable constant), then compute the logit and apply the update terms so
Q_star_1 and Q_star_0 are numerically stable.
Introduced an `eps` parameter to improve robustness of IPW and TMLE estimators against division by zero when handling propensity scores. Updated functions and documentation accordingly.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
CausalEstimate/estimators/functional/utils.py (1)
128-151: Guard against non-finite clever covariate values.Warn (or raise) if H contains NaN/inf; current checks only look at magnitude thresholds.
def _validate_clever_covariate(H: np.ndarray, estimand: str) -> None: @@ + if not np.all(np.isfinite(H)): + warnings.warn( + f"Non-finite values detected in clever covariate H for {estimand}. " + "Check ps bounds and eps.", RuntimeWarning + )
♻️ Duplicate comments (4)
CausalEstimate/estimators/functional/utils.py (1)
20-22: Add input validation and clamp denominators with np.clip (prior review still applies).Validate A/ps shapes and ranges; enforce 0 < clip_percentile ≤ 1; and use np.clip on ps and (1-ps) instead of adding eps. This prevents negative or >1 ps from producing misleading weights and matches IPW guards.
def compute_clever_covariate_ate( A: np.ndarray, ps: np.ndarray, - clip_percentile: float = 1, + clip_percentile: float = 1.0, eps: float = 1e-9, ) -> np.ndarray: @@ - # Unstabilized clever covariate components - H1_component = 1.0 / (ps + eps) - H0_component = 1.0 / (1 - ps + eps) + # Validate and stabilize inputs + A = np.asarray(A) + ps = np.asarray(ps) + if A.shape != ps.shape: + raise ValueError("A and ps must have the same shape.") + if not (0 < clip_percentile <= 1.0): + raise ValueError("clip_percentile must be in (0, 1].") + if not np.all(np.isfinite(ps)): + raise ValueError("ps contains non-finite values.") + if np.any((ps < 0) | (ps > 1)): + raise ValueError("ps must be within [0, 1].") + + # Unstabilized components with clamped denominators + denom1 = np.clip(ps, eps, 1 - eps) + denom0 = np.clip(1 - ps, eps, 1 - eps) + H1_component = 1.0 / denom1 + H0_component = 1.0 / denom0Also applies to: 41-44
CausalEstimate/estimators/tmle.py (1)
20-22: Validate clip_percentile/eps at init (fail fast).Move validation from _compute_effect to init, cast to float, and assert eps > 0.
- self.clip_percentile = clip_percentile - self.eps = eps + if not (0 < clip_percentile <= 1): + raise ValueError("clip_percentile must be in (0, 1].") + if not (eps > 0): + raise ValueError("eps must be > 0.") + self.clip_percentile = float(clip_percentile) + self.eps = float(eps)Also applies to: 49-51
CausalEstimate/estimators/functional/tmle_att.py (1)
46-49: Bug: missing A/(…) and (1−A)/(…) gating and unsafe logit on 0/1.Update terms must include A and (1−A) factors; also clip Y1_hat/Y0_hat before logit to avoid inf/NaN.
- update_term_1 = epsilon * (1.0 / (p_treated + eps)) + update_term_1 = epsilon * (A / (p_treated + eps)) @@ - weight_component = ps / (p_treated * (1 - ps) + eps) + denom0 = np.clip(1 - ps, eps, 1 - eps) + weight_component = (1 - A) * ps / (p_treated * denom0) @@ - update_term_0 = -epsilon * weight_component + update_term_0 = -epsilon * weight_component @@ - Q_star_1 = expit(logit(Y1_hat) + update_term_1) - Q_star_0 = expit(logit(Y0_hat) + update_term_0) + logit_eps = 1e-6 + Y1_c = np.clip(Y1_hat, logit_eps, 1 - logit_eps) + Y0_c = np.clip(Y0_hat, logit_eps, 1 - logit_eps) + Q_star_1 = expit(logit(Y1_c) + update_term_1) + Q_star_0 = expit(logit(Y0_c) + update_term_0)Also applies to: 55-66, 68-69
CausalEstimate/estimators/functional/tmle.py (1)
99-102: Bug: TMLE updates omit A/(…) and (1−A)/(…) factors; unsafe logit on 0/1.Multiply H1/H0 by A and (1−A) respectively, and clip Y1_hat/Y0_hat before logit. Also pass A into update_estimates.
- Q_star_1, Q_star_0 = update_estimates(ps, Y0_hat, Y1_hat, epsilon) + Q_star_1, Q_star_0 = update_estimates(A, ps, Y0_hat, Y1_hat, epsilon) @@ -def update_estimates( - ps: np.ndarray, +def update_estimates( + A: np.ndarray, + ps: np.ndarray, Y0_hat: np.ndarray, Y1_hat: np.ndarray, epsilon: float, eps: float = 1e-9, ) -> Tuple[np.ndarray, np.ndarray]: @@ - 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) + H1 = A / (np.clip(ps, eps, 1 - eps)) + H0 = -(1.0 - A) / (np.clip(1.0 - ps, eps, 1 - eps)) + + logit_eps = 1e-6 + Y1_c = np.clip(Y1_hat, logit_eps, 1 - logit_eps) + Y0_c = np.clip(Y0_hat, logit_eps, 1 - logit_eps) + Q_star_1 = expit(logit(Y1_c) + epsilon * H1) + Q_star_0 = expit(logit(Y0_c) + epsilon * H0)Also applies to: 106-124
🧹 Nitpick comments (10)
CausalEstimate/estimators/functional/utils.py (1)
243-246: Naming nit: avoid shadowing eps across modules.Consider renaming local
eps = 1e-6tologit_epsfor clarity with module-wide eps.CausalEstimate/estimators/functional/ipw.py (6)
20-24: Docstring still mentions stabilized weights (not implemented anymore).Update module docstring to reflect percentile-based clipping and remove stabilized weight mention.
-We also provide an option to use stabilized weights as described in: -... -Estimating causal effects from epidemiological data -https://pubmed.ncbi.nlm.nih.gov/16790829/ +This module uses percentile-based clipping and an eps guard for numerical stability.
37-45: Handle μ0 near-zero/NaN in RR robustly.Use isclose and NaN checks to avoid misleading finite divisions.
- if mu_0 == 0: + if (not np.isfinite(mu_0)) or np.isclose(mu_0, 0.0, atol=1e-12): warnings.warn( - "Risk in untreated group (mu_0) is 0, returning inf for Risk Ratio.", + "Risk in untreated group (mu_0) is NaN or ≈0; returning inf for Risk Ratio.", RuntimeWarning, ) rr = np.inf
86-97: RRT near-zero/NaN guard.Same robustness improvement as RR.
- if mu_0 == 0: + if (not np.isfinite(mu_0)) or np.isclose(mu_0, 0.0, atol=1e-12): warnings.warn( - "Risk in counterfactual untreated group (mu_0) is 0, returning inf for RRT.", + "Risk in counterfactual untreated group (mu_0) is NaN or ≈0; returning inf for RRT.", RuntimeWarning, ) rrt = np.inf
112-118: Docstring mismatch: this is Hajek (ratio) not “simple Horvitz–Thompson.”Either change wording or implement HT (divide by N). Recommend doc fix.
- Computes E[Y(1)] and E[Y(0)] for the ATE using the simple Horvitz-Thompson estimator, - with explicit checks for empty groups. + Computes E[Y(1)] and E[Y(0)] for the ATE using the Hajek (ratio) estimator, + with explicit checks for empty groups.Also applies to: 120-122
246-253: Optional: clamp denominators with np.clip for consistency with TMLE utils.This avoids systematically increasing denominators by eps when ps∈(0,1).
- weights = A / (ps + eps) + (1 - A) / (1 - ps + eps) + denom1 = np.clip(ps, eps, 1 - eps) + denom0 = np.clip(1 - ps, eps, 1 - eps) + weights = A / denom1 + (1 - A) / denom0 @@ - weights = A + (1 - A) * ps / (1 - ps + eps) + denom0 = np.clip(1 - ps, eps, 1 - eps) + weights = A + (1 - A) * ps / denom0
239-243: Ruff TRY003: shorten exception messages or centralize text.Minor style nit; consider shorter messages or constants to satisfy linter.
CausalEstimate/estimators/tmle.py (1)
74-76: Avoid late validation; now covered in init.Once moved to init, remove this check.
- if not (0 < self.clip_percentile <= 1): - raise ValueError("clip_percentile must be in (0, 1].")CausalEstimate/estimators/functional/tmle_att.py (2)
37-45: Remove editorial comments.Drop “incorrect/new code” commentary to keep codebase neutral.
- # --- 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. + # --- Step 2: Define separate update terms ---
85-89: Docstring nit: duplicated eps description.Streamline to a single-line note about eps.
- with optional clipping for the control group. - eps: float = 1e-9, - Guard against division by zero + with optional clipping for the control group. eps guards against division by zero.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
CausalEstimate/estimators/functional/ipw.py(5 hunks)CausalEstimate/estimators/functional/tmle.py(5 hunks)CausalEstimate/estimators/functional/tmle_att.py(2 hunks)CausalEstimate/estimators/functional/utils.py(3 hunks)CausalEstimate/estimators/ipw.py(4 hunks)CausalEstimate/estimators/tmle.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- CausalEstimate/estimators/ipw.py
🧰 Additional context used
🧬 Code graph analysis (3)
CausalEstimate/estimators/tmle.py (2)
CausalEstimate/estimators/functional/tmle.py (2)
compute_tmle_ate(19-42)compute_tmle_rr(45-83)CausalEstimate/estimators/functional/tmle_att.py (1)
compute_tmle_att(74-103)
CausalEstimate/estimators/functional/tmle_att.py (1)
CausalEstimate/estimators/functional/utils.py (2)
compute_clever_covariate_att(71-125)estimate_fluctuation_parameter(218-294)
CausalEstimate/estimators/functional/tmle.py (1)
CausalEstimate/estimators/functional/utils.py (3)
compute_clever_covariate_ate(17-68)compute_initial_effect(153-215)estimate_fluctuation_parameter(218-294)
🪛 Ruff (0.12.2)
CausalEstimate/estimators/tmle.py
75-75: Avoid specifying long messages outside the exception class
(TRY003)
CausalEstimate/estimators/functional/ipw.py
239-239: Avoid specifying long messages outside the exception class
(TRY003)
241-241: Avoid specifying long messages outside the exception class
(TRY003)
243-243: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (8)
CausalEstimate/estimators/functional/ipw.py (4)
58-66: LGTM.Forwarding clip_percentile/eps is consistent and clear.
72-80: LGTM.ATT path consistently forwards clip_percentile/eps.
123-126: LGTM.Centralizing weight computation keeps semantics consistent.
151-157: LGTM.ATT weighted outcomes wiring looks correct.
CausalEstimate/estimators/tmle.py (1)
79-87: LGTM.Consistent forwarding of clip_percentile and eps to functional estimators.
Also applies to: 90-98, 101-109
CausalEstimate/estimators/functional/tmle_att.py (1)
27-35: LGTM.Using compute_clever_covariate_att with clipping/eps is aligned.
CausalEstimate/estimators/functional/tmle.py (2)
26-34: LGTM.ATE path is consistent with compute_estimates API.
Also applies to: 38-42
52-60: LGTM.RR path handles near-zero denominator with isclose and returns initial effects.
Also applies to: 78-83
Summary by CodeRabbit
New Features
Refactor
Tests
Breaking Changes