The winner's curse: shrinkage and cumulative impact

%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import norm

from experiment_utils.shrinkage import (
    cumulative_impact,
    empirical_bayes_shrinkage,
    estimate_guardrail_rho,
    fit_t_prior_with_estimated_mean,
    nss_adjusted_cumulative_impact,
    nss_adjusted_cumulative_impact_mvn,
    t_prior_shrinkage,
)

pd.set_option("display.notebook_repr_html", False)
pd.set_option("display.float_format", lambda x: f"{x:.3f}")
plt.rcParams.update({"figure.dpi": 110, "font.size": 11})

rng = np.random.default_rng(20260815)

# --- A world of experiments (relative lifts, in %) ---
n_exp = 200
n_guards = 5
alpha_ship = 0.10
z_crit = norm.ppf(1.0 - alpha_ship / 2.0)   # two-sided hard bar (~1.645)

tau = 0.8          # spread of true primary effects
mu_true = 0.3      # portfolio helps a little on average
true_effect = rng.normal(mu_true, tau, n_exp)

se = rng.uniform(0.8, 2.5, n_exp)
observed = true_effect + rng.normal(0, se)

# Guardrails: true effects correlated with the primary (heterogeneous rho).
rhos = np.linspace(0.05, 0.45, n_guards)
true_guards = np.column_stack([
    rho * true_effect + tau * np.sqrt(1.0 - rho**2) * rng.normal(0, 1, n_exp)
    for rho in rhos
])
se_guards = rng.uniform(0.8, 2.5, size=(n_exp, n_guards))
obs_guards = true_guards + rng.normal(0, se_guards)

exp = pd.DataFrame({"true_effect": true_effect, "observed": observed, "se": se})
exp["z"] = exp["observed"] / exp["se"]
print(exp.head())
print(f"\nGuardrail correlations with primary (true): {np.round(rhos, 2)}")
print(f"Portfolio truth: mean true effect = {true_effect.mean():.2f}%")
print(f"Mean observed effect (all 200)     = {observed.mean():.2f}%")

primary_win = exp["z"].values > z_crit
z_guards = obs_guards / se_guards
guard_harm = z_guards < -z_crit                    # any significant harm?
blocked_by_guards = primary_win & guard_harm.any(axis=1)
shipped = primary_win & ~guard_harm.any(axis=1)

exp["primary_win"] = primary_win
exp["shipped"] = shipped
winners = exp[exp["shipped"]]

funnel = pd.DataFrame({
    "count": [n_exp, int(primary_win.sum()), int(blocked_by_guards.sum()),
              int(shipped.sum())],
}, index=["experiments", "primary hard wins", "blocked by guardrails", "shipped"])
print(funnel, "\n")

print(f"Shipped winners' average OBSERVED lift : {winners['observed'].mean():.2f}%")
print(f"Shipped winners' average TRUE lift     : {winners['true_effect'].mean():.2f}%")
inflation = winners["observed"].mean() / winners["true_effect"].mean()
print(f"\nReported lift is {inflation:.1f}x the truth among the shipped set.")

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(exp.loc[~exp["primary_win"], "true_effect"],
           exp.loc[~exp["primary_win"], "observed"],
           s=18, color="0.75", label="not a primary win")
ax.scatter(exp.loc[blocked_by_guards, "true_effect"],
           exp.loc[blocked_by_guards, "observed"],
           s=36, marker="x", color="#e67e22", label="blocked by guardrails")
ax.scatter(winners["true_effect"], winners["observed"],
           s=28, color="#c0392b", label="shipped (primary + guards)")
lims = [-4, 8]
ax.plot(lims, lims, "k--", lw=1, label="observed = truth")
ax.axhline(winners["observed"].mean(), color="#c0392b", lw=1.2, ls=":",
           label="shipped mean observed")
ax.axhline(winners["true_effect"].mean(), color="#2471a3", lw=1.2, ls=":",
           label="shipped mean truth")
ax.set_xlim(lims); ax.set_ylim(lims)
ax.set_xlabel("True primary lift (%)"); ax.set_ylabel("Observed primary lift (%)")
ax.set_title("Winner's curse after the shipping rule")
ax.legend(fontsize=9, loc="upper left")
plt.tight_layout()
plt.show()

t_prior = fit_t_prior_with_estimated_mean(
    exp["observed"].values, exp["se"].values, df=4.0
)

print(
    f"Fitted t-prior:  mu = {t_prior['prior_mean']:.2f}%  "
    f"[{t_prior['prior_mean_ci_lower']:.2f}, {t_prior['prior_mean_ci_upper']:.2f}]  "
    f"scale = {t_prior['scale']:.2f}  df = {t_prior['df']:.0f}"
)
print(f"(true values were mu = {mu_true:.2f}%   tau = {tau:.2f}%)\n")

shrunk = t_prior_shrinkage(
    exp["observed"].values,
    exp["se"].values,
    scale=t_prior["scale"],
    df=t_prior["df"],
    prior_mean=t_prior["prior_mean"],
)
exp["shrunk"] = shrunk["shrunk"]
exp["B"] = shrunk["shrinkage_factor"]

eb = empirical_bayes_shrinkage(
    exp["observed"].values, exp["se"].values, prior_mean=t_prior["prior_mean"]
)
exp["shrunk_eb"] = eb["shrunk"]

def rmse(a, b):
    return np.sqrt(np.mean((a - b) ** 2))

print("Accuracy vs the known truth (all 200 experiments):")
print(f"  RMSE of raw observed : {rmse(exp['observed'], exp['true_effect']):.2f}")
print(f"  RMSE of t-prior      : {rmse(exp['shrunk'],   exp['true_effect']):.2f}")
print(f"  RMSE of normal EB    : {rmse(exp['shrunk_eb'], exp['true_effect']):.2f}")

w = exp[exp["shipped"]]
summary = pd.DataFrame({
    "mean lift (%)": [w["observed"].mean(), w["shrunk"].mean(),
                      w["shrunk_eb"].mean(), w["true_effect"].mean()],
    "RMSE vs truth": [rmse(w["observed"], w["true_effect"]),
                      rmse(w["shrunk"],   w["true_effect"]),
                      rmse(w["shrunk_eb"], w["true_effect"]), 0.0],
}, index=["observed (reported)", "t-prior shrunk", "normal EB", "true"])
print(summary, "\n")

retained = w["shrunk"].mean() / w["observed"].mean()
print(f"t-prior shrinkage retains {retained:.0%} of the shipped reported mean lift.")
print(f"Average t-prior shrinkage factor among shipped: B = {w['B'].mean():.2f}")

fig, ax = plt.subplots(figsize=(7, 5))
order = w.sort_values("observed").reset_index(drop=True)
xpos = np.arange(len(order))
ax.vlines(xpos, order["shrunk"], order["observed"], color="0.8", lw=1)
ax.scatter(xpos, order["observed"], s=22, color="#c0392b", label="reported")
ax.scatter(xpos, order["shrunk"], s=22, color="#2471a3", label="t-prior shrunk")
ax.axhline(order["true_effect"].mean(), color="k", ls="--", lw=1,
           label="true mean of shipped")
ax.set_xlabel("shipped experiments (sorted by reported lift)")
ax.set_ylabel("lift (%)")
ax.set_title("Shrinkage pulls reported lifts toward the archive average")
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()

naive_stack = float(w["observed"].sum())
true_stack = float(w["true_effect"].sum())

cum_t = cumulative_impact(
    exp["observed"].values,
    exp["se"].values,
    shipped=exp["shipped"].values,
    prior=t_prior,
    aggregation="sum",
)

cum_eb = cumulative_impact(
    exp["observed"].values,
    exp["se"].values,
    shipped=exp["shipped"].values,
    prior_mean=t_prior["prior_mean"],
    aggregation="sum",
)

# NSS pick-one: companion with largest |rho| (known here from the DGP).
best_k = int(np.argmax(np.abs(rhos)))
cum_nss = nss_adjusted_cumulative_impact(
    exp["observed"].values,
    exp["se"].values,
    obs_guards[:, best_k],
    se_guards[:, best_k],
    shipped=exp["shipped"].values,
    rho=float(rhos[best_k]),
    prior_sd_primary=float(np.sqrt(eb["tau2"])),
    prior_sd_guard=float(np.sqrt(eb["tau2"])),
    prior_mean_primary=t_prior["prior_mean"],
    aggregation="sum",
)


# NSS MVN: joint shrink with all K companions (magnitude only).
cum_mvn = nss_adjusted_cumulative_impact_mvn(
    exp["observed"].values,
    exp["se"].values,
    obs_guards,
    se_guards,
    shipped=exp["shipped"].values,
    rho_primary=rhos,
    prior_sd_primary=float(np.sqrt(eb["tau2"])),
    prior_sd_guard=float(np.sqrt(eb["tau2"])),
    prior_mean_primary=t_prior["prior_mean"],
    rho_guardrails="factor",
    aggregation="sum",
)

stack = pd.DataFrame({
    "cumulative impact (pp)": [
        naive_stack, cum_t["cumulative"], cum_eb["cumulative"],
        cum_nss["cumulative"], cum_mvn["cumulative"], true_stack,
    ],
    "% of naive": [
        100,
        100 * cum_t["cumulative"] / naive_stack,
        100 * cum_eb["cumulative"] / naive_stack,
        100 * cum_nss["cumulative"] / naive_stack,
        100 * cum_mvn["cumulative"] / naive_stack,
        100 * true_stack / naive_stack,
    ],
}, index=[
    "naive (reported)",
    "t-prior (primary-only)",
    "normal EB (primary-only)",
    f"NSS pick-one (rho={rhos[best_k]:.2f})",
    "NSS MVN (all 5 guards)",
    "true additive",
])
print(stack)
print(
    f"\nSame ship mask for every stack: n_shipped = {cum_t['n_shipped']}  "
    f"(primary wins {int(primary_win.sum())}, "
    f"blocked by guards {int(blocked_by_guards.sum())})"
)
print(
    f"t-prior 95% CI: [{cum_t['ci_lower']:.1f}, {cum_t['ci_upper']:.1f}]  "
    f"pick-one: [{cum_nss['ci_lower']:.1f}, {cum_nss['ci_upper']:.1f}]  "
    f"MVN: [{cum_mvn['ci_lower']:.1f}, {cum_mvn['ci_upper']:.1f}]"
)

def one_world(rng, delta):
    """One portfolio under the primary + guardrail ship gate."""
    te = rng.normal(mu_true, tau, n_exp)
    s = rng.uniform(0.8, 2.5, n_exp)
    x = te + rng.normal(0, s)

    tg = np.column_stack([
        rho * te + tau * np.sqrt(1.0 - rho**2) * rng.normal(0, 1, n_exp)
        for rho in rhos
    ])
    sg = rng.uniform(0.8, 2.5, size=(n_exp, n_guards))
    xg = tg + rng.normal(0, sg)

    primary_ok = (x / s) > z_crit
    guard_ok = ~((xg / sg) < -z_crit).any(axis=1)
    shipped = primary_ok & guard_ok

    if shipped.sum() < 1:
        return None

    mu_hat = float(np.average(x, weights=1.0 / s**2))
    cum = cumulative_impact(
        x, s, shipped=shipped, prior_mean=mu_hat, aggregation="sum"
    )

    a = te[shipped]
    L = float(a.sum())
    share = 1 - (np.sum(a**2) / L**2) if L != 0 else 0.0
    holdout_sub = L * (1 - delta * share)

    return {
        "naive": float(x[shipped].sum()),
        "shrunk": float(cum["cumulative"]),
        "oracle": L,
        "holdout_add": L,
        "holdout_sub": holdout_sub,
        "n_primary": int(primary_ok.sum()),
        "n_blocked": int((primary_ok & ~guard_ok).sum()),
        "n_shipped": int(shipped.sum()),
    }

n_mc, delta = 2000, 0.30
rows = [one_world(np.random.default_rng(s), delta) for s in range(n_mc)]
sim = pd.DataFrame([r for r in rows if r is not None])

print("Scale-rule funnel (means across MC worlds)")
print(sim[["n_primary", "n_blocked", "n_shipped"]].mean().round(2), "\n")

def bias_row(est, truth):
    d = sim[est] - sim[truth]
    return [sim[est].mean(), d.mean(), d.abs().mean()]

print("ADDITIVE world  (holdout = sum of true parts)")
add = pd.DataFrame(
    [bias_row("naive", "holdout_add"),
     bias_row("shrunk", "holdout_add"),
     bias_row("oracle", "holdout_add")],
    columns=["mean stack (pp)", "bias vs holdout", "mean |error|"],
    index=["naive", "cumulative_impact (EB)", "oracle"])
print(add, "\n")

print(f"SUB-ADDITIVE world  (delta = {delta}: joint < sum of parts)")
sub = pd.DataFrame(
    [bias_row("naive", "holdout_sub"),
     bias_row("shrunk", "holdout_sub"),
     bias_row("oracle", "holdout_sub")],
    columns=["mean stack (pp)", "bias vs holdout", "mean |error|"],
    index=["naive", "cumulative_impact (EB)", "oracle"])
print(sub)
print(f"\nJoint holdout is on average "
      f"{100 * sim['holdout_sub'].mean() / sim['holdout_add'].mean():.0f}% "
      f"of the additive sum.")

labels = ["naive", "shrunk", "oracle", "holdout"]
means = [sim["naive"].mean(), sim["shrunk"].mean(), sim["oracle"].mean(),
         sim["holdout_sub"].mean()]
colors = ["#c0392b", "#2471a3", "#7f8c8d", "#27ae60"]

fig, ax = plt.subplots(figsize=(7, 4.5))
bars = ax.bar(labels, means, color=colors)
ax.axhline(sim["holdout_sub"].mean(), color="#27ae60", ls="--", lw=1)
for b, m in zip(bars, means):
    ax.text(b.get_x() + b.get_width() / 2, m + 0.3, f"{m:.1f}",
            ha="center", fontsize=10)
ax.set_ylabel("cumulative impact (pp)")
ax.set_title("Sub-additive world: shrinkage helps, but the stack still\n"
             "overstates the holdout when effects do not fully add")
plt.tight_layout()
plt.show()