Did ideology move candidates off their polls? Opponents, party money and labels given after the fact

import os

for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(name, "1")  # before NumPy: one BLAS thread per chain

import json
import logging
import sys
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import nutpie
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import xarray as xr
from matplotlib.lines import Line2D
from scipy import stats

from pymc_challenges import data

# ArviZ leaves a lazy stub under "jax", and Matplotlib's first plot tests it, which would import all of
# JAX for nothing (see E89). Nothing here uses JAX, so the stub goes.
sys.modules.pop("jax", None)

RANDOM_SEED = 118
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
plt.rcParams.update({"axes.labelsize": 10, "xtick.labelsize": 9, "ytick.labelsize": 9})   # several panels per row
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="arviz")
pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 30)
pd.set_option("display.max_rows", 80)
T_START = time.time()
TIMES = {}

INK, MUTED, LIGHT = "#222222", "#8a8a8a", "#d9d9d9"
LEFT_C, CEN_C, D_C, R_C = "#7b3294", "#e08e0b", "#2166ac", "#c0392b"
EVE = {2018: "2018-11-05", 2022: "2022-11-07", 2026: "2026-11-02"}    # last day of information
ELEC = {2018: "2018-11-06", 2022: "2022-11-08", 2026: "2026-11-03"}
WINDOW = 21                                                           # days before election day
MIN_BUCKET = 3                                                        # fewer nominees than this are counted, not averaged

P = race_polls.copy()
P["end"] = pd.to_datetime(P.end)
P["days"] = (P.cycle.map(lambda c: pd.Timestamp(ELEC[c])) - P.end).dt.days
P["two"] = 100 * (2 * P.dem / (P.dem + P.rep) - 1)                 # two-party margin, D minus R, points
P["sponsor"] = np.where(P.partisan == "D", 1.0, np.where(P.partisan == "R", -1.0, 0.0))
n_rows = P.groupby("cycle").size()
# one row per pollster, race and end date: the likely-voter version where a poll reports several
P["pop_rank"] = P.population.map({"lv": 0, "rv": 1, "v": 2, "a": 3}).fillna(4)
P = P.sort_values(["cycle", "race", "pollster", "end", "pop_rank"]).drop_duplicates(["cycle", "race", "pollster", "end"])


def poll_average(window=WINDOW, nonpartisan=False):
    """Plain mean of the two-party margins of the polls that ended in the last `window` days."""
    q = P[P.days <= window]
    if nonpartisan:
        q = q[q.sponsor == 0]
    g = q.groupby(["cycle", "race"])
    return pd.DataFrame({"avg": g.two.mean(), "n": g.size(), "s": g.sponsor.mean()}).reset_index()


OUTCOMES = {"21": (21, False), "42": (42, False), "70": (70, False), "np": (21, True)}
tab = []
for c in (2018, 2022, 2026):
    q = P[(P.cycle == c) & (P.days <= WINDOW)]
    tab.append({"cycle": c, "poll rows": int(n_rows[c]), "polls after merging versions": int((P.cycle == c).sum()),
                "races with any poll": P[P.cycle == c].race.nunique(),
                f"polls in the last {WINDOW} days": len(q), "races": q.race.nunique(),
                "D-sponsored": int((q.sponsor > 0).sum()), "R-sponsored": int((q.sponsor < 0).sum()),
                "538 internal flag": int(q.internal.sum()),
                "races with one poll": int((q.groupby("race").size() == 1).sum())})
print(pd.DataFrame(tab).to_string(index=False))

LEFT = ["cpc", "jd", "our_revolution", "bnc", "dsa_member"]
CENTRE = ["ndc", "blue_dog", "blue_dog_site", "blue_dog_pac", "newdem_fund", "newdem_watch"]
LIST_NAME = {"cpc": "Congressional Progressive Caucus", "jd": "Justice Democrats", "our_revolution": "Our Revolution",
             "bnc": "Brand New Congress", "dsa_member": "DSA officeholders (Wikipedia's list)",
             "ndc": "New Democrat Coalition", "blue_dog": "Blue Dog Coalition (Wikipedia roster)",
             "blue_dog_site": "Blue Dog Coalition (its own member page)", "blue_dog_pac": "Blue Dog PAC endorsements",
             "newdem_fund": "NewDem Action Fund candidates", "newdem_watch": "NewDem Action Fund watch list"}
ROSTERS = ["cpc", "dsa_member", "ndc", "blue_dog", "blue_dog_site"]      # the others are endorsement lists
IE_GROUP = {"hill_committee": "hill", "aligned_superpac": "aligned", "aligned_affiliate": "aligned",
            "party_other": "other", "other": "other"}
IE_COLS = [f"{s}_{g}_{t}" for s in "dr" for g in ("hill", "aligned", "other") for t in ("sup", "opp")]
PLAIN = {"fund": ["lean", "d_inc", "d_served"], "opp": ["r_inc", "r_served", "r_hfc", "r_rmsp", "r_receipts"]}
MONEY = (["d_receipts", "red_to_blue", "d_direct"] + [c for c in IE_COLS if c[0] == "d"]
         + ["r_direct"] + [c for c in IE_COLS if c[0] == "r"])


def any_flag(F, cols, prefix="dx_"):
    return np.any([(F[prefix + c] == True).to_numpy() for c in cols], axis=0)  # noqa: E712  (a blank is not True)


def z_within_cycle(v, cycle):
    """z-score within a cycle; a column with no variation in a cycle becomes 0."""
    out = np.zeros(len(v))
    for c in np.unique(cycle):
        k = cycle == c
        sd = np.nanstd(v[k])
        out[k] = (v[k] - np.nanmean(v[k])) / sd if sd > 0 else 0.0
    return out


def money_table():
    """Dollars per race by side helped, group and support/oppose (filings received by the eve), and
    party committees' direct spending up to the pre-general report."""
    e = party_ie[(party_ie.version == "eve") & party_ie.nominee & party_ie.sup_opp.isin(["S", "O"])].copy()
    e["col"] = e.pro.str.lower() + "_" + e.group.map(IE_GROUP) + "_" + e.sup_opp.map({"S": "sup", "O": "opp"})
    ie_ = e.pivot_table(index=["cycle", "race"], columns="col", values="amt", aggfunc="sum")
    d = party_direct[party_direct.version == "pre_general"].copy()
    d["col"] = d.side.str.lower() + "_direct"
    di = d.pivot_table(index=["cycle", "race"], columns="col", values="amt", aggfunc="sum")
    return ie_.join(di, how="outer").fillna(0.0).reset_index()


def make_frame():
    """One row per race with both parties on the ballot and a poll in the last 70 days, all three cycles."""
    F = races[races.cycle.isin([2018, 2022, 2026]) & races.contested & ~races.split].copy()
    for key, (w, nonp) in OUTCOMES.items():
        F = F.merge(poll_average(w, nonp).rename(columns={"avg": f"avg_{key}", "n": f"n_{key}", "s": f"s_{key}"}),
                    on=["cycle", "race"], how="left")
    F = F[F.n_70.notna()]
    for side in "dr":
        I = ideo[ideo.side == side.upper()].drop(columns=["office", "state", "side"]).add_prefix(f"{side}x_")
        F = F.merge(I, left_on=["cycle", "race"], right_on=[f"{side}x_cycle", f"{side}x_race"], how="left")
    F = F[~(F.d_is_ind == True) & ~F.dx_multi_cands & ~F.rx_multi_cands]  # noqa: E712
    F = F.merge(money_table(), on=["cycle", "race"], how="left").reset_index(drop=True)
    cyc, new = F.cycle.to_numpy(), {}
    usd = {c: F[c].fillna(0.0).to_numpy() for c in IE_COLS + ["d_direct", "r_direct"]}
    F = F.drop(columns=list(usd))
    new["m"] = 100 * (2 * F.dem2 - 1)                  # certified two-party margin; blank for 2026
    for key in OUTCOMES:
        new[f"y_{key}"] = new["m"] - F[f"avg_{key}"]
    new["left"], new["centre"] = any_flag(F, LEFT), any_flag(F, CENTRE)
    new["bucket"] = np.select([new["left"] & new["centre"], new["left"], new["centre"]], ["both", "left", "centre"], "no list")
    # seat incumbents by who was in office on the eve; House members who won a special election count
    for s in "dr":
        new[f"{s}_inc"] = (F[f"{s}x_in_office_on_eve"] & (F[f"{s}x_incumbent"] | (F.office == "H"))).astype(float)
        new[f"{s}_served"] = (F[f"{s}x_served_before"] & (new[f"{s}_inc"] == 0)).astype(float)
        new[f"{s}_receipts"] = z_within_cycle(np.log10(F[f"{s}_rec_eve"].to_numpy(float) + 5e4), cyc)
    for c, v in usd.items():
        new[c + "_usd"] = v
        new[c] = z_within_cycle(np.log10(v + 1e4), cyc)
    new["r_hfc"] = (F.rx_hfc == True).astype(float)      # noqa: E712
    new["r_rmsp"] = (F.rx_rmsp == True).astype(float)    # noqa: E712
    new["red_to_blue"] = F.dx_red_to_blue.astype(float)  # blank in 2026: no dated list
    new["lean"] = F.lean / 10
    return pd.concat([F.drop(columns=[c for c in new if c in F]), pd.DataFrame(new, index=F.index)], axis=1)


ALL = make_frame()
MAIN = ALL[ALL.cycle.isin([2018, 2022]) & ALL.n_21.notna()].reset_index(drop=True)
print("races with a poll in the last 70 days:", ALL.groupby("cycle").size().to_dict())
print(f"races with a poll in the last {WINDOW} days and a certified result:", MAIN.groupby("cycle").size().to_dict())
print(MAIN.groupby(["cycle", "office"]).y_21.agg(races="size", mean="mean", sd="std", lowest="min", highest="max").round(1).to_string())
g = MAIN.groupby(pd.cut(MAIN.s_21, [-1.01, -0.01, 0.01, 1.01], labels=["net R-sponsored", "no sponsor mark", "net D-sponsored"]),
                 observed=True).y_21
print(g.agg(races="size", mean="mean", sd="std").round(1).to_string())
print(f"least-squares slope of the poll error on the sponsor share: {np.polyfit(MAIN.s_21, MAIN.y_21, 1)[0]:+.1f} points")

hist = pollerr[(pollerr.cycle < 2018) & ~pollerr.generic & pollerr.type_simple.isin(["House-G", "Sen-G"])].copy()
hist["err"] = -hist.err                                # the file's sign is poll minus result
hist["dev"] = hist.err - hist.groupby(["cycle", "type_simple"]).err.transform("mean")
gen = pollerr[(pollerr.cycle < 2018) & pollerr.generic].groupby("cycle").err.mean()
sen = pollerr[(pollerr.cycle < 2018) & (pollerr.type_simple == "Sen-G")].groupby("cycle").err.mean()
SD_B = float(np.sqrt(np.mean(np.r_[gen.to_numpy(), sen.to_numpy()] ** 2)))    # shared error of a cycle, as in E94
NBINS = [0, 1, 3, 6, 1000]
NLAB = ["1", "2-3", "4-6", "7+"]
hist_sd = hist.groupby(pd.cut(hist.n, NBINS, labels=NLAB), observed=True).dev.agg(["std", "size"])
main_dev = MAIN.y_21 - MAIN.groupby(["cycle", "office"]).y_21.transform("mean")
main_sd = main_dev.groupby(pd.cut(MAIN.n_21, NBINS, labels=NLAB), observed=True).agg(["std", "size"])
print(f"cycle-level shared error, 1998-2016: root mean square {SD_B:.1f} points over {len(gen) + len(sen)} cycle means")
print("sd of a race's poll error around its cycle and office mean, by number of polls in its average")
print(pd.DataFrame({"1998-2016 sd": hist_sd["std"].round(1), "races": hist_sd["size"],
                    "2018 and 2022 sd": main_sd["std"].round(1), "races ": main_sd["size"]}).to_string())

fig, axes = plt.subplots(1, 3, figsize=(13, 3.9))
ax = axes[0]
bins = np.arange(-27, 19, 3)
for c, col in ((2018, INK), (2022, MUTED)):
    ax.hist(MAIN.y_21[MAIN.cycle == c], bins=bins, histtype="step", lw=1.8, color=col, label=f"{c}: {int((MAIN.cycle == c).sum())} races")
ax.axvline(0, color=MUTED, lw=0.7)
ax.set_xlabel("result minus final poll average (D minus R, points)")
ax.set_ylabel("races")
ax.legend(fontsize=8, loc="upper left")
ax.set_title("The outcome: how far each race landed from its polls", fontsize=10)
ax = axes[1]
ax.scatter(MAIN.s_21 + rng.uniform(-0.03, 0.03, len(MAIN)), MAIN.y_21, s=14, color=INK, alpha=0.45, edgecolor="none")
xs = np.array([-1, 1])
slope = np.polyfit(MAIN.s_21, MAIN.y_21, 1)
ax.plot(xs, np.polyval(slope, xs), color=CEN_C, lw=2)
ax.axhline(0, color=MUTED, lw=0.7)
ax.set_xticks([-1, -0.5, 0, 0.5, 1], ["all R", "", "none", "", "all D"])
ax.set_xlabel("sponsor share of the polls in the average")
ax.set_title(f"Sponsored polls lean to their sponsor\n(least-squares slope {slope[0]:+.1f} points)", fontsize=10)
ax = axes[2]
x = np.arange(4)
ax.plot(x, hist_sd["std"], "o-", color=MUTED, label=f"1998-2016 ({len(hist)} races)")
ax.plot(x, main_sd["std"], "s-", color=INK, label=f"2018 and 2022 ({len(MAIN)} races)")
ax.set_xticks(x, NLAB)
ax.set_ylim(0, 10)
ax.set_xlabel("polls in the average")
ax.set_ylabel("sd of the poll error, points")
ax.legend(fontsize=8)
ax.set_title("More polls, smaller error", fontsize=10);

def naive_contrast(y, left, centre):
    """Mean poll error of centre-list nominees minus that of left-list nominees (those on both are left out)."""
    a, b = y[centre & ~left], y[left & ~centre]
    return a.mean() - b.mean() if len(a) and len(b) else np.nan


def bootstrap_contrast(y, left, centre, n=4000, seed=RANDOM_SEED):
    g = np.random.default_rng(seed)
    a, b = y[centre & ~left], y[left & ~centre]
    return (g.choice(a, (n, len(a))).mean(1) - g.choice(b, (n, len(b))).mean(1))


yM, LM, CM = MAIN.y_21.to_numpy(), MAIN.left.to_numpy(), MAIN.centre.to_numpy()
bm = MAIN.groupby("bucket").y_21.agg(races="size", mean="mean", sd="std").round(2).loc[["left", "centre", "both", "no list"]]
bm.loc[bm.races < MIN_BUCKET, ["mean", "sd"]] = np.nan
print(bm.to_string(na_rep="too few"))
rows = []
for lab, k in (("2018", MAIN.cycle == 2018), ("2022", MAIN.cycle == 2022), ("both cycles", MAIN.cycle > 0),
               ("both cycles, House", MAIN.office == "H"), ("both cycles, not seat incumbents", MAIN.d_inc == 0)):
    k = k.to_numpy()
    bs = bootstrap_contrast(yM[k], LM[k], CM[k])
    rows.append({"races": lab, "left": int((LM & ~CM & k).sum()), "centre": int((CM & ~LM & k).sum()),
                 "centre minus left": round(naive_contrast(yM[k], LM[k], CM[k]), 2),
                 "90% lower": round(np.quantile(bs, 0.05), 1), "90% upper": round(np.quantile(bs, 0.95), 1)})
NAIVE = pd.DataFrame(rows).set_index("races")
print(NAIVE.to_string())
NAIVE_BS = bootstrap_contrast(yM, LM, CM)

fig, ax = plt.subplots(figsize=(10, 3.6))
order = [("left", LEFT_C), ("centre", CEN_C), ("both", INK), ("no list", MUTED)]
for i, (b, col) in enumerate(order):
    s = MAIN[MAIN.bucket == b]
    for c, mk, dx in ((2018, "o", -0.14), (2022, "^", 0.14)):
        v = s.y_21[s.cycle == c]
        ax.scatter(i + dx + rng.uniform(-0.09, 0.09, len(v)), v, s=22, marker=mk, color=col, alpha=0.55, edgecolor="none")
    if len(s) >= MIN_BUCKET:
        ax.plot([i - 0.3, i + 0.3], [s.y_21.mean()] * 2, color=INK, lw=2)
        ax.text(i + 0.33, s.y_21.mean(), f"{s.y_21.mean():+.1f}", va="center", fontsize=9)
ax.axhline(0, color=MUTED, lw=0.7)
ax.set_xticks(range(4), [f"{b}\n{int((MAIN.bucket == b).sum())} races" for b, _ in order])
ax.set_ylabel("result minus final poll average\n(points; above 0 = beat the polls)")
ax.legend(handles=[Line2D([], [], marker="o", ls="", color=MUTED, label="2018"), Line2D([], [], marker="^", ls="", color=MUTED, label="2022")],
          fontsize=8, loc="lower right")
ax.set_title("Poll error of Democratic nominees by dated list, 2018 and 2022 (bars: bucket means)", fontsize=10);

def design(F, key="21", ref=None):
    """Arrays for the model from rows of the frame. `ref` (a fitted design) supplies the centring means; a
    covariate that is blank (no dated source in that cycle) is set to the mean, so it contributes nothing."""
    plain = PLAIN["fund"] + PLAIN["opp"]
    Xp, Xm = F[plain].to_numpy(float), F[MONEY].to_numpy(float)
    mp = np.nanmean(Xp, 0) if ref is None else ref["mean_p"]
    mm = np.nanmean(Xm, 0) if ref is None else ref["mean_m"]
    Xp, Xm = np.where(np.isnan(Xp), mp, Xp) - mp, np.where(np.isnan(Xm), mm, Xm) - mm
    cycles = sorted(F.cycle.unique())
    sc = (F.cycle.astype(str) + " " + F.state).to_numpy(dtype=object)
    scl = sorted(set(sc))
    return {"F": F, "y": F[f"y_{key}"].to_numpy(float), "n": F[f"n_{key}"].to_numpy(float), "s": F[f"s_{key}"].to_numpy(float),
            "L": F.left.to_numpy(float), "C": F.centre.to_numpy(float), "Xp": Xp, "Xm": Xm, "mean_p": mp, "mean_m": mm,
            "plain": plain, "cycles": cycles, "ci": F.cycle.map({c: i for i, c in enumerate(cycles)}).to_numpy(),
            "oi": (F.office == "S").to_numpy().astype(int), "scl": scl, "si": np.array([scl.index(v) for v in sc]),
            "race": (F.cycle.astype(str) + " " + F.race).to_numpy(dtype=object)}


def build_model(D, blocks=("fund", "opp", "party"), likelihood="t", scales="shares", masked=False, extra=None):
    """Poll error of a race: cycle and office shifts, sponsor lean, covariate blocks, ideology lists, a
    state-and-cycle effect and noise that shrinks with the number of polls.

    scales="direct" samples three standard deviations; "shares" samples a total and its split (section 4.2).
    masked=True writes the likelihood with 0/1 row weights so one compiled model serves many refits."""
    plain = [c for b in ("fund", "opp") if b in blocks for c in PLAIN[b]]
    jp = [D["plain"].index(c) for c in plain]
    coords = {"cycle": D["cycles"], "office": ["H", "S"], "state_cycle": D["scl"], "race": D["race"],
              "bucket": ["left", "centre"], "part": ["state", "race", "poll"]}
    if plain:
        coords["x"] = plain
    if "party" in blocks:
        coords["money"] = MONEY
    if extra is not None:
        coords["extra"] = list(extra.columns)
    ok = ~np.isnan(D["y"])
    with pm.Model(coords=coords) as m:
        L = pm.Data("L", D["L"], dims="race")
        C = pm.Data("C", D["C"], dims="race")
        y = pm.Data("y_obs", np.where(ok, D["y"], 0.0), dims="race")
        n = pm.Data("n_polls", np.where(ok, D["n"], 1.0), dims="race")
        s = pm.Data("sponsor", np.where(ok, D["s"], 0.0), dims="race")
        w = pm.Data("w", ok.astype(float), dims="race")
        b = pm.Normal("b", 0, SD_B, dims="cycle")                               # shared polling error
        d = pm.ZeroSumNormal("d", sigma=2.0, dims=("cycle", "office"))          # House against Senate
        delta = pm.Normal("delta", 2.0, 2.0)                                    # sponsor lean
        theta = pm.Normal("theta", 0, 3.0, dims="bucket")
        pm.Deterministic("contrast", theta[1] - theta[0])                       # centre minus left
        mu = b[D["ci"]] + d[D["ci"], D["oi"]] - delta * s + theta[0] * L + theta[1] * C
        if plain:
            beta = pm.Normal("beta", 0, 2.0, dims="x")
            mu = mu + pt.dot(D["Xp"][:, jp], beta)
        if "party" in blocks:
            s_money = pm.HalfNormal("s_money", 1.0)
            gamma = pm.Deterministic("gamma", s_money * pm.Normal("z_money", 0, 1, dims="money"), dims="money")
            mu = mu + pt.dot(D["Xm"], gamma)
        if extra is not None:
            xe = extra.to_numpy(float)
            mu = mu + pt.dot(xe - xe.mean(0), pm.Normal("beta_extra", 0, 2.0, dims="extra"))
        u = pm.Normal("u", 0, 1, dims="state_cycle")
        if scales == "direct":
            sigma = pm.HalfNormal("sigma", 4.0, dims="office")
            tau = pm.HalfNormal("tau", 6.0)
            mu = mu + pm.HalfNormal("s_state", 3.0) * u[D["si"]]
            sd = pt.sqrt(sigma[D["oi"]] ** 2 + tau**2 / n)
        else:
            S = pm.HalfNormal("S", 7.0, dims="office")                         # scale of a one-poll race
            share = pm.Dirichlet("share_of_variance", np.array([2.0, 2.0, 2.0]), dims="part")
            mu = mu + S[D["oi"]] * pt.sqrt(share[0]) * u[D["si"]]
            sd = S[D["oi"]] * pt.sqrt(share[1] + share[2] / n)
        dist = pm.Normal.dist(mu, sd) if likelihood == "normal" else pm.StudentT.dist(nu=pm.Gamma("nu", 2.0, 0.1), mu=mu, sigma=sd)
        if masked:
            pm.Potential("likelihood", pt.sum(w * pm.logp(dist, y)))
        elif likelihood == "normal":
            pm.Normal("y", mu, sd, observed=D["y"], dims="race")
        else:
            pm.StudentT("y", nu=m["nu"], mu=mu, sigma=sd, observed=D["y"], dims="race")
    return m


def diagnostics(idata):
    """Divergences, and the worst r_hat and ESS over variables, one variable at a time (see E94)."""
    rh, es = az.rhat(idata.posterior), az.ess(idata.posterior)
    return {"divergences": int(idata.sample_stats["diverging"].sum()),
            "max r_hat": round(max(float(rh[v].max()) for v in rh.data_vars), 3),
            "worst r_hat": max(rh.data_vars, key=lambda v: float(rh[v].max())),
            "min ESS": int(min(float(es[v].min()) for v in es.data_vars)),
            "lowest ESS": min(es.data_vars, key=lambda v: float(es[v].min())),
            "tuning": int(idata.posterior.attrs.get("tuning_steps", -1))}


def fit(D, **kw):
    t0 = time.time()
    with build_model(D, **kw):
        idata = pm.sample(random_seed=RANDOM_SEED, progressbar=False, target_accept=0.95)
    return idata, {"seconds": round(time.time() - t0, 1), **diagnostics(idata)}


def contrast_row(c):
    c = np.asarray(c).ravel()
    return {"centre minus left": round(float(np.median(c)), 2), "90% lower": round(float(np.quantile(c, 0.05)), 1),
            "90% upper": round(float(np.quantile(c, 0.95)), 1), "P(centre > left)": round(float((c > 0).mean()), 2)}


def ess_of(c):
    """Bulk ESS of flat draws stored chain after chain (4 chains)."""
    return int(az.ess(xr.DataArray(np.asarray(c).reshape(4, -1), dims=("chain", "draw"))))


def tail_probabilities(c, by=(2, 3)):
    c = np.asarray(c).ravel()
    return {f"P(centre ahead by {k}+ points)": round(float((c >= k).mean()), 3) for k in by} | {
        f"P(left ahead by {k}+ points)": round(float((c <= -k).mean()), 3) for k in by}


DM = design(MAIN)
print(f"{len(MAIN)} races, {len(DM['scl'])} state-and-cycle groups, "
      f"{len(DM['plain'])} fundamentals and opponent covariates, {len(MONEY)} party-support covariates")
sizes = pd.Series(DM["si"]).value_counts()
print(f"state-and-cycle groups holding a single race: {int((sizes == 1).sum())} of {len(sizes)}; "
      f"races with a single poll: {int((DM['n'] == 1).sum())} of {len(MAIN)}")
print(f"prior sd of a cycle's shared error: {SD_B:.1f} points (from 1998-2016)")

def robust_sd(a, axis=-1):
    q75, q25 = np.nanquantile(a, [0.75, 0.25], axis=axis)
    return (q75 - q25) / 1.349


def prior_spread(**kw):
    """Spread of simulated race errors with one poll and with four or more, and the prior of the contrast."""
    with build_model(DM, **kw):
        prior = pm.sample_prior_predictive(draws=2000, random_seed=RANDOM_SEED)
    yp = prior.prior_predictive["y"].stack(s=("chain", "draw")).transpose("s", ...).values
    yp = np.where(np.isfinite(yp), yp, np.nan)
    return (robust_sd(yp[:, DM["n"] == 1]), robust_sd(yp[:, DM["n"] >= 4]), prior.prior["contrast"].values.ravel(),
            np.array([naive_contrast(v, LM, CM) for v in yp[:500]]))


def show_spread(lab, one, many):
    print(f"{lab}: one poll {np.median(one):.1f} [{np.quantile(one, 0.05):.1f}, {np.quantile(one, 0.95):.1f}], "
          f"four or more polls {np.median(many):.1f} [{np.quantile(many, 0.05):.1f}, {np.quantile(many, 0.95):.1f}]")


pp_one, pp_many, pc, pn = prior_spread(likelihood="normal", scales="direct")
h_one, h_many = robust_sd(hist.dev[hist.n == 1].to_numpy()), robust_sd(hist.dev[hist.n >= 4].to_numpy())
print("spread of race poll errors across a data set (IQR / 1.349, points): median and 90% interval over prior draws")
show_spread("prior predictive", pp_one, pp_many)
print(f"1998-2016:        one poll {h_one:.1f}, four or more polls {h_many:.1f}")
print(f"prior of the contrast (centre minus left): sd {pc.std():.1f}, 90% interval [{np.quantile(pc, 0.05):.1f}, {np.quantile(pc, 0.95):.1f}]")
print(f"prior predictive of the difference in means: 90% interval [{np.nanquantile(pn, 0.05):.1f}, {np.nanquantile(pn, 0.95):.1f}]")

fig, axes = plt.subplots(1, 2, figsize=(10, 3.3))
ax = axes[0]
bins = np.linspace(0, 25, 51)
ax.hist(np.clip(pp_one, 0, 25), bins=bins, color=INK, alpha=0.55, density=True, label="one poll")
ax.hist(np.clip(pp_many, 0, 25), bins=bins, color=CEN_C, alpha=0.55, density=True, label="four or more polls")
ax.axvline(h_one, color=INK, lw=2, ls="--")
ax.axvline(h_many, color=CEN_C, lw=2, ls="--")
ax.set_xlabel("spread of race poll errors in a simulated data set, points")
ax.set_yticks([])
ax.legend(fontsize=8)
ax.set_title("Prior predictive spread (dashed: 1998-2016)", fontsize=10)
ax = axes[1]
ax.hist(pc, bins=60, color=MUTED, alpha=0.7, density=True)
ax.axvline(0, color=INK, lw=0.7)
ax.set_xlabel("centre minus left, points")
ax.set_yticks([])
ax.set_title("Prior of the ideology contrast", fontsize=10);

def draws_of(idata, n=2000):
    """Posterior draws as plain arrays, first axis the draw."""
    post = az.extract(idata, num_samples=n, random_seed=RANDOM_SEED)
    T = lambda v, *dims: post[v].transpose("sample", *dims).values  # noqa: E731
    out = {"b": T("b", "cycle"), "d": T("d", "cycle", "office"), "delta": T("delta"), "theta": T("theta", "bucket"),
           "S": T("S", "office"), "share": T("share_of_variance", "part")}
    for v, dim in (("beta", "x"), ("gamma", "money"), ("beta_extra", "extra")):
        if v in post:
            out[v] = T(v, dim)
    if "nu" in post:
        out["nu"] = T("nu")
    return out


def predict(dr, D, new_cycle=None, zero_ideology=False, seed=RANDOM_SEED):
    """Replicated poll errors (draws x races). State-and-cycle effects are always drawn afresh: keeping the
    fitted ones would let each lone race explain itself. new_cycle=sd draws the cycle's shared error and its
    House-Senate split from their priors, for a cycle the model has not seen."""
    g = np.random.default_rng(seed)
    k = len(dr["delta"])
    if new_cycle is not None:
        dh = g.normal(0, 2.0 / np.sqrt(2), k)[:, None]             # ZeroSumNormal(2) over two offices
        level = g.normal(0, new_cycle, k)[:, None] + np.where(D["oi"] == 1, -dh, dh)
    else:
        level = dr["b"][:, D["ci"]] + dr["d"][:, D["ci"], D["oi"]]
    mu = level - dr["delta"][:, None] * D["s"]
    if not zero_ideology:
        mu = mu + dr["theta"][:, [0]] * D["L"] + dr["theta"][:, [1]] * D["C"]
    if "beta" in dr:
        mu = mu + dr["beta"] @ D["Xp"].T
    if "gamma" in dr:
        mu = mu + dr["gamma"] @ D["Xm"].T
    S = dr["S"][:, D["oi"]]
    mu = mu + S * np.sqrt(dr["share"][:, [0]]) * g.standard_normal((k, len(D["scl"])))[:, D["si"]]
    eps = g.standard_t(dr["nu"][:, None], size=mu.shape) if "nu" in dr else g.standard_normal(mu.shape)
    return mu + S * np.sqrt(dr["share"][:, [1]] + dr["share"][:, [2]] / D["n"]) * eps


def check(rep, y):
    lo, hi = np.quantile(rep, [0.05, 0.95], 0)
    lo5, hi5 = np.quantile(rep, [0.25, 0.75], 0)
    out = {"90% interval covers": round(float(np.mean((y >= lo) & (y <= hi))), 3),
           "50% interval covers": round(float(np.mean((y >= lo5) & (y <= hi5))), 3)}
    for lab, f in (("largest |error|", lambda a: np.abs(a).max(-1)), ("sd", lambda a: a.std(-1)),
                   ("excess kurtosis", lambda a: stats.kurtosis(a, axis=-1))):
        r = f(rep)
        out[lab] = f"{f(y):.1f} against {np.median(r):.1f} [{np.quantile(r, 0.05):.1f}, {np.quantile(r, 0.95):.1f}], p = {np.mean(r >= f(y)):.3f}"
    return out


t0 = time.time()
FULL, diag_full = fit(DM, likelihood="t")
TIMES["full model"] = time.time() - t0
DR_N, DR = draws_of(fit_normal), draws_of(FULL)
rep_n, rep_t = predict(DR_N, DM), predict(DR, DM)
print(pd.DataFrame({"normal": check(rep_n, yM), "Student-t": check(rep_t, yM)}).to_string())
print(az.summary(FULL, var_names=["nu"], round_to=1).to_string())
print("normal   ", contrast_row(fit_normal.posterior["contrast"].values))
print("Student-t", contrast_row(FULL.posterior["contrast"].values))
NORMAL_CONTRAST = fit_normal.posterior["contrast"].values.ravel()
del fit_normal

fig, axes = plt.subplots(1, 3, figsize=(13, 3.4))
for ax, (lab, rep, col) in zip(axes[:2], (("normal", rep_n, MUTED), ("Student-t", rep_t, CEN_C))):
    pit = (rep < yM).mean(0)
    ax.hist(pit, bins=10, range=(0, 1), color=col, alpha=0.8, density=True)
    ax.axhline(1, color=INK, lw=0.7, ls=":")
    ax.set_xlabel("PIT of each race's poll error")
    ax.set_title(f"{lab} errors", fontsize=10)
axes[0].set_ylabel("density")
ax = axes[2]
bins = np.linspace(10, 50, 41)
ax.hist(np.clip(np.abs(rep_n).max(1), 10, 50), bins=bins, color=MUTED, alpha=0.7, density=True, label="normal")
ax.hist(np.clip(np.abs(rep_t).max(1), 10, 50), bins=bins, color=CEN_C, alpha=0.6, density=True, label="Student-t")
ax.axvline(np.abs(yM).max(), color=INK, lw=2)
ax.set_xlabel("largest poll error in a replicated data set, points")
ax.set_yticks([])
ax.legend(fontsize=8)
ax.set_title("The largest miss (line: observed)", fontsize=10);