Skip to content

Feat/update interface - #83

Merged
kirilklein merged 9 commits into
mainfrom
feat/update_interface
Sep 13, 2025
Merged

Feat/update interface#83
kirilklein merged 9 commits into
mainfrom
feat/update_interface

Conversation

@kirilklein

@kirilklein kirilklein commented Sep 13, 2025

Copy link
Copy Markdown
Owner

Summary by CodeRabbit

  • New Features

    • Added clip_percentile to bound extreme weights/covariates across IPW and TMLE (ATE/ATT/RR) and introduced eps for numerical stability.
  • Refactor

    • Unified, vectorized IPW weighting with group-aware clipping; replaced stabilization toggle with clipping-based TMLE/covariate workflows.
  • Tests

    • Added comprehensive clipping tests (including deterministic exact-value checks) and removed stabilization-focused tests.
  • Breaking Changes

    • Removed stabilized parameter from public APIs; new clip_percentile and eps parameters replace it.

…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.
@coderabbitai

coderabbitai Bot commented Sep 13, 2025

Copy link
Copy Markdown
Contributor

Walkthrough

Replaces 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

Cohort / File(s) Summary
IPW functional weighting
CausalEstimate/estimators/functional/ipw.py
Replaced per-branch weight logic with a unified, vectorized implementation for ATE/ATT that uses an eps guard; added compute_ipw_weights with input validation (weight_type, clip_percentile, matching shapes); unified percentile-based, group-aware clipping; propagated eps through all weight/estimator helpers and updated signatures/docstrings.
TMLE functional (ATE / RR)
CausalEstimate/estimators/functional/tmle.py
Removed stabilized branching; added clip_percentile and eps parameters; compute path now builds clever covariates using 1/(ps+eps) and 1/(1-ps+eps) with optional percentile clipping; simplified update flow by removing pi/stabilized logic; updated function signatures and docs.
TMLE functional (ATT)
CausalEstimate/estimators/functional/tmle_att.py
Replaced stabilized-based ATT update with clipping/eps-based covariate and update logic; added degenerate-treated handling (no treated units); control update uses a weight component ps/(p_treated*(1-ps)+eps) with optional clipping; updated signatures and docs.
Clever covariates utilities
CausalEstimate/estimators/functional/utils.py
Replaced stabilized toggle with clip_percentile and eps for both ATE/ATT clever covariates; compute per-group components (1/(ps+eps), 1/(1-ps+eps), and ATT control weight) and apply per-group percentile clipping when requested; updated docstrings and signatures.
Estimator classes
CausalEstimate/estimators/ipw.py, CausalEstimate/estimators/tmle.py
Removed stabilized constructor parameter; added/propagated clip_percentile and new eps attribute; added validation for clip_percentile; internal calls updated to pass clip_percentile and eps to functional helpers.
IPW clipping tests
tests/test_functional/test_ipw_clipping.py
Added extensive tests validating ATE/ATT clipping behavior (percentile-bound assertions, deterministic exact-value checks, group-aware assertions); minor comment edits retained.
TMLE clipping tests (new)
tests/test_functional/test_tmle_clipping.py
New test module validating clipping for TMLE ATE/ATT/RR (clever-covariate clipping, finiteness, no-clipping equivalence, edge-ps scenarios); added deterministic checks and varied percentile cases.
TMLE stabilization tests (removed)
tests/test_functional/test_tmle/test_stabilization.py
Deleted legacy tests for stabilized TMLE behavior (stabilized=True paths).
Utils tests updated
tests/test_functional/test_utils.py
Rewrote tests to exercise clip_percentile semantics instead of stabilized paths; added boundary/extreme-value checks and removed stabilized-specific assertions.

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
Loading
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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60–90 minutes

Possibly related PRs

  • Fix/stabilized weights #82 — Overlaps changes to compute_ipw_weights, IPW/TMLE utility functions, and migration from stabilized logic to clipping-based logic.
  • Feat/stabilized weights #80 — Refactors IPW weighting and consolidates weight computation; closely related to the new vectorized/centralized IPW implementation.
  • 68 add rr #69 — Modifies TMLE functional APIs (compute_estimates/update_estimates/compute_tmle_*); intersects with TMLE signature and internal changes here.

Poem

I nibble on ps and hop through code,
Clip-percentile carrots line my road.
Weights now tidy, eps keeps them kind,
Covariates clipped so nothing’s misaligned.
Tests stamp my paw — hop, clip, compute, unwind. 🥕🐇

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check ✅ Passed The title "Feat/update interface" accurately reflects the PR's primary intent — broad interface/signature changes across multiple estimators (IPW, TMLE, utils) — so it is related to the main changeset. It is concise and not misleading, but relatively generic and does not call out the most important specifics (e.g., replacing the "stabilized" flag with clip_percentile/eps and adding eps to IPW). For a quick scan it communicates the purpose but could be more descriptive for future readers.
Docstring Coverage ✅ Passed Docstring coverage is 84.21% which is sufficient. The required threshold is 80.00%.
✨ Finishing touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch feat/update_interface

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 legacy stabilized kwarg in TMLE.init

README references stabilized=True (README.md:69) — accept legacy stabilized via **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 share

Division 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_hat
CausalEstimate/estimators/functional/tmle.py (1)

70-75: Nit: make RR “unrealistically large” threshold configurable

Hardcoding 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

📥 Commits

Reviewing files that changed from the base of the PR and between 5a1a2d7 and 163d053.

📒 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 APIs

ATE/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 basis

You 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 diagnostics

Return payload includes EFFECT and initial/adjustment fields via compute_initial_effect; looks good.

Comment on lines +66 to 68
# --- 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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

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.

Suggested change
# --- 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.

Comment thread CausalEstimate/estimators/functional/tmle.py Outdated
Comment thread CausalEstimate/estimators/functional/utils.py
Comment thread CausalEstimate/estimators/tmle.py
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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 / denom0

Also 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-6 to logit_eps for 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

📥 Commits

Reviewing files that changed from the base of the PR and between 3b3b6e9 and a287e57.

📒 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

Comment thread CausalEstimate/estimators/functional/ipw.py
Comment thread CausalEstimate/estimators/functional/utils.py
@kirilklein
kirilklein merged commit 198bae4 into main Sep 13, 2025
4 checks passed
@kirilklein
kirilklein deleted the feat/update_interface branch September 13, 2025 14:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant