Fix unequal-fold PPI weighting

pappi/core/pipeline.py

@@ -190,6 +190,7 @@ class Split:
 
 ## Percentile interval from `n` replicates: labeled rows resampled by `groups` (default:
 ## the correction split's groups, else rows), deployment rows by row; propensity refit.
+## Resampling preserves correction folds and keeps label clusters intact.
 ## `arms` ("x", "z", "xz", "pi") limits it; an IWCV row follows its arm.
 @dataclass(eq=False)
 class Bootstrap:
@@ -436,14 +437,26 @@ def _groups(frame, col, rows):
 ## Draws run in order here, so n_jobs never changes results. Copies keep their cluster id,
 ## so a duplicate is never scored by a model that saw its twin.
 def _bootstrap(n: int, n_jobs: int, clusters: np.ndarray, n_u: int, rng,
-               fn) -> tuple[np.ndarray, np.ndarray]:
+               fn, fold_l: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
     _, inv = np.unique(clusters, return_inverse=True)
     members = np.split(np.argsort(inv, kind="stable"), np.cumsum(np.bincount(inv))[:-1])
     m = len(members)
+    pools = [np.arange(m)]
+    n_folds = 0 if fold_l is None else len(np.unique(fold_l))
+    if fold_l is not None:
+        cluster_folds = [np.unique(fold_l[rows]) for rows in members]
+        if all(len(f) == 1 for f in cluster_folds):
+            ids = np.array([f[0] for f in cluster_folds])
+            pools = [np.flatnonzero(ids == f) for f in np.unique(ids)]
 
     def draws():
         for _ in range(n):
-            il = np.concatenate([members[j] for j in rng.integers(0, m, m)])
+            while True:
+                selected = np.concatenate([rng.choice(pool, len(pool)) for pool in pools])
+                il = np.concatenate([members[j] for j in selected])
+                ## Cross-fold clusters stay intact; redraw if a correction fold disappears.
+                if fold_l is None or len(np.unique(fold_l[il])) == n_folds:
+                    break
             iu = rng.integers(0, n_u, n_u)
             yield il, iu
 
@@ -512,6 +525,7 @@ def _fit(data, z_cols, bar, fit_model, fit_propensity, metric, in_domain,
         if sp is correction:
             p, p_u, models = _oof(fit_model, x_l, y, ids, k, x_u)
             p_l = p
+            fold_l = ids
         else:
             p = _oof(fit_model, x_l, y, ids, k)
         if sp in in_domain:
@@ -521,15 +535,16 @@ def _fit(data, z_cols, bar, fit_model, fit_propensity, metric, in_domain,
     yl_hat, yu_hat = p_l.argmax(1), p_u.argmax(2)
 
     def est(w, il=slice(None), iu=slice(None)):
-        return ppi_estimate(metric, p_l[il], yl_hat[il], y[il], w, p_u[:, iu], yu_hat[:, iu], k)
+        return ppi_estimate(metric, p_l[il], yl_hat[il], y[il], w, p_u[:, iu], yu_hat[:, iu],
+                            k, fold_l[il])
 
     rng = np.random.default_rng(7)
 
     ## Estimators run on each arm's weights, as (key prefix, row names, fn): PPI, then IWCV.
     ests = [("", ARMS, est)]
     if iwcv:
         ests.append(("iwcv_", IWCV, lambda w, il=slice(None), iu=slice(None):
-                     iwcv_estimate(metric, yl_hat[il], y[il], w, k)))
+                     iwcv_estimate(metric, yl_hat[il], y[il], w, k, fold_l[il])))
 
     def points(w, il, iu):
         return [fn(w, il, iu)["point"] for *_, fn in ests]
@@ -559,7 +574,7 @@ def rep(il, iu, F=F):
                            p_assign(len(rr), rr if pg is None else pg[rr]), 2)[:, 1]
                 return points(inverse_odds_weights(q_b[:len(il)]), il, iu)
 
-            ci = _bootstrap(interval.n, n_jobs, clusters, n_u, rng, rep)
+            ci = _bootstrap(interval.n, n_jobs, clusters, n_u, rng, rep, fold_l)
         rows += arm_rows(key, inverse_odds_weights(q[:n_l]), ci,
                          auc=float(roc_auc_score(s, q)), q_l_median=float(np.median(q[:n_l])),
                          q_u_median=float(np.median(q[n_l:])),
@@ -571,7 +586,7 @@ def rep(il, iu, F=F):
         if interval is not None and interval.covers("pi"):
             ## Serial: a known-weights replicate is too cheap to ship to a process.
             ci = _bootstrap(interval.n, 1, clusters, n_u, rng,
-                            lambda il, iu: points(w_pi[il], il, iu))
+                            lambda il, iu: points(w_pi[il], il, iu), fold_l)
         rows += arm_rows("pi", w_pi, ci)
 
     if any(row.get("clipped") for row in rows):

pappi/core/ppi.py

@@ -116,8 +116,16 @@ def score(self, y, y_hat, k):
         return float(np.mean(y == y_hat))
 
 
-## Hájek-weighted mean of the rows of r, its covariance, and Kish's n_eff.
-def _hajek(w: np.ndarray, r: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
+## Hájek-weighted mean, covariance, and Kish's n_eff. With fold_l, normalize within
+## each correction fold, then combine fold means equally and their variances by 1/m**2.
+def _hajek(w: np.ndarray, r: np.ndarray,
+           fold_l: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray, float]:
+    if fold_l is not None:
+        parts = [_hajek(w[fold_l == f], r[fold_l == f]) for f in np.unique(fold_l)]
+        means, covariances, effective = zip(*parts)
+        m = len(parts)
+        return (np.mean(means, axis=0), np.sum(covariances, axis=0) / m**2,
+                float(m**2 / np.sum(1 / np.asarray(effective))))
     w_sum = float(w.sum())
     mu = (w[:, None] * r).sum(axis=0) / w_sum
     c = r - mu
@@ -136,7 +144,8 @@ def _delta(metric: Metric, mu: np.ndarray, Sigma: np.ndarray, K: int) -> dict:
 
 ## Eq. (2) of the paper: imputation over U + Hájek-weighted
 ## residual over L, delta-method 95% CI. Cross-fit: p_hat_u (models, n_u, k) and
-## y_hat_u (models, n_u) hold each fold model's predictions; U imputes their mean row.
+## y_hat_u (models, n_u) hold each fold model's predictions. fold_l identifies each
+## label's correction fold; imputation and correction both average folds equally.
 def ppi_estimate(
     metric: Metric,
     p_hat_l: np.ndarray,
@@ -146,8 +155,14 @@ def ppi_estimate(
     p_hat_u: np.ndarray,
     y_hat_u: np.ndarray,
     k_classes: int,
+    fold_l: np.ndarray | None = None,
 ) -> dict:
     K = k_classes
+    if p_hat_u.ndim == 3 and len(p_hat_u) > 1:
+        if fold_l is None:
+            raise ValueError("cross-fit estimates require fold_l for the correction labels")
+        if len(np.unique(fold_l)) != len(p_hat_u):
+            raise ValueError("every deployment fold model needs held-out correction labels")
     g_l = metric.expected(p_hat_l, y_hat_l, K)
     h_l = metric.observed(y_l, y_hat_l, K)
     r_l = h_l - g_l
@@ -159,7 +174,7 @@ def ppi_estimate(
     centered_u = g_u - mu_imp
     S_imp = centered_u.T @ centered_u / N
 
-    mu_res, V_res, n_eff = _hajek(w_l, r_l)
+    mu_res, V_res, n_eff = _hajek(w_l, r_l, fold_l)
     mu_total = mu_imp + mu_res
     return _delta(metric, mu_total, S_imp / N + V_res, K) | {
         "rates": mu_total, "rectifier": mu_res, "imputation": mu_imp, "n_eff": n_eff}
@@ -168,13 +183,15 @@ def ppi_estimate(
 ## IWCV: the same weights on the observed rates over L alone, with no deployment term.
 ## Delta-method interval as in ppi_estimate, without the imputation variance.
 def iwcv_estimate(metric: Metric, y_hat_l: np.ndarray, y_l: np.ndarray, w_l: np.ndarray,
-                  k_classes: int) -> dict:
-    mu, V, n_eff = _hajek(w_l, metric.observed(y_l, y_hat_l, k_classes))
+                  k_classes: int, fold_l: np.ndarray | None = None) -> dict:
+    mu, V, n_eff = _hajek(w_l, metric.observed(y_l, y_hat_l, k_classes), fold_l)
     return _delta(metric, mu, V, k_classes) | {"n_eff": n_eff}
 
 
-def ppi_macro_f1(p_hat_l, y_hat_l, y_l, w_l, p_hat_u, y_hat_u, k_classes: int) -> dict:
-    return ppi_estimate(MacroF1(), p_hat_l, y_hat_l, y_l, w_l, p_hat_u, y_hat_u, k_classes)
+def ppi_macro_f1(p_hat_l, y_hat_l, y_l, w_l, p_hat_u, y_hat_u, k_classes: int,
+                 fold_l: np.ndarray | None = None) -> dict:
+    return ppi_estimate(MacroF1(), p_hat_l, y_hat_l, y_l, w_l, p_hat_u, y_hat_u,
+                        k_classes, fold_l)
 
 
 ## w = (1-q)/q is proportional to p_U/p_L; 1/q is not.

tests/test_audit.py

@@ -5,7 +5,7 @@
 import pandas as pd
 import pytest
 
-from pappi import Accuracy, Protocol, Split, audit
+from pappi import Accuracy, Bootstrap, MacroF1, Protocol, Split, audit
 from pappi.cli import main
 from pappi.demo import synthetic_campaign, true_score
 
@@ -61,6 +61,50 @@ def predict_proba(self, X):
     assert row.ci_lo <= truth <= row.ci_hi
 
 
+@pytest.mark.parametrize("metric, truth", [
+    (Accuracy(), 0.5),
+    (MacroF1(), (8 / 13 + 2 / 7) / 2),
+])
+@pytest.mark.parametrize("clusters", [None, "fold", "cluster", "mixed"])
+def test_unequal_correction_folds_keep_model_target(metric, truth, clusters):
+    fold_ids = np.r_[np.full(900, -7), np.full(100, 13)]
+    frame = pd.DataFrame({
+        "x_0": np.zeros(2000),
+        "label": np.r_[np.tile(["A", "A", "A", "A", "B"], 200), [""] * 1000],
+        "fold": np.r_[fold_ids, np.zeros(1000)],
+        "cluster": np.arange(2000) % 50,
+        "mixed": np.r_[np.zeros(899), np.ones(101), np.zeros(1000)],
+        "p_selected": np.r_[np.ones(900), np.full(100, 0.2), np.ones(1000)],
+    })
+
+    class Constant:
+        def __init__(self, probabilities):
+            self.probabilities = probabilities
+
+        def predict_proba(self, X):
+            return np.tile(self.probabilities, (len(X), 1))
+
+    ## The models always predict A and B, respectively, on an 80/20 population.
+    def trainer(X, y):
+        return Constant([0.9, 0.1] if len(y) == 100 else [0.4, 0.6])
+
+    split = Split("unequal", groups="fold" if clusters == "fold" else None,
+                  assign=lambda n, g: fold_ids)
+    report = audit(
+        frame, fit_model=trainer, fit_propensity=lambda X, y: Constant([0.5, 0.5]),
+        in_domain=[split], inclusion="p_selected", metric=metric, iwcv=True,
+        interval=Bootstrap(30, groups=clusters, arms=("x", "pi")),
+    )
+    rows = report.frame.set_index("key")
+    for key in ("x", "pi", "iwcv_x", "iwcv_pi"):
+        row = rows.loc[key]
+        assert row.score == pytest.approx(truth)
+        assert row.ci_lo <= truth <= row.ci_hi
+        if isinstance(metric, Accuracy):
+            assert row.sigma == pytest.approx(np.sqrt(0.8 * 0.2 / 4 * (1 / 900 + 1 / 100)))
+        assert row.boot_lo - 1e-12 <= truth <= row.boot_hi + 1e-12
+
+
 def test_custom_propensity_assign_without_groups_is_used():
     calls = []
 

tests/test_ppi.py

@@ -2,8 +2,9 @@
 # Licensed under the MIT License.
 
 import numpy as np
+import pytest
 
-from pappi.core.ppi import MacroF1, inclusion_weights, macro_f1, ppi_estimate
+from pappi.core.ppi import Accuracy, MacroF1, inclusion_weights, macro_f1, ppi_estimate
 
 
 ## The estimator alone, with known inclusion weights. The model is calibrated on near rows
@@ -32,3 +33,11 @@ def estimate(w_fn, rng):
     assert 0.90 <= weighted[:, 1].mean() <= 0.99
     assert abs(weighted[:, 0].mean() - truth) < 0.005
     assert unweighted[:, 1].mean() < 0.5
+
+
+def test_crossfit_estimate_requires_correction_fold_ids():
+    p = np.tile([0.9, 0.1], (2, 1))
+    p_u = np.stack([p, p])
+    with pytest.raises(ValueError, match="fold_l"):
+        ppi_estimate(Accuracy(), p, p.argmax(1), np.array([0, 1]), np.ones(2),
+                     p_u, p_u.argmax(2), 2)