Enforce the AIPW weight normalization contract (#551)

Summary: Pull Request resolved: #551

Differential Revision: D114518701

Pulled By: talgalili

fbshipit-source-id: 57d5926d0eab905a50b6c07483d21893c1452f24

Author: neuralsorcerer

CHANGELOG.md

@@ -4,6 +4,14 @@
 
 - **Breaking:** the previously-inert `predicted_outcome_columns` parameter (`Sample.from_frame` / `SampleFrame.from_frame` / `SampleFrame._create`) and the `SampleFrame.predicted_outcome_columns` property are renamed to `outcomes_hat_columns` (internal `_column_roles` key `"predicted"` → `"outcomes_hat"`). The old names are removed outright with no alias — the role was reserved scaffolding, never populated or consumed, so no migration is provided.
 
+- **Breaking:** `BalanceFrame.aipw()` now requires responder weights produced by
+  `adjust()`, validates both weight vectors as non-empty, one-dimensional,
+  finite, and non-negative with positive finite totals, and verifies that the
+  responder total matches the target-weight total to a relative tolerance of
+  `1e-6`. This enforces the estimator's same-population-scale normalization
+  contract instead of accepting arbitrary, uncalibrated weights; call
+  `adjust(...)` before `aipw()` and do not rescale its output.
+
 ## New Features
 
 - **Outcome-model input and replay validation is hardened.** `fit_outcome_model(..., sample_weight=...)` rejects incorrectly shaped arrays and zero, negative, NaN, infinite, non-real, or non-numeric weights with actionable `ValueError` messages; accepted weights must be one-dimensional, finite, strictly positive real numbers aligned to the covariates. Replaying categorical covariates now explicitly maps novel levels to missing before constructing the frozen categorical dtype, avoiding the pandas deprecation warning while preserving the existing novel-level behavior.
@@ -74,7 +82,7 @@
   scored.outcomes_hat().mean()      # μ̂_OM on the holdout target via train_bf's fitted model
   ```
 
-- **`BalanceFrame.aipw()` — doubly-robust (AIPW) estimate `μ̂_DR`.** New `BalanceFrame.aipw()` (and, via the MRO, `Sample.aipw()`) returns the augmented / one-sample AIPW estimate of the target-population mean, per outcome column, combining the fitted outcome model `ĝ` with the balance weights `w`: `μ̂_DR = wmean(ĝ(X_T), w_T) + wmean(Y − ĝ(X_S), w)` (the augmentation runs over responders with an observed `Y`). It is **doubly robust** — consistent if *either* the outcome model *or* the weighting model is correct — and completes the estimator trio alongside `outcomes().mean()` (`μ̂_IPW`) and `outcomes_hat().mean()` (`μ̂_OM`); equivalently it is a GREG (model-assisted) estimator with the balance weights as the design weights. It requires a fitted outcome model (`fit_outcome_model(...)`) **and** a target (`set_target(...)`), accepts any balance weights, and **warns** when the responder weights are constant (no weighting fitted → `μ̂_DR` reduces to `μ̂_OM`). The point-estimate arithmetic lives in the pure `balance.outcome_models.aipw_point_estimate(...)`. **This is the point estimate only** — no confidence interval yet; an honest AIPW interval must jointly capture the weighting- and outcome-model uncertainty (see the TODOs in `balance/outcome_models/aipw.py`: cross-fitting, an analytic influence-function / sandwich SE, and an end-to-end joint bootstrap).
+- **`BalanceFrame.aipw()` — doubly-robust (AIPW) estimate `μ̂_DR`.** New `BalanceFrame.aipw()` (and, via the MRO, `Sample.aipw()`) returns the augmented / one-sample AIPW estimate of the target-population mean, per outcome column, combining the fitted outcome model `ĝ` with the balance weights `w`: `μ̂_DR = wmean(ĝ(X_T), w_T) + wmean(Y − ĝ(X_S), w)` (the augmentation runs over responders with an observed `Y`). It is **doubly robust** — consistent if *either* the outcome model *or* the weighting model is correct — and completes the estimator trio alongside `outcomes().mean()` (`μ̂_IPW`) and `outcomes_hat().mean()` (`μ̂_OM`); equivalently it is a GREG (model-assisted) estimator with the balance weights as the design weights. It requires a fitted outcome model (`fit_outcome_model(...)`), a target (`set_target(...)`), and `adjust()`-calibrated responder weights whose total matches the target-weight total within relative tolerance `1e-6`; this same-population-scale requirement makes the doubly-robust claim asymptotic under the Hájek normalization used by balance (ratio bias is `O(1/n)`). It **warns** when the calibrated responder weights are constant (`μ̂_DR` reduces to `μ̂_OM`). The point-estimate arithmetic lives in the pure `balance.outcome_models.aipw_point_estimate(...)`. **This is the point estimate only** — no confidence interval yet; an honest AIPW interval must jointly capture the weighting- and outcome-model uncertainty (see the TODOs in `balance/outcome_models/aipw.py`: cross-fitting, an analytic influence-function / sandwich SE, and an end-to-end joint bootstrap).
 
   ```python
   bf = sample.adjust(method="ipw").set_target(target)   # or any balance weights

balance/balance_frame.py

@@ -3540,9 +3540,12 @@ def aipw(self) -> pd.Series:
         ``outcomes().mean()`` gives ``μ̂_IPW`` and ``outcomes_hat().mean()`` gives
         ``μ̂_OM``; this gives ``μ̂_DR``.
 
-        Requires a fitted outcome model **and** a target (:meth:`set_target`).
-        Uses whatever weight columns are present (any balance weights). With
-        constant responder weights it warns and reduces to ``μ̂_OM``.
+        Requires a fitted outcome model, a target (:meth:`set_target`), and
+        responder weights produced by :meth:`adjust`. The adjusted responder
+        weights and target weights must have matching totals (within a relative
+        tolerance of ``1e-6``), ensuring that both terms use weights on the
+        same target-population scale. With constant responder weights it warns
+        and reduces to ``μ̂_OM``.
 
         Point estimate only -- no confidence interval (an honest AIPW interval
         must jointly capture the weighting- and outcome-model uncertainty; see
@@ -3553,7 +3556,9 @@ def aipw(self) -> pd.Series:
 
         Raises:
             ValueError: If no outcome model has been fit, if no target is set,
-                or if the responders carry no observed outcomes.
+                if the responders have not been adjusted, if the responder and
+                target weight totals are not on the same scale, or if the
+                responders carry no observed outcomes.
 
         Examples:
             >>> import pandas as pd
@@ -3569,12 +3574,15 @@ def aipw(self) -> pd.Series:
             >>> tgt = SampleFrame.from_frame(
             ...     pd.DataFrame({"id": [5, 6], "x": [15.0, 35.0],
             ...                   "weight": [1.0, 1.0]}))
-            >>> bf = BalanceFrame(sample=resp, target=tgt)
+            >>> bf = BalanceFrame(sample=resp, target=tgt).adjust(method="ipw")
             >>> _ = bf.fit_outcome_model(model=LinearRegression())
             >>> bf.aipw().index.tolist()
             ['y']
         """
-        from balance.outcome_models.aipw import aipw_point_estimate
+        from balance.outcome_models.aipw import (
+            _validate_aipw_weight_scale,
+            aipw_point_estimate,
+        )
 
         model = self.outcome_model
         if model is None:
@@ -3586,6 +3594,12 @@ def aipw(self) -> pd.Series:
             raise ValueError(
                 "aipw() requires a target population; call set_target(...) first."
             )
+        if not self.is_adjusted:
+            raise ValueError(
+                "aipw() requires adjust()-calibrated responder weights; call "
+                "adjust(...) before aipw() so responder and target weights are "
+                "on the same population scale."
+            )
         observed_outcomes = self._outcome_columns
         if observed_outcomes is None:
             raise ValueError(
@@ -3595,24 +3609,25 @@ def aipw(self) -> pd.Series:
 
         target = _assert_type(self._sf_target)
         sample_weight = self.weight_series
-        if (
-            sample_weight is not None
-            and len(sample_weight) > 1
-            and sample_weight.nunique() == 1
-        ):
+        target_weight = target.weight_series
+        if sample_weight is None or target_weight is None:
+            raise ValueError(
+                "aipw() requires responder and target weight columns on the same "
+                "population scale."
+            )
+        _validate_aipw_weight_scale(sample_weight, target_weight)
+        if len(sample_weight) > 1 and sample_weight.nunique() == 1:
             logger.warning(
-                "aipw(): responder weights are constant -- it appears no "
-                "weighting model was fit (adjust() was not run, or it produced "
-                "uniform weights); the AIPW estimate reduces to the "
-                "outcome-model estimate mu_OM."
+                "aipw(): adjusted responder weights are constant; the AIPW "
+                "estimate reduces to the outcome-model estimate mu_OM."
             )
 
         estimates = aipw_point_estimate(
             self._sf_sample.df_covars,
             observed_outcomes,
             sample_weight,
             target.df_covars,
-            target.weight_series,
+            target_weight,
             model,
         )
         return pd.Series(estimates, dtype=float)

balance/outcome_models/aipw.py

@@ -30,11 +30,18 @@
 
 This module provides the **point estimate only** (see the TODOs below for
 cross-fitting and honest variance/CI).
+
+The public :meth:`balance.balance_frame.BalanceFrame.aipw` entry point enforces
+the estimator's normalization contract: responder weights must come from
+``adjust()`` and their total must match the target-weight total. Direct callers
+of :func:`aipw_point_estimate` are responsible for supplying weights on that
+same target-population scale.
 """
 
 from __future__ import annotations
 
 import logging
+import math
 from typing import Any, Dict, List
 
 import numpy as np
@@ -44,6 +51,8 @@
 
 logger: logging.Logger = logging.getLogger(__package__)
 
+_AIPW_WEIGHT_SUM_RTOL: float = 1e-6
+
 # TODO (cross-fitting): the augmentation uses in-sample ĝ(X_S) — the model was
 # fit on these same responders — which is optimistic for flexible learners. Add
 # K-fold cross-fitted (out-of-fold) predictions for the residual term, and
@@ -68,6 +77,70 @@
 #   harness; it is a prerequisite for a fully honest .summary() interval.
 
 
+def _validate_aipw_weight_scale(
+    sample_weight: pd.Series | np.ndarray,
+    target_weight: pd.Series | np.ndarray,
+) -> None:
+    """Validate the same-population-scale contract for public AIPW estimates.
+
+    Zero-valued row weights are valid (for example, uncovered cells can receive
+    zero weight), but both vectors must be non-empty, one-dimensional, finite,
+    non-negative, and have positive totals. Their totals must differ by less
+    than the internal ``1e-6`` tolerance, relative to the target total.
+
+    Args:
+        sample_weight: Adjusted responder weights.
+        target_weight: Target design weights.
+
+    Raises:
+        ValueError: If either vector or the relationship between their totals
+            violates the AIPW normalization contract.
+    """
+
+    arrays: dict[str, np.ndarray] = {}
+    for name, weight in (
+        ("responder", sample_weight),
+        ("target", target_weight),
+    ):
+        try:
+            array = np.asarray(weight, dtype=float)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"aipw() requires numeric {name} weights.") from exc
+        if array.ndim != 1 or array.size == 0:
+            raise ValueError(
+                f"aipw() requires a non-empty, one-dimensional {name} weight vector."
+            )
+        if not np.isfinite(array).all():
+            raise ValueError(f"aipw() requires finite {name} weights.")
+        if (array < 0).any():
+            raise ValueError(f"aipw() requires non-negative {name} weights.")
+        arrays[name] = array
+
+    try:
+        sample_weight_total = math.fsum(arrays["responder"])
+        target_weight_total = math.fsum(arrays["target"])
+    except OverflowError as exc:
+        raise ValueError(
+            "aipw() requires finite responder and target weight totals."
+        ) from exc
+    if not math.isfinite(sample_weight_total) or not math.isfinite(target_weight_total):
+        raise ValueError("aipw() requires finite responder and target weight totals.")
+    if sample_weight_total <= 0 or target_weight_total <= 0:
+        raise ValueError("aipw() requires positive responder and target weight totals.")
+
+    relative_difference = (
+        abs(sample_weight_total - target_weight_total) / target_weight_total
+    )
+    if relative_difference >= _AIPW_WEIGHT_SUM_RTOL:
+        raise ValueError(
+            "aipw() requires adjust()-calibrated responder and target weights "
+            "on the same population scale: the relative weight-total "
+            f"difference is {relative_difference:.6g}, which must be less than "
+            f"{_AIPW_WEIGHT_SUM_RTOL:g}. Re-run adjust(...) without changing "
+            "the resulting weights."
+        )
+
+
 def aipw_point_estimate(
     sample_covars: pd.DataFrame,
     outcomes: pd.DataFrame,
@@ -91,7 +164,8 @@ def aipw_point_estimate(
             ``model["outcome_columns"]``. ``NaN`` rows are dropped from the
             residual term (weights realigned), matching ``fit_outcome_model``.
         sample_weight: Responder (balance) weights ``w``, or ``None`` for an
-            unweighted augmentation.
+            unweighted augmentation. Direct callers must ensure that these are
+            on the same population scale as ``target_weight``.
         target_covars: Target covariates ``X_T``.
         target_weight: Target weights ``w_T``, or ``None`` for a simple mean.
         model: A fitted model dict from :func:`fit_outcome_model`.

docs/architecture/architecture_0_23_0.md

@@ -18,6 +18,7 @@
 
 | Date | Decision | Rationale |
 |------|----------|-----------|
+| 2026-08-01 | **AIPW normalization contract: implemented.** The augmentation is valid only when `w_R` and `w_T` are on the same population scale; Hájek self-normalization gives *asymptotic* DR (ratio bias `O(1/n)`). `aipw()` requires valid, non-negative `adjust()`-calibrated responder weights and asserts `|Σw_R−Σw_T|/Σw_T < 1e-6`. | Enforces the correctness precondition at the estimator boundary and reports an actionable error for malformed, arbitrary, or subsequently rescaled weights. |
 | 2026-07-15 | **`outcomes_hat` is the canonical spelling everywhere — public and internal — as a clean rename (no alias, no `FutureWarning`, no migration).** Public: `outcomes_hat()`, `df_outcomes_hat`, `outcomes_hat_columns`. Internal: `_column_roles` key `"predicted"` → `"outcomes_hat"`, the `_create()`/`from_frame()` parameter `predicted_outcome_columns` → `outcomes_hat_columns`, internal locals (`predicted_list` → `outcomes_hat_list`, etc.), the overlap-validation dict key, and the protocol member `_outcomes_hat_columns`. The old `predicted_outcome_columns` param/property and the `"predicted"` role key are **removed outright**. | One vocabulary end-to-end; removes the overload with predicted-*weights* (`predict_weights`); distinct from `.outcomes()` (observed Y). No deprecation needed — the `predicted` role was reserved scaffolding **unused by any code or user** (only storage/validation/tests referenced it) and was never populated, so there is nothing to keep back-compatible and nothing to migrate. |
 | 2026-07-15 | **Scope: outcome model now, AIPW later.** Ship IPW/Hájek (exists) + outcome-model / g-computation estimate (`μ̂_OM`) in 0.23.0; defer explicit AIPW/DR to a follow-up. Phase 1 delivers **no general DR** — only the linear-WLS special case (§1). | Smaller, reviewable diff stack that delivers the core ask (average `outcomes_hat` on the target). AIPW variance/CI + the normalization contract get their own design pass. (superseded 2026-07-18: explicit AIPW shipped as `bf.aipw()` — general for any learner, not only the linear-WLS special case; only AIPW variance/CI + cross-fitting remain deferred.) |
 | 2026-07-15 | **Estimate lives on the `outcomes_hat()` view:** `μ̂_OM = bf.outcomes_hat().mean()` (target row). Not under `.outcomes()`, no separate `outcome_estimate(method=)` dispatcher for phase 1. | Keeps the estimate where the predicted outcomes live; reuses the existing weighted-mean/CI machinery for free. |
@@ -813,7 +814,7 @@ absolute GitHub URLs; the tutorial notebook (created in diff 7, extended in 8–
 end-to-end in CI.
 
 **Phase 2 (separate later stack):**
-- **`[balance] Add doubly-robust AIPW estimate bf.aipw() + R oracles`** (D112679814) — **shipped** the AIPW point estimate (combine `outcomes_hat` + IPW weights). Still deferred on top of it: the same-scale `w_R`/`w_T` normalization **assert**, **cross-fitted** (out-of-fold) responder residuals, and AIPW variance/CI.
+- **`[balance] Add doubly-robust AIPW estimate bf.aipw() + R oracles`** (D112679814) — **shipped** the AIPW point estimate (combine `outcomes_hat` + IPW weights) and its same-scale `w_R`/`w_T` normalization assertion. Still deferred on top of it: **cross-fitted** (out-of-fold) responder residuals and AIPW variance/CI.
 - Optional follow-ons: weighted/Bayesian bootstrap; BCa intervals.
 
 ---
@@ -845,7 +846,7 @@ guard; CLI out-of-scope; expanded tests/errors/docs (§12).
 default — `summary()` reports the fit-weights and *scopes* any DR statement to them; `outcomes_hat`
 columns use the `<outcome>_hat` convention + a `from_frame` leak-warning; `mean_with_ci` defaults to
 bootstrap and **raises on a lone/target-less view**; the AIPW normalization contract (same-scale
-weights; asymptotic Hájek DR) is documented for phase 2; and lifecycle fixes — deepcopy preserves
+weights; asymptotic Hájek DR) is enforced; and lifecycle fixes — deepcopy preserves
 the model (**reference-sharing** the estimator), `set_target` preserves it, re-fit drops stale Ŷ,
 and the bootstrap refits with the stored fit-weighting. Plus coherence fixes: Hájek sweep, demoted
 the special case, marked `mean_with_ci` overridden, disambiguated `w_R^fit` vs `w_R`, and added a
@@ -863,9 +864,9 @@ Deferred to implementation / later phases (not blocking this design):
 - **AIPW / doubly-robust** — the point estimate has **shipped** as `bf.aipw()` (BalanceFrame/Sample;
   pure `aipw_point_estimate`), general for any learner; only its **variance/CI** and **cross-fitting**
   remain deferred (each needs its own design).
-- **AIPW normalization contract** — **not yet enforced**: the shipped `aipw()` uses the present
-  balance weights and only warns on constant responder weights; the same-scale `w_R`/`w_T`
-  requirement + the `|Σw_R−Σw_T|/Σw_T < tol` tolerance assert remain a TODO (treat DR as asymptotic, Hájek `O(1/n)`).
+- **AIPW normalization contract** — **implemented**: `aipw()` requires an adjusted frame and
+  rejects responder/target weight totals when `|Σw_R−Σw_T|/Σw_T >= 1e-6`, enforcing same-scale
+  `w_R`/`w_T` (with DR asymptotic under Hájek normalization and ratio bias `O(1/n)`).
 - **Weighted / Bayesian bootstrap** — a robustness upgrade over the nonparametric bootstrap.
 - **BCa intervals** — an accuracy upgrade over percentile CIs if warranted.
 - **Cross-fitting** — to reduce own-observation bias when the same responders fit `ĝ` and enter