Ragged information: polls, fundamentals and money in the US midterms

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
os.environ.setdefault("JAX_PLATFORMS", "cpu")  # nothing here uses a GPU

import gzip
import hashlib
import json
import logging
import sys
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
from IPython.display import HTML
from matplotlib import animation
from matplotlib.collections import LineCollection, PolyCollection
from matplotlib.lines import Line2D
from matplotlib.ticker import MaxNLocator
from scipy.stats import gaussian_kde

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 = 94
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*sample_stats.*")
pd.set_option("display.width", 160)
pd.set_option("display.max_columns", 20)
T_START = time.time()

D_C, R_C, INK, MUTED, LIGHT, GOLD = "#2166ac", "#c0392b", "#222222", "#8a8a8a", "#d9d9d9", "#e08e0b"
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"}
END_30G = {2018: "2018-11-26", 2022: "2022-11-28", 2026: "2026-11-23"}  # post-general coverage ends
PRES = {1978: "D", 1982: "R", 1986: "R", 1990: "R", 1994: "D", 1998: "D", 2002: "R", 2006: "R",
        2010: "D", 2014: "D", 2018: "R", 2022: "D", 2026: "R"}         # president's party at each midterm
TRAIN = {2018: [2016], 2022: [2016, 2018, 2020], 2026: [2016, 2018, 2020, 2022]}
TIMES = {}

races = data.load("e94_races")
polls = data.load("e94_polls")
reports = data.load("e94_reports")
ie = data.load("e94_ie")
rematch = data.load("e94_rematch")
pollerr = data.load("e94_pollerr")
build_log = json.load(open(data.path("e94_build_log")))
for c in ("cvg_start", "cvg_end", "receipt"):
    reports[c] = pd.to_datetime(reports[c])
for c in ("date", "first_receipt"):
    ie[c] = pd.to_datetime(ie[c])
polls["end"] = pd.to_datetime(polls["end"])

NAT = {int(k): v for k, v in build_log["national_house_dem2"].items()}
NAT.update({int(k): v for k, v in build_log["national_house_dem2_medsl"].items() if int(k) < 2016})
MARGIN = {y: 100 * (2 * v - 1) for y, v in NAT.items()}          # national House margin, D minus R
MARGIN_2024 = 100 * (2 * build_log["national_house_2024_wikipedia"]["dem2"] - 1)


def money_signal(df):
    """log((D receipts + $50k) / (R receipts + $50k)): eve totals where reports exist, else cycle totals.

    A candidate with an FEC ID and nothing reported counts as $0. A 2026 nominee with no FEC ID found is
    missing (NaN), not $0: the model imputes the signal (see prepare)."""
    eve = df.cycle.isin([2018, 2022, 2026]).to_numpy()
    d = np.where(eve, df.d_rec_eve, df.d_rec_final).astype(float)
    r = np.where(eve, df.r_rec_eve, df.r_rec_final).astype(float)
    miss = (df.cycle == 2026).to_numpy() & (np.isnan(d) | np.isnan(r))
    out = np.clip(np.log((np.nan_to_num(d) + 5e4) / (np.nan_to_num(r) + 5e4)), -5, 5)
    return np.where(miss, np.nan, out)


R = races.copy()
R["inc"] = np.where(R.d_inc, 1.0, np.where(R.r_inc, -1.0, 0.0))
R["money"] = money_signal(R)
R["m"] = 100 * (2 * R.dem2 - 1)                                   # D minus R margin, points
R["leanx"] = R.lean / 10
R["modelled"] = R.seat & R.contested & ~R.split
print(R.groupby(["cycle", "office"]).agg(races=("race", "size"), contested=("contested", "sum"),
                                         modelled=("modelled", "sum")).T.to_string())

fig, axes = plt.subplots(2, 3, figsize=(13, 6.6), sharex=True, height_ratios=[2.2, 1])
cols = {"independent": INK, "partisan only": GOLD, "none": LIGHT}
for j, (c, t) in enumerate(rag.items()):
    ax = axes[0, j]
    for off, mk, sz in (("H", "o", 16), ("S", "*", 70)):
        for k, col in cols.items():
            s = t[(t.office == off) & (t.kind == k)]
            y = np.log1p(s.n) + (rng.uniform(-0.12, 0.12, len(s)) if k == "none" else 0)
            ax.scatter(s.lean, y, s=sz, marker=mk, color=col, alpha=0.8 if k != "none" else 0.6,
                       edgecolor="none" if off == "H" else "white", linewidth=0.4)
    ax.set_yticks(np.log1p([0, 1, 3, 10, 30, 100]), ["0", "1", "3", "10", "30", "100"])
    ax.set_title(f"{c}: {int((t[t.office == 'H'].n > 0).sum())} of {int((t.office == 'H').sum())} "
                 f"contested districts polled", fontsize=10)
    ax.axvspan(-10, 10, color=LIGHT, alpha=0.25, zorder=-1)
    ax.set_ylabel("race polls" if j == 0 else "")
    ax2 = axes[1, j]
    h = t[t.office == "H"]
    bins = np.arange(-60, 61, 5)
    cut = pd.cut(h.lean.clip(-59.9, 59.9), bins)
    share = h.groupby(cut, observed=False).apply(lambda s: (s.n > 0).mean() if len(s) else np.nan)
    part = h.groupby(cut, observed=False).apply(lambda s: (s.kind == "partisan only").mean() if len(s) else np.nan)
    mids = bins[:-1] + 2.5
    ax2.bar(mids, share.to_numpy(), width=4.4, color=INK, alpha=0.8, label="any poll")
    ax2.bar(mids, part.to_numpy(), width=4.4, color=GOLD, label="partisan polls only")
    ax2.set_ylim(0, 1)
    ax2.set_ylabel("share of districts" if j == 0 else "")
axes[1, 0].legend(fontsize=8, loc="upper left")
handles = [Line2D([], [], marker="o", ls="", color=cols[k], label=k) for k in cols]
handles += [Line2D([], [], marker="*", ls="", color=INK, ms=10, label="Senate race")]
axes[0, 2].legend(handles=handles, fontsize=8, loc="upper right")
fig.supxlabel("partisan lean of the race (D minus R, points); grey band: |lean| < 10", fontsize=10)
fig.suptitle("Ragged information: race polls by partisan lean (House: dots, Senate: stars)", fontsize=11);

def midterm_prior(target):
    """Previous House margin plus the mean swing against the president's party in earlier midterms."""
    sw = {y: (MARGIN[y] - MARGIN[y - 2]) * (1 if PRES[y] == "R" else -1)
          for y in PRES if y < target and y in MARGIN and y - 2 in MARGIN}
    prev = MARGIN_2024 if target == 2026 else MARGIN[target - 2]
    s = np.array(list(sw.values()))
    sign = 1 if PRES[target] == "R" else -1
    return {"prev": prev, "swings": sw, "mu": prev + sign * s.mean(), "sd": s.std(ddof=1),
            "mean_swing": s.mean(), "n": len(s)}


def poll_error_scales(target):
    """Shared and race-level polling-error scales from FiveThirtyEight's raw polls of earlier cycles."""
    pe = pollerr[pollerr.cycle < target]
    gen = pe[pe.generic].groupby("cycle").err.mean()
    sen = pe[~pe.generic & (pe.type_simple == "Sen-G")]
    sen_cyc = sen.groupby("cycle").err.mean()
    shared = float(np.sqrt(np.mean(np.r_[gen.to_numpy(), sen_cyc.to_numpy()] ** 2)))
    race = float((sen.err - sen.cycle.map(sen_cyc)).std())
    return {"shared": shared, "race": race, "n_cycles": len(gen) + len(sen_cyc)}


PRIOR = {c: midterm_prior(c) for c in (2018, 2022, 2026)}
PERR = {c: poll_error_scales(c) for c in (2018, 2022, 2026)}
for c in PRIOR:
    p, e = PRIOR[c], PERR[c]
    print(f"{c}: previous margin {p['prev']:+.1f}, mean swing against the president's party "
          f"{p['mean_swing']:.1f} (sd {p['sd']:.1f}, {p['n']} midterms) -> N ~ Normal({p['mu']:+.1f}, {p['sd']:.1f}); "
          f"poll error sd: shared {e['shared']:.1f}, race {e['race']:.1f}")
print("actual national margin: 2018 %+.1f, 2022 %+.1f" % (MARGIN[2018], MARGIN[2022]))
v = np.array(list(PRIOR[2026]["swings"].values()))
print(f"swings 1978-2022 range from {v.min():+.1f} to {v.max():+.1f}")

def prepare(target, mode="full"):
    """Training races (earlier cycles), target races, national and race polls for one target cycle."""
    train = R[R.cycle.isin(TRAIN[target]) & R.modelled & R.lean.notna()]
    tgt = R[(R.cycle == target) & R.seat].copy()
    tm = tgt[tgt.modelled].reset_index(drop=True)
    P = polls[polls.cycle == target].copy()
    P["week"] = ((pd.Timestamp(ELEC[target]) - P.end).dt.days // 7).astype(int)
    P = P[(P.week >= 0) & (P.week < 22)]
    P["two"] = 100 * (2 * P.dem / (P.dem + P.rep) - 1)
    nn = P.n.fillna(600).clip(200, 5000) * (P.dem + P.rep) / 100
    q = P.dem / (P.dem + P.rep)
    P["se"] = 200 * np.sqrt(q * (1 - q) / nn)
    if mode == "fund":
        P = P.iloc[0:0]
    gp = P[P.kind == ("generic_avg" if target == 2026 else "generic")].copy()
    if mode == "agg" and target < 2026:   # a stand-in aggregator: a 14-day mean of the polls, once a week
        rows = [{"week": w, "two": gp.two[(gp.week >= w) & (gp.week <= w + 1)].mean(), "pollster": "rolling mean"}
                for w in range(22) if ((gp.week >= w) & (gp.week <= w + 1)).sum() >= 3]
        gp = pd.DataFrame(rows)
    if target == 2026 or mode == "agg":
        gp["se"] = 1.0                    # an assumed sd for an average of several polls
    rp = P[P.kind.isin(["house", "senate"]) & P.race.isin(tm.race) & (P.week < 10)].copy()
    # A missing money signal gets a prior from the same cycle's other races: a regression on lean,
    # incumbency and office (all known on the eve), with its residual sd.
    miss = np.flatnonzero(tm.money.isna().to_numpy())
    Z = np.c_[np.ones(len(tm)), tm.leanx, tm.inc, (tm.office == "S").astype(float)]
    ok = ~tm.money.isna().to_numpy()
    coef, *_ = np.linalg.lstsq(Z[ok], tm.money.to_numpy()[ok], rcond=None)
    miss_sd = float(np.std(tm.money.to_numpy()[ok] - Z[ok] @ coef, ddof=Z.shape[1]))
    return {"target": target, "mode": mode, "train": train, "tgt": tgt, "tm": tm, "gp": gp, "rp": rp,
            "miss": miss, "miss_mu": np.clip(Z[miss] @ coef, -5, 5), "miss_sd": miss_sd,
            "pollsters": sorted(set(gp.pollster) | set(rp.pollster)), "polled": sorted(set(rp.race)),
            "states": sorted(set(R.state)), "prior": PRIOR[target], "perr": PERR[target]}


def build_model(d, rotate=True):
    train, tm, gp, rp = d["train"], d["tm"], d["gp"], d["rp"]
    cycles, states, pollsters, polled = TRAIN[d["target"]], d["states"], d["pollsters"], d["polled"]
    s_idx = {s: i for i, s in enumerate(states)}
    isS_tr = (train.office == "S").to_numpy().astype(int)
    isS = (tm.office == "S").to_numpy().astype(int)
    X_tr = np.c_[train.leanx, train.inc, train.money]
    X = np.c_[tm.leanx, tm.inc, np.nan_to_num(tm.money.to_numpy())]
    c_tr = train.cycle.map({c: i for i, c in enumerate(cycles)}).to_numpy()
    st_tr, st = train.state.map(s_idx).to_numpy(), tm.state.map(s_idx).to_numpy()
    natm_tr = np.array([MARGIN[c] for c in cycles])
    mu_N, sd_N, sd_b = d["prior"]["mu"], d["prior"]["sd"], d["perr"]["shared"]
    mu_b = d.get("bias_mean", 0.0)   # the shared error's prior mean: 0 except in section 7.1
    coords = {"office": ["H", "S"], "x": ["lean", "inc", "money"], "cycle": cycles, "state": states,
              "week": np.arange(22)}
    if pollsters:
        coords.update(pollster=pollsters)
    if polled:
        coords.update(polled=polled)
    miss = d["miss"]
    if len(miss):
        coords.update(missing_race=tm.race.to_numpy()[miss])
    with pm.Model(coords=coords) as m:
        if len(miss):   # unknown money signals, as parameters with the regression prior from prepare()
            mm = pm.Normal("money_missing", d["miss_mu"], d["miss_sd"], dims="missing_race")
            X = pt.set_subtensor(pt.as_tensor_variable(X)[miss, 2], mm)
        # fundamentals, learned on earlier cycles where the national margin is known
        beta = pm.Normal("beta", [5.0, 3.0, 2.0], [3.0, 3.0, 2.0], dims=("office", "x"))
        sig = pm.HalfNormal("sig", 10.0, dims="office")
        s_state = pm.HalfNormal("s_state", 4.0)
        gamma = pm.Normal("gamma", 0, 3, dims="office")
        s_int = pm.HalfNormal("s_int", 2.0)
        z_int = pm.Normal("z_int", 0, 1, dims=("cycle", "office"))
        z_st = pm.Normal("z_st", 0, 1, dims=("cycle", "state"))
        mu_tr = ((natm_tr[:, None] + gamma[None, :] + s_int * z_int)[c_tr, isS_tr]
                 + (beta[isS_tr] * X_tr).sum(-1) + s_state * z_st[c_tr, st_tr])
        pm.Normal("m_train", mu_tr, sig[isS_tr], observed=train.m.to_numpy())
        # the target cycle: national margin N and the shared polling error, seen by the polls as N + bias
        if rotate:
            zr = pm.Normal("z_rot", 0, 1, shape=2)
            h = np.hypot(sd_N, sd_b)
            N = pm.Deterministic("N", mu_N + sd_N * (sd_N / h * zr[0] - sd_b / h * zr[1]))
            bias = pm.Deterministic("bias", mu_b + sd_b * (sd_b / h * zr[0] + sd_N / h * zr[1]))
        else:
            N = pm.Deterministic("N", mu_N + sd_N * pm.Normal("z_N", 0, 1))
            bias = pm.Deterministic("bias", mu_b + sd_b * pm.Normal("z_bias", 0, 1))
        s_rw = pm.HalfNormal("s_rw", 1.0)
        z_rw = pm.Normal("z_rw", 0, 1, dims="week")
        walk = N + s_rw * pt.concatenate([pt.zeros(1), pt.cumsum(z_rw[1:])])  # week 0 = election week
        if pollsters:
            house = pm.ZeroSumNormal("house", sigma=pm.HalfNormal("s_house", 2.0), dims="pollster")
            pidx = {p: i for i, p in enumerate(pollsters)}
        z_int_t = pm.Normal("z_int_t", 0, 1, dims="office")
        icpt_t = N + gamma + s_int * z_int_t
        z_st_t = pm.Normal("z_st_t", 0, 1, dims="state")
        if len(gp):
            s_x = pm.HalfNormal("s_x_gen", 2.0)
            mu_g = walk[gp.week.to_numpy()] + bias + house[gp.pollster.map(pidx).to_numpy()]
            pm.Normal("gen", mu_g, pt.sqrt(gp.se.to_numpy() ** 2 + s_x**2), observed=gp.two.to_numpy())
        if polled:
            pol_i = np.array([tm.index[tm.race == r][0] for r in polled])
            e_pol = pm.Normal("e_pol", 0, 1, dims="polled")
            m_pol = pm.Deterministic(
                "m_pol", icpt_t[isS[pol_i]] + (beta[isS[pol_i]] * X[pol_i]).sum(-1)
                + s_state * z_st_t[st[pol_i]] + sig[isS[pol_i]] * e_pol, dims="polled")
            b_race = pm.Normal("b_race", 0, d["perr"]["race"], dims="polled")
            d_part = pm.Normal("d_part", 2.0, 2.0)
            k = rp.race.map({r: i for i, r in enumerate(polled)}).to_numpy()
            spons = np.where(rp.partisan.fillna("") == "D", 1.0, np.where(rp.partisan.fillna("") == "R", -1.0, 0.0))
            s_xr = pm.HalfNormal("s_x_race", 3.0)
            mu_r = (m_pol[k] + (walk[rp.week.to_numpy()] - N) + bias + b_race[k] + d_part * spons
                    + house[rp.pollster.map(pidx).to_numpy()])
            pm.Normal("racepoll", mu_r, pt.sqrt(rp.se.to_numpy() ** 2 + s_xr**2), observed=rp.two.to_numpy())
    return m


def diagnostics(idata):
    """Divergences, and the worst r_hat and ESS over variables (one variable at a time: see AUTHORING)."""
    rh, es = az.rhat(idata.posterior), az.ess(idata.posterior)
    worst = max(rh.data_vars, key=lambda v: float(rh[v].max()))
    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": worst,
            "min ESS": int(min(float(es[v].min()) for v in es.data_vars)),
            "tuning": int(idata.posterior.attrs.get("tuning_steps", -1))}


def simulate(d, idata, n=2000, seed=RANDOM_SEED):
    """Draws of every modelled race's margin: polled races from the sampler, the rest simulated."""
    tm = d["tm"]
    post = az.extract(idata, num_samples=n, random_seed=seed)
    g = np.random.default_rng(seed)
    isS = (tm.office == "S").to_numpy().astype(int)
    X = np.broadcast_to(np.c_[tm.leanx, tm.inc, np.nan_to_num(tm.money.to_numpy())], (n, len(tm), 3)).copy()
    st = tm.state.map({s: i for i, s in enumerate(d["states"])}).to_numpy()
    T = lambda v, *dims: post[v].transpose("sample", *dims).values  # noqa: E731
    if len(d["miss"]):
        X[:, d["miss"], 2] = T("money_missing", "missing_race")
    icpt = post["N"].values[:, None] + T("gamma", "office") + post["s_int"].values[:, None] * T("z_int_t", "office")
    mu = (icpt[:, isS] + (T("beta", "office", "x")[:, isS, :] * X).sum(-1)
          + (T("z_st_t", "state") * post["s_state"].values[:, None])[:, st])
    sim = mu + T("sig", "office")[:, isS] * g.standard_normal(mu.shape)
    if d["polled"]:
        pol_i = np.array([tm.index[tm.race == r][0] for r in d["polled"]])
        sim[:, pol_i] = T("m_pol", "polled")
    walk = post["N"].values[:, None] + post["s_rw"].values[:, None] * np.c_[
        np.zeros(n), np.cumsum(T("z_rw", "week")[:, 1:], axis=1)]
    return {"sim": sim, "N": post["N"].values, "bias": post["bias"].values, "walk": walk,
            "beta": T("beta", "office", "x"), "sig": T("sig", "office")}


def seats(d, sim):
    """D seats per draw: modelled races won plus seats with only one side on the ballot."""
    tgt, tm = d["tgt"], d["tm"]
    out = {}
    for off in ("H", "S"):
        sel = (tm.office == off).to_numpy()
        other = tgt[(tgt.office == off) & ~tgt.modelled]
        if d["target"] < 2026:   # one side on the ballot, or a jungle ballot called by lean
            fixed = int(np.sum(np.where(other.split, other.lean > 0, other.winner == "D")))
        else:
            fixed = int((other.d_name.fillna("") != "").sum())
        out[off] = (sim[:, sel] > 0).sum(1) + fixed
    return out


def fit_cycle(target, mode="full", rotate=True, keep=False, bias_mean=0.0):
    d = {**prepare(target, mode), "bias_mean": bias_mean}
    t0 = time.time()
    with build_model(d, rotate):
        idata = pm.sample(random_seed=RANDOM_SEED, progressbar=False, target_accept=0.9)
    out = {"d": d, "seconds": time.time() - t0, **diagnostics(idata), **simulate(d, idata)}
    out["seats"] = seats(d, out["sim"])
    if keep:
        out["idata"] = idata
    return out


d18 = prepare(2018)
print(f"2018: {len(d18['train'])} training races (2016), {len(d18['tm'])} modelled target races, "
      f"{len(d18['gp'])} generic polls, {len(d18['rp'])} race polls in {len(d18['polled'])} races, "
      f"{len(d18['pollsters'])} pollsters")
d26 = prepare(2026)
print(f"2026: {d26['gp'].pollster.nunique()} aggregators, {len(d26['gp'])} weekly snapshots in the last 22 weeks; "
      f"{len(d26['rp'])} race polls in {len(d26['polled'])} races")

def race_scores(sim, y):
    """Scores of predictive draws sim (draws x races) against outcomes y (D margin, points)."""
    S = sim.shape[0]
    pD = (sim > 0).mean(0)
    win = y > 0
    p_true = np.where(win, ((sim > 0).sum(0) + 0.5) / (S + 1), ((sim <= 0).sum(0) + 0.5) / (S + 1))
    out = {f"cover {int(100 * p)}%": float(np.mean((y >= np.quantile(sim, (1 - p) / 2, 0))
                                                     & (y <= np.quantile(sim, (1 + p) / 2, 0))))
           for p in (0.5, 0.8, 0.9)}
    out.update({"Brier": float(np.mean((pD - win) ** 2)), "log score": float(np.mean(np.log(p_true))),
                "MAE": float(np.mean(np.abs(np.median(sim, 0) - y))), "wrong calls": int(np.sum((pD > 0.5) != win))})
    return out, (sim < y).mean(0), pD


LABEL = {"full": "full model", "agg": "aggregate only", "fund": "fundamentals only", "polls": "polls only"}
score_rows, PIT, PD = [], {}, {}
for (c, mode), f in FITS.items():
    tm = f["d"]["tm"]
    y = tm.m.to_numpy()
    sc, PIT[c, mode], PD[c, mode] = race_scores(f["sim"], y)
    act = {off: int((f["d"]["tgt"][f["d"]["tgt"].office == off].winner == "D").sum()) for off in ("H", "S")}
    row = {"cycle": c, "model": LABEL[mode]}
    for off in ("H", "S"):
        q = np.quantile(f["seats"][off], [0.05, 0.5, 0.95])
        row[f"{off} seats"] = f"{q[1]:.0f} [{q[0]:.0f}, {q[2]:.0f}]"
        row[f"{off} actual"] = act[off]
    row["P(D House)"] = round(float((f["seats"]["H"] >= 218).mean()), 2)
    qn = np.quantile(f["N"], [0.05, 0.5, 0.95])
    row["national"] = f"{qn[1]:+.1f} [{qn[0]:+.1f}, {qn[2]:+.1f}]"
    row["actual N"] = round(MARGIN[c], 1)
    row.update({k: round(v, 3) if isinstance(v, float) else v for k, v in sc.items()})
    score_rows.append(row)
scores = pd.DataFrame(score_rows).sort_values(["cycle", "model"], key=lambda s: s.map(
    {"full model": 0, "aggregate only": 1, "polls only": 2, "fundamentals only": 3}) if s.name == "model" else s)
print(scores.drop(columns=["cover 50%", "cover 80%", "cover 90%", "Brier", "log score", "MAE", "wrong calls"]).to_string(index=False))
print()
print(scores[["cycle", "model", "cover 50%", "cover 80%", "cover 90%", "Brier", "log score", "MAE", "wrong calls"]].to_string(index=False))