diff --git a/pappi/core/pipeline.py b/pappi/core/pipeline.py index 77ad91e..8b2ff6a 100644 --- a/pappi/core/pipeline.py +++ b/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,7 +535,8 @@ 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) @@ -529,7 +544,7 @@ def est(w, il=slice(None), iu=slice(None)): 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): diff --git a/pappi/core/ppi.py b/pappi/core/ppi.py index f96a231..ee07045 100644 --- a/pappi/core/ppi.py +++ b/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. diff --git a/tests/test_audit.py b/tests/test_audit.py index b27937f..5c3f4ef 100644 --- a/tests/test_audit.py +++ b/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 = [] diff --git a/tests/test_ppi.py b/tests/test_ppi.py index 4545586..7ffdfb5 100644 --- a/tests/test_ppi.py +++ b/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)