1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Paired inference between arms (docs/evaluation_protocol.md §4).34* Block bootstrap: blocks = mun_code × calendar month (fallback mun_code, then i.i.d.), B = 2000, percentile CI5 and a centred bootstrap p-value for Δ = metric(B) − metric(A) computed on the same resampled blocks for both arms.6* Holm step-down correction for a family of contrasts.7* Diebold–Mariano test on the per-record loss differential d_i = |e_i^B| − |e_i^A| with a Newey–West (HAC) variance.8 The HAC variance is only meaningful when the series is ordered in time; the function therefore sorts by `date`9 and refuses to run when no date is available (`dm=False` in `compare_arms` skips it).10"""11from __future__ import annotations1213from collections.abc import Callable, Sequence14from typing import Any1516import numpy as np17import pandas as pd18from scipy import stats1920from ._common import arm_filter, arm_label, block_ids21from .metrics import PRIMARY_METRIC, metric_statistic, per_record2223Statistic = Callable[[pd.DataFrame], float]242526def _resolve_statistic(metric: str | Statistic) -> tuple[str, Statistic]:27 if callable(metric):28 return getattr(metric, "__name__", "custom"), metric29 return metric, metric_statistic(metric)303132def paired_records(preds_all: Any, truth: Any, spec_a: str | dict, spec_b: str | dict) -> tuple[pd.DataFrame, pd.DataFrame]:33 """Per-record frames of arms A and B restricted to the common tx_ids, in the same order."""34 d = per_record(preds_all, truth)35 a = arm_filter(d, spec_a).set_index("tx_id")36 b = arm_filter(d, spec_b).set_index("tx_id")37 common = a.index.intersection(b.index)38 if len(common) == 0:39 raise ValueError(f"arms {arm_label(spec_a)} and {arm_label(spec_b)} share no tx_id")40 return a.loc[common].reset_index(), b.loc[common].reset_index()414243def block_bootstrap_contrast(44 a: pd.DataFrame,45 b: pd.DataFrame,46 statistic: str | Statistic = PRIMARY_METRIC,47 *,48 B: int = 2000,49 seed: int = 20261027,50 alpha: float = 0.05,51 blocks: pd.Series | None = None,52) -> dict[str, float]:53 """Paired block bootstrap of Δ = stat(b) − stat(a) (negative Δ = arm B better for a loss).5455 `a` and `b` are per-record frames aligned on the same tx_ids (same order). Blocks are resampled with56 replacement (Künsch-style moving blocks are not used: blocks are the natural clusters municipality × month).57 Returns point estimate, percentile CI, centred bootstrap p-value and the bootstrap SE.58 """59 if len(a) != len(b) or not (a["tx_id"].to_numpy() == b["tx_id"].to_numpy()).all():60 raise ValueError("a and b must be aligned on the same tx_ids")61 name, stat = _resolve_statistic(statistic)62 blk = (blocks if blocks is not None else block_ids(a)).to_numpy()63 uniq, inv = np.unique(blk, return_inverse=True)64 n_blocks = len(uniq)65 members = [np.flatnonzero(inv == k) for k in range(n_blocks)]66 rng = np.random.default_rng(seed)67 delta_hat = stat(b) - stat(a)68 nan_result = {"metric": name, "n": float(len(a)), "n_blocks": float(n_blocks), "B": 0.0, "value_a": float(stat(a)), "value_b": float(stat(b)),69 "diff": float(delta_hat), "ci_lo": np.nan, "ci_hi": np.nan, "se_boot": np.nan, "p_boot": np.nan}70 if not np.isfinite(delta_hat): # e.g. coverage contrast involving an arm without intervals71 return nan_result72 draws = np.empty(B)73 for r in range(B):74 picked = rng.integers(0, n_blocks, size=n_blocks)75 idx = np.concatenate([members[k] for k in picked])76 draws[r] = stat(b.iloc[idx]) - stat(a.iloc[idx])77 draws = draws[np.isfinite(draws)]78 if len(draws) < 10:79 return nan_result80 lo, hi = np.percentile(draws, [100 * alpha / 2, 100 * (1 - alpha / 2)])81 # centred bootstrap p-value (Efron & Tibshirani 1993 §15): shift the bootstrap distribution to the null82 centred = draws - delta_hat83 p = float(np.mean(np.abs(centred) >= abs(delta_hat)))84 p = max(p, 1.0 / (len(draws) + 1)) # never report an exact zero85 return {86 "metric": name, "n": float(len(a)), "n_blocks": float(n_blocks), "B": float(len(draws)),87 "value_a": float(stat(a)), "value_b": float(stat(b)), "diff": float(delta_hat),88 "ci_lo": float(lo), "ci_hi": float(hi), "se_boot": float(np.std(draws, ddof=1)), "p_boot": p,89 }909192def holm(pvals: Sequence[float]) -> np.ndarray:93 """Holm (1979) step-down adjusted p-values (monotone, capped at 1). NaNs are passed through."""94 p = np.asarray(pvals, dtype=float)95 out = np.full_like(p, np.nan)96 ok = np.isfinite(p)97 pv = p[ok]98 m = len(pv)99 if m == 0:100 return out101 order = np.argsort(pv)102 adj = np.empty(m)103 running = 0.0104 for rank, i in enumerate(order):105 val = min(1.0, (m - rank) * pv[i])106 running = max(running, val)107 adj[i] = running108 out[ok] = adj109 return out110111112def newey_west_variance(d: np.ndarray, lag: int | None = None) -> tuple[float, int]:113 """Long-run variance of the mean of `d` with Bartlett kernel; default bandwidth floor(4 (n/100)^{2/9})."""114 d = np.asarray(d, dtype=float)115 n = len(d)116 if lag is None:117 lag = int(np.floor(4 * (n / 100.0) ** (2.0 / 9.0)))118 dc = d - d.mean()119 gamma0 = float(dc @ dc) / n120 lrv = gamma0121 for k in range(1, lag + 1):122 w = 1.0 - k / (lag + 1.0)123 lrv += 2.0 * w * float(dc[k:] @ dc[:-k]) / n124 return max(lrv, 1e-300), lag125126127def diebold_mariano(128 loss_a: np.ndarray, loss_b: np.ndarray, *, order: np.ndarray | None = None, lag: int | None = None, harvey: bool = True129) -> dict[str, float]:130 """DM test of H0: E[loss_b − loss_a] = 0 with HAC variance.131132 ASSUMPTION: the loss differential is a covariance-stationary time series ordered by date. Pass `order`133 (e.g. argsort of the sale dates) — the series is reordered before computing autocovariances. Harvey, Leybourne134 & Newbold (1997) small-sample correction applied by default (t_{n−1} reference distribution).135 """136 la, lb = np.asarray(loss_a, dtype=float), np.asarray(loss_b, dtype=float)137 if order is not None:138 la, lb = la[order], lb[order]139 d = lb - la140 n = len(d)141 lrv, h = newey_west_variance(d, lag)142 dm = d.mean() / np.sqrt(lrv / n)143 if harvey:144 corr = np.sqrt((n + 1 - 2 * (h + 1) + (h + 1) * h / n) / n)145 dm = dm * corr146 p = 2 * stats.t.sf(abs(dm), df=n - 1)147 else:148 p = 2 * stats.norm.sf(abs(dm))149 return {"dm_stat": float(dm), "dm_p": float(p), "dm_lag": float(h), "dm_mean_diff": float(d.mean())}150151152def compare_arms(153 preds_all: Any,154 truth: Any,155 pairs: Sequence[tuple[str | dict, str | dict]],156 metric: str | Statistic | Sequence[str | Statistic] = PRIMARY_METRIC,157 *,158 B: int = 2000,159 seed: int = 20261027,160 alpha: float = 0.05,161 dm: bool = True,162 holm_family: bool = True,163) -> pd.DataFrame:164 """Tidy table of paired contrasts. Δ = metric(B) − metric(A): for a loss, Δ < 0 means arm B is better.165166 One row per (pair, metric). Holm correction is applied within each metric across the supplied pairs167 (the pre-registered family) when `holm_family` is True. DM columns are NaN when no `date` is available.168 """169 metrics = list(metric) if isinstance(metric, (list, tuple)) else [metric]170 rows = []171 for spec_a, spec_b in pairs:172 a, b = paired_records(preds_all, truth, spec_a, spec_b)173 for m in metrics:174 res = block_bootstrap_contrast(a, b, m, B=B, seed=seed, alpha=alpha)175 row = {"arm_a": arm_label(spec_a), "arm_b": arm_label(spec_b)}176 row.update(res)177 row.update({"dm_stat": np.nan, "dm_p": np.nan, "dm_lag": np.nan})178 if dm and "date" in a.columns and a["date"].notna().all() and res["metric"] in {"mean_abs_log_error", "mae_cad", "rmse_log", "pinball_log", "crps_log"}:179 col = {"mean_abs_log_error": "abs_e", "mae_cad": "abs_err_cad", "rmse_log": "e", "pinball_log": "pinball_log", "crps_log": "crps_log"}[res["metric"]]180 la, lb = a[col].to_numpy(dtype=float), b[col].to_numpy(dtype=float)181 if col == "e":182 la, lb = la**2, lb**2 # DM on squared errors for RMSE183 ok = np.isfinite(la) & np.isfinite(lb)184 if ok.sum() >= 10:185 order = np.argsort(pd.to_datetime(a["date"]).to_numpy()[ok], kind="stable")186 row.update(diebold_mariano(la[ok], lb[ok], order=order))187 rows.append(row)188 out = pd.DataFrame(rows)189 if out.empty:190 return out191 out["p_holm"] = np.nan192 if holm_family:193 for m, idx in out.groupby("metric").groups.items():194 out.loc[idx, "p_holm"] = holm(out.loc[idx, "p_boot"].to_numpy())195 out["rel_diff"] = out["diff"] / out["value_a"]196 cols = ["arm_a", "arm_b", "metric", "n", "n_blocks", "B", "value_a", "value_b", "diff", "rel_diff", "ci_lo", "ci_hi",197 "se_boot", "p_boot", "p_holm", "dm_stat", "dm_p", "dm_lag"]198 return out[[c for c in cols if c in out.columns]]199