Skip to content

Commit 564527e

Browse files
Fix summary() crash when outcome model is fit without adjust()
_outcome_estimates_summary() unconditionally called aipw(), which raises when the frame has not been adjusted (aipw() needs adjust()-calibrated responder weights on the target population scale). The changelog documents the no-adjust() outcome-model workflow, so summary() on such a frame crashed with a ValueError instead of reporting. Omit the mu_DR/AIPW block when not adjusted while still showing mu_IPW and mu_OM, and add a regression test.
1 parent 39fe706 commit 564527e

2 files changed

Lines changed: 31 additions & 4 deletions

File tree

balance/balance_frame.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3752,8 +3752,6 @@ def _outcome_estimates_summary(self) -> str | None:
37523752
for col in model["outcome_columns"]
37533753
if f"{col}_hat" in target_predictions.columns
37543754
}
3755-
dr_estimates = {str(col): float(v) for col, v in self.aipw().items()}
3756-
37573755
def _fmt_points(estimates: dict[str, float]) -> str:
37583756
return "\n".join(
37593757
f" {col} {val:.3f}" for col, val in estimates.items()
@@ -3770,9 +3768,20 @@ def _fmt_points(estimates: dict[str, float]) -> str:
37703768
"Outcome model / g-computation (mu_OM), point estimate "
37713769
"(CI not shown; not cached):",
37723770
_fmt_points(om_estimates),
3773-
"Doubly robust / AIPW (mu_DR), point estimate:",
3774-
_fmt_points(dr_estimates),
37753771
]
3772+
3773+
# μ̂_DR requires adjust()-calibrated responder weights on the same
3774+
# population scale as the target (``aipw()`` enforces this). A frame
3775+
# with a fitted outcome model but no adjust() is a supported workflow
3776+
# (see the changelog), so omit the AIPW block instead of crashing.
3777+
if self.is_adjusted:
3778+
dr_estimates = {str(col): float(v) for col, v in self.aipw().items()}
3779+
blocks.extend(
3780+
[
3781+
"Doubly robust / AIPW (mu_DR), point estimate:",
3782+
_fmt_points(dr_estimates),
3783+
]
3784+
)
37763785
return "Outcome estimates:\n" + "\n".join(blocks)
37773786

37783787
def summary(self) -> str:

tests/test_aipw.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,24 @@ def test_summary_shows_ipw_om_dr_when_model_fit(self) -> None:
315315
# the rich section replaces the plain "Outcome weighted means" block
316316
self.assertNotIn("Outcome weighted means", summary)
317317

318+
def test_summary_omits_dr_when_model_fit_but_unadjusted(self) -> None:
319+
"""summary() with a fitted outcome model but NO adjust() must not crash.
320+
321+
The μ̂_DR/AIPW block requires adjust()-calibrated weights, so it is
322+
omitted (μ̂_IPW and μ̂_OM are still shown) rather than raising.
323+
"""
324+
sample_df, target_df = _make_aipw_fixture()
325+
s = Sample.from_frame(
326+
sample_df, id_column="id", weight_column="weight", outcome_columns=["y"]
327+
)
328+
t = Sample.from_frame(target_df, id_column="id", weight_column="weight")
329+
st = s.set_target(t).fit_outcome_model(model=LinearRegression())
330+
summary = st.summary() # previously raised ValueError from aipw()
331+
self.assertIn("Outcome estimates:", summary)
332+
self.assertIn("mu_IPW", summary)
333+
self.assertIn("mu_OM", summary)
334+
self.assertNotIn("mu_DR", summary)
335+
318336
def test_summary_unchanged_without_outcome_model(self) -> None:
319337
sample_df, target_df = _make_aipw_fixture()
320338
s = Sample.from_frame(

0 commit comments

Comments
 (0)