The Overton Paradox, Updated: Technical Appendix

import sys

if "google.colab" in sys.modules:
    %pip install -q "pymc>=6" "arviz>=1.1" nutpie

import gzip
import json
import re
from os.path import basename, exists
from urllib.request import urlretrieve

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from scipy.optimize import brentq
from scipy.special import expit
from scipy.stats import norm

DATA_URL = "https://github.com/AllenDowney/OvertonParadox/releases/download/data-v1/"
LOCAL = "../release/data-v1/"  # staged by scripts/stage_release.sh, when running in a clone of the repo


def download(filename):
    """Download a file from the data release, unless it's already here."""
    if exists(LOCAL + filename):
        return LOCAL + filename
    if not exists(filename):
        urlretrieve(DATA_URL + filename, filename)
    return filename

manifest = json.load(open(download("manifest.json")))
raw = gzip.decompress(open(download("respondents.bin.gz"), "rb").read())


def decode(raw, columns):
    """The excerpt as a DataFrame; -1 in the integer columns means missing."""
    out = {}
    for col in columns:
        dtype = np.dtype(col["dtype"]).newbyteorder("<")
        a = np.frombuffer(raw, dtype=dtype, count=col["length"], offset=col["offset"])
        out[col["name"]] = np.where(a == -1, np.nan, a.astype(float)) if dtype.kind == "i" else a.astype(float)
    return pd.DataFrame(out)


gss = decode(raw, manifest["columns"])
gss["year"] = gss["year"].astype(int)
gss.shape

table = json.load(open(download("questions.json")))
meta = pd.DataFrame(table["questions"]).T
len(meta)

def polviews3(x):
    return np.select([x <= 3, x == 4, x >= 5], ["Liberal", "Moderate", "Conservative"], None)


gss["polviews3"] = pd.Series(polviews3(gss["polviews"]), index=gss.index).where(gss["polviews"].notna())
gss["decade"] = gss["year"] // 10 * 10
gss.groupby("decade")["polviews3"].value_counts(normalize=True).unstack().round(2)

candidates = [v for v in meta.index if meta.loc[v, "rank"] is not None and not pd.isna(meta.loc[v, "rank"])]
len(candidates)

NOT_A_POSITION = ["other", "depends", "depends on age", "neither", "wdnt belong", "wouldnt vote",
                  "r pays no income tax (vol.)"]
SCALE_ORDER = {
    "divlaw": [1, 3, 2],  # easier, stay as is, more difficult
    "courts": [1, 3, 2],  # too harshly, about right, not harshly enough
}


def response_differences(gss, varname):
    """Conservative minus liberal percentage choosing each answer, averaged over years."""
    data = gss[["year", "polviews3", "weight", varname]].dropna()
    data = data[data["polviews3"].isin(["Conservative", "Liberal"])]
    diffs, counts = [], []
    for year, group in data.groupby("year"):
        shares = group.pivot_table(index=varname, columns="polviews3", values="weight", aggfunc="sum", fill_value=0)
        if not {"Conservative", "Liberal"} <= set(shares.columns):
            continue
        shares = shares / shares.sum()
        diffs.append((shares["Conservative"] - shares["Liberal"]) * 100)
        counts.append(len(group))
    table = pd.concat(diffs, axis=1).fillna(0)
    return (table * counts).sum(axis=1) / sum(counts)


def best_cut(positions, order):
    """The contiguous split of an ordered scale that best separates the groups."""
    cum = positions.reindex(order).cumsum().iloc[:-1]
    k = cum.abs().values.argmax()
    low, high = list(order[: k + 1]), list(order[k + 1 :])
    return (low, cum.iloc[k]) if cum.iloc[k] > 0 else (high, -cum.iloc[k])


def code_question(gss, varname):
    labels = {float(code): label for code, label in meta.loc[varname, "value_labels"].items()}
    diff = response_differences(gss, varname)
    positions = diff.drop([v for v in diff.index if labels.get(v, "").lower() in NOT_A_POSITION])
    order = [float(v) for v in SCALE_ORDER.get(varname, sorted(positions.index))]
    return best_cut(positions, order)

response_differences(gss, "cappun").round(1)

book_codes = {v: [float(c) for c in codes] for v, codes in table["book_codes"].items()}
codes, separation = {}, {}
for v in candidates:
    side, size = code_question(gss, v)
    codes[v] = book_codes.get(v, sorted(side))
    separation[v] = size

saved = {v: [float(c) for c in meta.loc[v, "conservative"]] for v in candidates}
assert all(codes[v] == saved[v] for v in candidates)
assert max(abs(separation[v] - meta.loc[v, "separation"]) for v in candidates) < 1e-6

def make_coded(gss, codes):
    """1 for a conservative answer, 0 for another answer, NaN if not answered."""
    return pd.DataFrame({v: gss[v].isin(cs).astype(float).where(gss[v].notna()) for v, cs in codes.items()})


coded = make_coded(gss, codes)

w_mean1 = gss["weight"] / gss["weight"].mean()


def polviews_coefficient(v):
    data = pd.DataFrame(dict(y=coded[v], year=gss["year"], cohort=gss["cohort"], polviews=gss["polviews"],
                             w=w_mean1)).dropna()
    model = smf.glm("y ~ year + cohort + polviews", data=data, family=sm.families.Binomial(),
                    freq_weights=data["w"])
    return model.fit().params["polviews"]


param = pd.Series({v: polviews_coefficient(v) for v in candidates}).sort_values(ascending=False)
param.head(10).round(3)

assert np.allclose(param, meta.loc[param.index, "polviews_param"].astype(float), atol=1e-6)

MIN_WAVES = 16

answered_with_polviews = coded.notna() & gss["polviews"].notna().values[:, None]
waves = (answered_with_polviews.groupby(gss["year"]).sum() >= 100).sum()
assert (waves[candidates] == meta.loc[candidates, "waves"].astype(int)).all()
coverage_ok = waves >= MIN_WAVES
coverage_ok.sum()

REVERSAL_POINTS = 5
MIN_RESPONDENTS = 200
Z90 = norm.ppf(0.95)


def weighted_share(y, w):
    """Weighted mean and its standard error, using Kish's effective sample size."""
    p = np.average(y, weights=w)
    n_eff = w.sum() ** 2 / (w**2).sum()
    return p, np.sqrt(p * (1 - p) / n_eff)


def decade_gaps(v):
    df = pd.DataFrame(dict(y=coded[v], p=gss["polviews3"], w=gss["weight"], decade=gss["decade"])).dropna()
    df = df[df["p"].isin(["Liberal", "Conservative"])]
    rows = []
    for d, group in df.groupby("decade"):
        con, lib = group[group["p"] == "Conservative"], group[group["p"] == "Liberal"]
        if len(group) < MIN_RESPONDENTS or len(con) == 0 or len(lib) == 0:
            continue
        pc, sc = weighted_share(con["y"].values, con["w"].values)
        pl, sl = weighted_share(lib["y"].values, lib["w"].values)
        gap, se = (pc - pl) * 100, np.hypot(sc, sl) * 100
        rows.append(dict(decade=d, gap=gap, lo=gap - Z90 * se, hi=gap + Z90 * se))
    return pd.DataFrame(rows)


def reverses(gaps):
    con_ahead = (gaps["lo"] > 0) & (gaps["gap"] >= REVERSAL_POINTS)
    lib_ahead = (gaps["hi"] < 0) & (gaps["gap"] <= -REVERSAL_POINTS)
    return bool(con_ahead.any() and lib_ahead.any())


gaps = {v: decade_gaps(v) for v in candidates}
reversal = pd.Series({v: reverses(g) for v, g in gaps.items()})
reversal[reversal].index.tolist()

gaps["confed"].round(1)

def gap_string(g):
    return " ".join(f"{d}s:{x:+.0f}" for d, x in zip(g["decade"], g["gap"]))


assert all(gap_string(gaps[v]) == meta.loc[v, "gap_by_decade"] for v in candidates)

meta.loc[candidates].groupby("topic").size().loc[lambda s: s > 1]

CONFIDENCE = r"^con(press|labor|clerg|educ|bus|army|sci|tv|finan|fed|legis|judge|medic)$"
confidence = pd.Series({v: bool(re.match(CONFIDENCE, v)) for v in candidates})
confidence.sum()

def select(drop=()):
    """Candidates in rank order that pass coverage, stability, wording, and confidence, one per topic."""
    keep = [v for v in param.index
            if coverage_ok[v] and not reversal[v] and not confidence[v]
            and not (v.startswith("nat") and v.endswith("y")) and v not in drop]
    topics = meta.loc[keep, "topic"]
    return list(topics[~topics.duplicated()].index)


before_drift = select()
len(before_drift)

XS = np.linspace(-6, 6, 21)
PRIOR = norm.pdf(XS, 0, 4)
PRIOR /= PRIOR.sum()
MIN_ANSWERED = 5


def log_likelihoods(a, d, xs=XS):
    """Each question's log likelihood on the grid: shape (questions, respondents, grid).

    a: array of 1/0/NaN answers (respondents x questions); d: difficulties
    """
    P = expit(xs[None, :] - d[:, None])
    out = np.zeros((a.shape[1], a.shape[0], len(xs)))
    for j in range(a.shape[1]):
        col = a[:, j, None]
        out[j] = np.where(col == 1, np.log(P[j]), np.where(col == 0, np.log(1 - P[j]), 0))
    return out


def grid_posterior(a, d, prior=PRIOR, xs=XS):
    """Each respondent's posterior on the grid."""
    loglike = log_likelihoods(a, d, xs).sum(axis=0)
    posterior = np.exp(loglike - loglike.max(axis=1, keepdims=True)) * prior
    return posterior / posterior.sum(axis=1, keepdims=True)

def em_difficulties(coded, weights, tol=1e-4, max_iter=100):
    a = coded.values
    w = np.asarray(weights, dtype=float)
    asked = ~np.isnan(a)
    share = np.nansum(a * w[:, None], axis=0) / (asked * w[:, None]).sum(axis=0)
    d = -np.log(share / (1 - share))  # start: minus the logit of the prevalence
    for _ in range(max_iter):
        posterior = grid_posterior(a, d)
        new = d.copy()
        for j in range(a.shape[1]):
            mass = (w[asked[:, j], None] * posterior[asked[:, j]]).sum(axis=0)
            target = (w[asked[:, j]] * a[asked[:, j], j]).sum()
            new[j] = brentq(lambda dj: (mass * expit(XS - dj)).sum() - target, -10, 10)
        change = np.abs(new - d).max()
        d = new
        if change < tol:
            break
    return pd.Series(d, coded.columns)


def calibrate(names):
    """EM difficulties for a set of questions, from respondents who answered at least MIN_ANSWERED."""
    sub = coded[names]
    keep = (sub.notna().sum(axis=1) >= MIN_ANSWERED).values
    return em_difficulties(sub[keep], gss.loc[keep, "weight"])

DRIFT_THRESHOLD = 0.5  # logits


def drift(names, only=None):
    """Decade effects for each question, controlling for its leave-one-out score."""
    d = calibrate(names)
    a = coded[names].values
    ll = log_likelihoods(a, d.values)
    total = ll.sum(axis=0)
    answered = ~np.isnan(a)
    n_answered = answered.sum(axis=1)
    rows = []
    for j, v in enumerate(names):
        if only and v not in only:
            continue
        keep = answered[:, j] & (n_answered - 1 >= MIN_ANSWERED)
        loo = total[keep] - ll[j, keep]
        post = np.exp(loo - loo.max(axis=1, keepdims=True)) * PRIOR
        score = (post / post.sum(axis=1, keepdims=True)) @ XS
        y, w, decade = a[keep, j], gss["weight"].values[keep], gss["decade"].values[keep]
        counts = pd.Series(decade).value_counts()
        ok = np.isin(decade, counts.index[counts >= MIN_RESPONDENTS])
        decades = np.sort(np.unique(decade[ok]))
        X = np.column_stack([score[ok]] + [(decade[ok] == k).astype(float) for k in decades])
        res = sm.GLM(y[ok], X, family=sm.families.Binomial(), var_weights=w[ok] / w[ok].mean()).fit(cov_type="HC1")
        k = len(decades)
        C = np.zeros((k, k + 1))
        C[:, 1:] = np.eye(k) - 1 / k  # each decade relative to the mean of the decades
        effect = C @ res.params
        se = np.sqrt(np.diag(C @ res.cov_params() @ C.T))
        listed = ((np.abs(effect) - Z90 * se > 0) & (np.abs(effect) >= DRIFT_THRESHOLD)).any()
        rows.append(dict(variable=v, description=meta.loc[v, "description"], slope=res.params[0], listed=listed,
                         **{f"{dec}s": e for dec, e in zip(decades, effect)}))
    return pd.DataFrame(rows).set_index("variable")


drift_table = drift(before_drift)
listed = drift_table[drift_table["listed"]]
len(listed)

listed.drop(columns="listed").round(2)

drift(select(drop=["librac", "natspac"]), only=["colrac"]).drop(columns="listed").round(2)

DRIFT_EXCLUDED = ["librac", "natspac", "colrac", "spkrac"]
primary = select(drop=DRIFT_EXCLUDED)
len(primary)

assert sorted(primary) == sorted(table["presets"]["primary"])

meta.loc[primary, ["description", "first", "last", "waves"]]

asked = answered_with_polviews[primary].groupby(gss["year"]).sum() >= 100
by_decade = asked.groupby(asked.index // 10 * 10).any()
shared = by_decade.astype(int) @ by_decade.T.astype(int)
shared.index = shared.columns = [f"{d}s" for d in by_decade.index]
shared

asked.sum(axis=1).rename("questions asked").to_frame().T

d_primary = calibrate(primary)
assert np.allclose(d_primary, pd.Series(table["preset_settings"]["primary"]["difficulties"])[primary], atol=1e-9)


def percent_conservative(names, d):
    """Posterior mean percent conservative responses, for every respondent."""
    posterior = grid_posterior(coded[names].values, d[names].values)
    expected = expit(XS[None, :] - d[names].values[:, None]).mean(axis=0) * 100  # at each grid point
    return posterior @ expected


answered_primary = coded[primary].notna().sum(axis=1)
keep = (answered_primary >= MIN_ANSWERED) & gss["cohort"].notna() & gss["age"].notna()
df = gss.loc[keep, ["year", "age", "cohort", "polviews", "polviews3", "weight"]].copy()
df["y"] = percent_conservative(primary, d_primary)[keep.values] / 100  # the models work with proportions
len(df)

yearly = df.groupby("year").apply(lambda g: np.average(g["y"], weights=g["weight"]) * 100)
yearly.plot(xlabel="Survey year", ylabel="Percent conservative responses", style="o-");

fixture = json.load(open(download("primary_fixture.json")))
sample = pd.DataFrame(fixture["sample"])
mine = gss.loc[keep, ["year", "id"]].assign(y=df["y"].values).merge(sample, on=["year", "id"], suffixes=("", "_saved"))
assert len(mine) == len(sample) and np.allclose(mine["y"], mine["y_saved"], atol=1e-9)
saved_yearly = pd.Series(fixture["yearly_mean_y"]).rename(index=int) * 100
assert np.allclose(yearly, saved_yearly.loc[yearly.index], atol=1e-9)

def weighted_cells(df, keys):
    def summarize(g):
        w, y = g["weight"].values, g["y"].values
        mean = np.average(y, weights=w)
        return pd.Series(dict(p_weighted=mean, var_weighted=np.average((y - mean) ** 2, weights=w), n=len(g),
                              n_eff=w.sum() ** 2 / (w**2).sum()))

    return df.groupby(keys).apply(summarize).reset_index()

import pymc as pm


def rw2_trend_basis(n):
    """Eigenbasis of the centered RW2 prior, with the linear ramp split out."""
    L = np.zeros((n, n - 2))
    for k in range(2, n):
        L[k, : k - 1] = np.arange(k - 1, 0, -1)
    P = np.eye(n) - 1 / n
    lam, U = np.linalg.eigh(P @ L @ L.T @ P)
    keep = lam > 1e-9 * lam.max()
    M = U[:, keep] * np.sqrt(lam[keep])
    r = P @ np.arange(n, dtype=float)
    a = M.T @ r / (r @ r)
    return M - np.outer(r, a), r, a


def rw2(name, n, sigma, slope_init_scale, dims):
    """Centered RW2 with a free initial slope, in the trend form."""
    M_perp, r, a = rw2_trend_basis(n)
    w = pm.Normal(f"{name}_w", 0.0, 1.0, shape=n - 2)
    beta = pm.Flat(f"{name}_trend")
    v0 = beta - sigma * pm.math.dot(a, w)
    pm.Potential(f"{name}_slope_init_prior", pm.logp(pm.Normal.dist(0.0, slope_init_scale), v0))
    return pm.Deterministic(name, sigma * pm.math.dot(M_perp, w) + beta * r, dims=dims)

PERIOD_SETTINGS = dict(sigma_t_scale=0.125, slope_init_scale=0.2, sigma_short_scale=0.5)
COHORT_SETTINGS = dict(sigma_c_scale=0.125, sigma_t_scale=0.125, slope_init_scale=0.2, sigma_short_scale=0.2,
                       sigma_resid_scale=0.5)


def series_model(cells, index, grid, s=PERIOD_SETTINGS):
    idx = cells[index].map({v: i for i, v in enumerate(grid)}).values
    se = np.sqrt(np.maximum(cells["var_weighted"].values, 1e-4) / np.maximum(cells["n_eff"].values, 1.0))
    with pm.Model(coords={"grid": grid}) as m:
        alpha = pm.Normal("alpha", 0.0, 1.0)
        sigma_t = pm.Exponential("sigma_t", 1.0 / s["sigma_t_scale"])
        trend = rw2("trend", len(grid), sigma_t, s["slope_init_scale"], dims="grid")
        eta = pm.Deterministic("eta", alpha + trend, dims="grid")
        pm.Deterministic("mu", pm.math.sigmoid(eta), dims="grid")
        sigma_short = pm.HalfNormal("sigma_short", s["sigma_short_scale"])
        epsilon_z = pm.ZeroSumNormal("epsilon_z", sigma=1.0, dims="grid")
        epsilon = pm.Deterministic("epsilon", sigma_short * epsilon_z, dims="grid")
        pm.Normal("y_mean", mu=pm.math.sigmoid(eta[idx] + epsilon[idx]), sigma=se, observed=cells["p_weighted"].values)
    return m


def period_cohort_model(cells, years, cohorts, s=COHORT_SETTINGS):
    t = cells["year"].map({v: i for i, v in enumerate(years)}).values
    c = cells["cohort"].map({v: i for i, v in enumerate(cohorts)}).values
    with pm.Model(coords={"year": years, "cohort_grid": cohorts}) as m:
        alpha = pm.Normal("alpha", 0.0, 1.0)
        sigma_c = pm.Exponential("sigma_c", 1.0 / s["sigma_c_scale"])
        cohort = rw2("cohort", len(cohorts), sigma_c, s["slope_init_scale"], dims="cohort_grid")
        sigma_t = pm.Exponential("sigma_t", 1.0 / s["sigma_t_scale"])
        period = rw2("period", len(years), sigma_t, s["slope_init_scale"], dims="year")
        sigma_short = pm.Exponential("sigma_short", 1.0 / s["sigma_short_scale"])
        epsilon_z = pm.ZeroSumNormal("epsilon_z", sigma=1.0, dims="year")
        epsilon = pm.Deterministic("epsilon", sigma_short * epsilon_z, dims="year")
        p_obs = pm.math.sigmoid(alpha + cohort[c] + period[t] + epsilon[t])
        sigma_resid = pm.HalfNormal("sigma_resid", s["sigma_resid_scale"], dims="year")
        pm.Normal("y_mean", mu=p_obs, sigma=sigma_resid[t] / np.sqrt(cells["n_eff"].values),
                  observed=cells["p_weighted"].values)
    return m

try:
    import nutpie  # noqa: F401
    SAMPLER = "nutpie"
except ImportError:
    SAMPLER = "pymc"


def sample(m, target_accept, seed=17):
    with m:
        return pm.sample(tune=1000, draws=1000, chains=2, target_accept=target_accept, random_seed=seed,
                         nuts_sampler=SAMPLER, progressbar=False)


def draws_of(trace, name):
    values = trace.posterior[name].values
    return values.reshape(-1, *values.shape[2:])


def interval(x, axis=0):
    return np.mean(x, axis=axis), np.quantile(x, 0.05, axis=axis), np.quantile(x, 0.95, axis=axis)

def fit_series(sub, index, target_accept=0.99):
    grid = list(range(int(sub[index].min()), int(sub[index].max()) + 1))
    cells = weighted_cells(sub.assign(**{index: sub[index].astype(int)}), [index])
    trace = sample(series_model(cells, index, grid), target_accept)
    mu = draws_of(trace, "mu") * 100
    x = np.asarray(grid, float)
    observed = np.isin(x, cells[index].values)
    xs = x[observed] - x[observed].mean()
    slope = (mu[:, observed] - mu[:, observed].mean(axis=1, keepdims=True)) @ xs / (xs @ xs) * 10
    mean, lo, hi = interval(mu)
    curve = pd.DataFrame({index: grid, "mean": mean, "lo90": lo, "hi90": hi})
    return curve, dict(zip(["slope_per_decade", "slope_lo90", "slope_hi90"], map(float, interval(slope)))), trace

def fit_period_cohort(sub, target_accept=0.99):
    sub = sub.assign(cohort=sub["cohort"].astype(int))
    years = list(range(sub["year"].min(), sub["year"].max() + 1))
    cohorts = list(range(sub["cohort"].min(), sub["cohort"].max() + 1))
    cells = weighted_cells(sub, ["cohort", "year"])
    trace = sample(period_cohort_model(cells, years, cohorts), target_accept)
    alpha = draws_of(trace, "alpha")[:, None]
    cohort, period = draws_of(trace, "cohort"), draws_of(trace, "period")

    survey_years = sorted(sub["year"].unique())
    ti = {y: years.index(y) for y in survey_years}
    comp = sub.groupby(["year", "cohort"])["weight"].sum().unstack(fill_value=0).reindex(columns=cohorts, fill_value=0)
    comp = comp.div(comp.sum(axis=1), axis=0).loc[survey_years].values

    def level(mix, t):
        """Mean percent with cohort mix `mix` and the period at index t."""
        return (expit(alpha + cohort + period[:, t, None]) * mix).sum(axis=1) * 100

    y0, y1 = survey_years[0], survey_years[-1]
    p0, p1 = level(comp[0], ti[y0]), level(comp[-1], ti[y1])
    fix_cohort, fix_period = level(comp[0], ti[y1]), level(comp[-1], ti[y0])
    total = p1 - p0
    replacement = 0.5 * ((fix_period - p0) + (total - (fix_cohort - p0)))
    summary = {}
    for label, x in [("total_change", total), ("generational_replacement", replacement),
                     ("changing_minds", total - replacement)]:
        summary.update(dict(zip([label, f"{label}_lo90", f"{label}_hi90"], map(float, interval(x)))))

    pooled = sub.groupby("cohort")["weight"].sum().reindex(cohorts, fill_value=0).values
    std_period = np.stack([level(pooled / pooled.sum(), j) for j in range(len(years))], axis=1)
    x = np.array(survey_years, float) - np.mean(survey_years)
    sp = std_period[:, [ti[y] for y in survey_years]]
    slope = (sp - sp.mean(axis=1, keepdims=True)) @ x / (x @ x) * 10
    summary.update(zip(["period_slope_per_decade", "period_slope_lo90", "period_slope_hi90"],
                       map(float, interval(slope))))

    mean, lo, hi = interval(std_period)
    period_df = pd.DataFrame(dict(year=years, mean=mean, lo90=lo, hi90=hi))
    mean, lo, hi = interval(expit(alpha + cohort + period.mean(axis=1, keepdims=True)) * 100)
    cohort_df = pd.DataFrame(dict(cohort=cohorts, mean=mean, lo90=lo, hi90=hi))
    return period_df, cohort_df, summary, trace

saved = pd.read_csv(download("apc_summary.csv")).set_index("run")

age_fits = {}
for start in range(1970, 2030, 10):
    window = df[df["year"].between(start, start + 9)]
    age_fits[start] = fit_series(window.assign(age=window["age"].astype(int)), "age")

for start, (curve, _, _) in age_fits.items():
    plt.plot(curve["age"], curve["mean"], label=f"{start}s")
plt.xlabel("Age")
plt.ylabel("Percent conservative responses")
plt.legend();

windows = {1970: "1972-1979", 1980: "1980-1989", 1990: "1990-1999", 2000: "2000-2009", 2010: "2010-2019",
           2020: "2020-2024"}
pd.DataFrame({f"{start}s": dict(here=s["slope_per_decade"],
                                production=saved.loc[f"age_primary_window-{windows[start]}", "slope_per_decade"])
              for start, (_, s, _) in age_fits.items()}).T.round(2)

group_fits = {g: fit_series(df[df["polviews3"] == g], "year") for g in ["Liberal", "Moderate", "Conservative"]}
for g, color in zip(group_fits, ["C0", "C4", "C3"]):
    curve = group_fits[g][0]
    plt.plot(curve["year"], curve["mean"], color=color, label=g)
    plt.fill_between(curve["year"], curve["lo90"], curve["hi90"], color=color, alpha=0.2)
plt.xlabel("Survey year")
plt.ylabel("Percent conservative responses")
plt.legend();

period_df, cohort_df, pc_summary, pc_trace = fit_period_cohort(df)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for ax, comp, x, label in [(axes[0], cohort_df, "cohort", "Year of birth"), (axes[1], period_df, "year", "Survey year")]:
    ax.plot(comp[x], comp["mean"])
    ax.fill_between(comp[x], comp["lo90"], comp["hi90"], alpha=0.2)
    ax.set_xlabel(label)
axes[0].set_ylabel("Percent conservative responses")
axes[0].set_title("Cohort component")
axes[1].set_title("Period component (within cohorts)");

keys = ["period_slope_per_decade", "period_slope_lo90", "period_slope_hi90", "total_change",
        "generational_replacement", "changing_minds"]
pd.DataFrame(dict(here=[pc_summary[k] for k in keys],
                  production=saved.loc["period_cohort_primary", keys].astype(float))).round(2)

groups = saved[(saved["model"] == "period_cohort") & (saved["set"] == "primary") & saved["cohorts"].notna()]
decades = groups[groups["cohorts"].str.match(r"^\d{3}0-\d{3}9$")]
decades.set_index("cohorts")[["period_slope_per_decade", "period_slope_lo90", "period_slope_hi90",
                               "respondents"]].round(2)

XS_FINE = np.linspace(-6, 6, 121)
GROUPS = ["Liberal", "Moderate", "Conservative"]
LAST_YEAR = 2024


def latent_normals(answers, weights, d, groups, tol=1e-6, max_iter=500):
    """Mean and SD of a normal latent distribution in each group, difficulties fixed."""
    w = np.asarray(weights, dtype=float)
    labels, codes = np.unique(np.asarray(groups), return_inverse=True)
    loglike = log_likelihoods(answers, d, XS_FINE).sum(axis=0)
    like = np.exp(loglike - loglike.max(axis=1, keepdims=True))
    mu, sd = np.zeros(len(labels)), np.ones(len(labels))
    for _ in range(max_iter):
        posterior = like * norm.pdf(XS_FINE[None, :], mu[:, None], sd[:, None])[codes]
        posterior /= posterior.sum(axis=1, keepdims=True)
        total = np.stack([w[codes == k] @ posterior[codes == k] for k in range(len(labels))])
        m1 = total @ XS_FINE / total.sum(axis=1)
        m2 = total @ XS_FINE**2 / total.sum(axis=1)
        new_sd = np.sqrt(m2 - m1**2)
        change = max(np.abs(m1 - mu).max(), np.abs(new_sd - sd).max())
        mu, sd = m1, new_sd
        if change < tol:
            break
    return pd.DataFrame(dict(mu=mu, sd=sd), index=labels)


answers = coded.loc[keep, primary].values
dfp = df.reset_index(drop=True)
in_groups = dfp["polviews3"].isin(GROUPS).values
by_year = latent_normals(answers, dfp["weight"], d_primary.values, dfp["year"].astype(str))
by_group = latent_normals(answers[in_groups], dfp.loc[in_groups, "weight"], d_primary.values,
                          dfp.loc[in_groups, "year"].astype(str) + "|" + dfp.loc[in_groups, "polviews3"])

REF_YEAR = 1974
ref = by_group.loc[f"{REF_YEAR}|Liberal", "mu"]
rows = []
for y in sorted(dfp["year"].unique()):
    dist = by_year.loc[str(y)]
    row = dict(year=y, ref_liberal=100 * norm.cdf(ref, dist["mu"], dist["sd"]))
    for g in GROUPS:
        if f"{y}|{g}" in by_group.index:
            row[g.lower()] = 100 * norm.cdf(by_group.loc[f"{y}|{g}", "mu"], dist["mu"], dist["sd"])
    rows.append(row)
position = pd.DataFrame(rows).set_index("year")

position[["liberal", "moderate", "conservative"]].plot(alpha=0.6, color=["C0", "C4", "C3"])
plt.plot(position.index, position["ref_liberal"], "ko-", ms=3, label=f"{REF_YEAR} liberal")
plt.axhline(50, color="0.6", lw=1)
plt.ylabel("Percentile of the year's distribution")
plt.xlabel("Survey year")
plt.legend();

def headline_numbers(sub, sub_answers):
    in_g = sub["polviews3"].isin(GROUPS).values
    year = latent_normals(sub_answers, sub["weight"], d_primary.values, sub["year"].astype(str)).loc[str(LAST_YEAR)]
    grp = latent_normals(sub_answers[in_g], sub.loc[in_g, "weight"], d_primary.values,
                         sub.loc[in_g, "year"].astype(str) + "|" + sub.loc[in_g, "polviews3"])
    out = dict(ref_liberal=100 * norm.cdf(grp.loc[f"{REF_YEAR}|Liberal", "mu"], year["mu"], year["sd"]))
    for g in GROUPS:
        out[g.lower()] = 100 * norm.cdf(grp.loc[f"{LAST_YEAR}|{g}", "mu"], year["mu"], year["sd"])
    return out


two = dfp["year"].isin([REF_YEAR, LAST_YEAR]).values
sub, sub_answers = dfp[two].reset_index(drop=True), answers[two]
point = headline_numbers(sub, sub_answers)
rng = np.random.default_rng(17)
draws = []
for _ in range(200):
    rows = np.concatenate([rng.choice(idx, len(idx)) for idx in sub.groupby("year").indices.values()])
    draws.append(headline_numbers(sub.iloc[rows].reset_index(drop=True), sub_answers[rows]))
draws = pd.DataFrame(draws)
pd.DataFrame(dict(percentile=point, lo90=draws.quantile(0.05), hi90=draws.quantile(0.95))).round(1)

saved_headline = pd.read_csv(download("position_headline.csv")).set_index("set").loc["primary"]
assert all(abs(point[k] - saved_headline[k]) < 1e-6 for k in point)

rows = {}
for name in ["primary", "with_confidence", "p48", "social", "econA"]:
    run = saved.loc[f"period_cohort_{name}"]
    rows[name] = run[["n_questions", "respondents", "period_slope_per_decade", "period_slope_lo90",
                      "period_slope_hi90", "generational_replacement", "changing_minds"]]
pd.DataFrame(rows).T.round(2)

diag = saved[saved["set"].isin(["primary", "with_confidence", "social", "econA"])]
diag[["model", "set", "by", "window", "cohorts", "target_accept", "rhat_max", "ess_bulk_min", "divergences"]]

pd.Series(dict(rhat_max=diag["rhat_max"].max(), ess_bulk_min=diag["ess_bulk_min"].min(),
           divergences=diag["divergences"].sum(), fits_with_divergences=(diag["divergences"] > 0).sum())).round(3)

import arviz as az

rows = {}
for label, trace in [("period-cohort", pc_trace)] + [(f"period {g}", t) for g, (_, _, t) in group_fits.items()] + \
        [(f"age {s}s", t) for s, (_, _, t) in age_fits.items()]:
    summary = az.summary(trace, var_names=["alpha", "sigma_t", "sigma_short"], kind="diagnostics")
    rows[label] = dict(rhat_max=summary["r_hat"].max(), ess_bulk_min=summary["ess_bulk"].min(),
                       divergences=int(trace.sample_stats["diverging"].values.sum()))
pd.DataFrame(rows).T