The path we were on: years of life Europe lost to COVID, and whether it came back

import os

for v_ in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ[v_] = "1"  # nutpie runs chains as threads: one BLAS pool per chain would fight

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
from IPython.display import Markdown, display
from matplotlib.colors import TwoSlopeNorm

from pymc_challenges import data
from pymc_challenges import public as pub

RANDOM_SEED = 111
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="arviz")
warnings.filterwarnings("ignore", message="There are no gridspecs with layoutgrids")
sys.modules.pop("jax", None)  # arviz leaves a lazy jax stub that makes matplotlib import all of JAX
pd.set_option("display.width", 160)
T_START = time.time()
TIMES = {}
INK, MUTED, LIGHT, HI, LO = pub.INK, pub.MUTED, pub.LIGHT, pub.HI, pub.LO
print(f"PyMC {pm.__version__}, ArviZ {az.__version__}, nutpie {nutpie.__version__}")

EDGES = [0, 1] + list(range(5, 95, 5))  # 0, 1-4, 5-9, ..., 85-89, 90+: 20 groups
LABELS = ["0", "1-4"] + [f"{a}-{a + 4}" for a in range(5, 90, 5)] + ["90+"]
WIDTH = np.diff(EDGES + [np.inf])
MID = np.array([0.5, 3.0] + [x + 2.5 for x in range(5, 90, 5)] + [92.5])
NAMES = {"AT": "Austria", "BE": "Belgium", "BG": "Bulgaria", "CY": "Cyprus", "CZ": "Czechia",
         "DE": "Germany", "DK": "Denmark", "EE": "Estonia", "EL": "Greece", "ES": "Spain",
         "FI": "Finland", "FR": "France", "HR": "Croatia", "HU": "Hungary", "IE": "Ireland",
         "IT": "Italy", "LT": "Lithuania", "LU": "Luxembourg", "LV": "Latvia", "MT": "Malta",
         "NL": "Netherlands", "PL": "Poland", "PT": "Portugal", "RO": "Romania", "SE": "Sweden",
         "SI": "Slovenia", "SK": "Slovakia", "IS": "Iceland", "NO": "Norway", "CH": "Switzerland"}


def cube(df):
    """Deaths and exposure as (geo, year, age group) arrays; a group with any missing age is NaN."""
    geos = list(dict.fromkeys(df["geo"]))
    years = np.sort(df["year"].unique())
    g = df.assign(grp=np.searchsorted(EDGES, df["age"], side="right") - 1,
                  miss=df["deaths"].isna() | df["exposure"].isna())
    agg = g.groupby(["geo", "year", "grp"])[["deaths", "exposure"]].sum(min_count=1)
    agg[g.groupby(["geo", "year", "grp"])["miss"].any()] = np.nan
    agg = agg.reindex(pd.MultiIndex.from_product([geos, years, range(len(LABELS))]))
    shape = (len(geos), len(years), len(LABELS))
    e0 = df.groupby(["geo", "year"])["e0_eurostat"].first()
    e0 = e0.reindex(pd.MultiIndex.from_product([geos, years]))
    return (agg["deaths"].to_numpy().reshape(shape), agg["exposure"].to_numpy().reshape(shape),
            geos, years, e0.to_numpy().reshape(shape[:2]))


D, E, GEOS, YEARS, E0_EUROSTAT = cube(DF)
C, NY, A = D.shape
YI = {int(y): i for i, y in enumerate(YEARS)}
print(f"{C} countries x {NY} years x {A} age groups; missing cells: {int(np.isnan(D).sum())} "
      f"(all before 2013: {bool(np.isnan(D[:, YI[2013]:]).sum() == 0)})")
print("deaths in the 30 countries (millions):",
      {y: round(float(np.nansum(D[:, YI[y]])) / 1e6, 2)
       for y in (2000, 2010, 2019, 2020, 2021, 2022, 2023, 2024)})
print("smallest countries in 2019, people (thousands) and deaths:",
      {g: (round(E[GEOS.index(g), YI[2019]].sum() / 1e3),
           round(float(D[GEOS.index(g), YI[2019]].sum()))) for g in ("IS", "MT", "LU", "CY", "EE")})
print(f"Germany 2019: {D[GEOS.index('DE'), YI[2019]].sum():,.0f} deaths")

POP = DF.pivot_table(index=["geo", "year"], columns="age", values="pop")  # 1 January, ages 0..90+
DTH = DF.pivot_table(index=["geo", "year"], columns="age", values="deaths", dropna=False)
GRP = np.searchsorted(EDGES, np.arange(91), side="right") - 1


def break_ratio(g, t, adjust=False):
    """Census level over the old level, by age group, for a break between 1 January t-1 and t.

    Also returns the step in the whole population and the deaths-weighted step at ages 60-89.
    adjust=True adds each cohort's deaths back, so that the ratios measure migration and the
    break only, not a year of unusual mortality."""
    X, Dx = POP.loc[g], DTH.loc[g]

    def coh(y):  # cohort ratio P_y(a) / P_(y-1)(a-1): survival and migration over year y-1
        p0 = X.loc[y - 1, 0:88].to_numpy()
        r = X.loc[y, 1:89].to_numpy() / p0
        r90 = X.loc[y, 90] / X.loc[y - 1, [89, 90]].sum()
        if adjust:  # the cohort's deaths in year y-1: half those at its two ages, all of 90+
            r = r + 0.5 * (Dx.loc[y - 1, 0:88].to_numpy() + Dx.loc[y - 1, 1:89].to_numpy()) / p0
            r90 = r90 + (0.5 * Dx.loc[y - 1, 89] + Dx.loc[y - 1, 90]) / X.loc[y - 1, [89, 90]].sum()
        return np.r_[np.nan, r, r90]

    with np.errstate(divide="ignore", invalid="ignore"):  # Hungary's 2006-11 have no ages 85+
        r = coh(t) / ((coh(t - 1) + coh(t + 1)) / 2)
    r[0] = r[1]
    w = X.loc[t].to_numpy()
    yrs = [t - 2, t - 1, t, t + 1]
    tot = X.loc[yrs, X.loc[yrs].notna().all(axis=0)].sum(axis=1)  # the ages present in all four
    step = (tot[t] / tot[t - 1]) / ((tot[t - 1] / tot[t - 2] + tot[t + 1] / tot[t]) / 2)
    dw = Dx.loc[t, 60:89].to_numpy()
    ok = np.isfinite(r[60:90]) & np.isfinite(dw)
    old = np.average(r[60:90][ok], weights=dw[ok])
    rg = np.array([np.nanmean(r[GRP == j]) if np.isnan(w[GRP == j]).any()
                   else np.average(r[GRP == j], weights=w[GRP == j]) for j in range(A)])
    return rg, step, old


E_RAW, E = E, E.copy()  # the published exposures, and a copy we correct
fixed = {}
for (g, t), _ in DF[DF["pop_break"] == 1].groupby(["geo", "year"]):
    if t - 2 < YEARS[0]:
        print(f"break {g} {t}: not measurable (no year t-2 in the data)")
        continue
    r, step, old = break_ratio(g, t)
    old_adj = break_ratio(g, t, adjust=True)[2]
    print(f"break {g} {t}: whole population {step - 1:+.1%} beyond the usual change; ages 60-89 "
          f"{old - 1:+.1%} ({old_adj - 1:+.1%} with deaths added back)", end="")
    if abs(step - 1) > 0.01:
        c = GEOS.index(g)
        E[c, YEARS < t - 1] *= r          # years whose exposure is all before the break
        E[c, YI[t - 1]] *= (1 + r) / 2    # the year that straddles it
        fixed[f"{g} {t}"] = r - 1
        print("  -> corrected", end="")
    print()
print(pd.DataFrame(fixed, index=LABELS).T.iloc[:, [3, 5, 7, 9, 13, 15, 17, 19]].round(3)
      .to_string())
E0_OBS = e0_abridged(D / E)  # (geo, year), observed, no uncertainty
e0_raw = e0_abridged(D / E_RAW)
for g, y in (("PL", 2019), ("CZ", 2019), ("DE", 2009)):
    print(f"{NAMES[g]} {y}: life expectancy {e0_raw[GEOS.index(g), YI[y]]:.2f} on the published "
          f"populations, {E0_OBS[GEOS.index(g), YI[y]]:.2f} on the census level")
for g in ("CZ", "PL"):  # the 2021 steps measured with 2020's and 2021's deaths taken out
    c, r_adj = GEOS.index(g), break_ratio(g, 2021, adjust=True)[0]
    Ea = E_RAW[c, YI[2019]] * r_adj
    print(f"{NAMES[g]} 2021 step with deaths added back: 2019 life expectancy "
          f"{e0_abridged(D[c, YI[2019]] / Ea):.2f}")
show =["ES", "IT", "SE", "DE", "PL", "BG", "SK", "NO"]
print(pd.DataFrame({y: E0_OBS[[GEOS.index(g) for g in show], YI[y]]
                    for y in (2010, 2019, 2020, 2021, 2022, 2024)},
                   index=[NAMES[g] for g in show]).round(2).to_string())

FIT0, FIT1 = 2000, 2019
fit_years = (YEARS >= FIT0) & (YEARS <= FIT1)


def svd_lee_carter(Dc, Ec):
    """Lee-Carter by SVD for one country: a (age), b (age, sums to 1), k (year)."""
    lm = np.log((Dc + 0.5) / Ec)
    ok = np.isfinite(lm)
    a = np.nanmean(lm, axis=0)
    R = np.where(ok, lm - a, 0.0)
    U, s, Vt = np.linalg.svd(R, full_matrices=False)
    b, k = Vt[0], U[:, 0] * s[0]
    if b.sum() < 0:
        b, k = -b, -k
    return a, b / b.sum(), k * b.sum(), s[0] ** 2 / (s ** 2).sum()


cE = GEOS.index("ES")
a_s, b_s, k_s, share = svd_lee_carter(D[cE, fit_years], E[cE, fit_years])
d_s = np.diff(k_s)
drift_s, sig_s = d_s.mean(), d_s.std(ddof=1)
print(f"Spain 2000-2019: the first component explains {share:.1%} of the variation in log rates "
      "around a_x")
print(f"k falls by {drift_s:.2f} a year (sd of yearly changes {sig_s:.2f}); b is largest at "
      f"{LABELS[int(np.argmax(b_s))]} ({b_s.max():.3f}) and smallest at "
      f"{LABELS[int(np.argmin(b_s))]} ({b_s.min():.3f})")

fig, axes = plt.subplots(1, 3, figsize=(12, 3.6))
axes[0].plot(MID, a_s, "o-", color=INK, ms=3)
axes[0].set(title="a: the average log rate", xlabel="age")
axes[1].plot(MID, b_s, "o-", color=HI, ms=3)
axes[1].set(title="b: how much each age follows k", xlabel="age")
axes[2].plot(YEARS[fit_years], k_s, "o-", color=LO, ms=3)
axes[2].set(title="k: the period index", xlabel="year");

def lee_carter_model(Dg, Eg, geos, years, lik="nb", pooled=True, zero_sum=True):
    """Bayesian Lee-Carter on a dense (geo, year, age) grid; NaN cells are left out."""
    Cg, Tg, Ag = Dg.shape
    ok = np.isfinite(Dg) & np.isfinite(Eg)
    Dd = np.where(ok, np.round(np.nan_to_num(Dg)), 0).astype(np.int64)
    lE = np.log(np.where(ok, Eg, 1.0))
    coords = {"geo": geos, "age": LABELS, "year": years, "step": years[1:]}
    with pm.Model(coords=coords) as model:
        if pooled:
            a_eu = pm.Normal("a_eu", -5.0, 3.0, dims="age")
            s_a = pm.HalfNormal("s_a", 0.5, dims="age")
            a = pm.Normal("a", a_eu, s_a, dims=("geo", "age"))  # centred: the data are strong
            beta_eu = pm.ZeroSumNormal("beta_eu", sigma=1.5, dims="age")
            centre = beta_eu[None, :]
        else:
            a = pm.Normal("a", -5.0, 3.0, dims=("geo", "age"))
            centre = 0.0
        # log-sensitivities with zero-sum rows (the softmax ignores a common shift), and an RW1
        # smoothness penalty over age on their deviation from the European pattern
        if zero_sum:
            beta_r = pm.Flat("beta_r", shape=(Cg, Ag - 1))
            beta = pm.Deterministic(
                "beta", pt.concatenate([beta_r, -beta_r.sum(axis=1, keepdims=True)], axis=1),
                dims=("geo", "age"))
        else:  # Failure 1: every row of beta free, so a constant added to a row changes nothing
            beta = pm.Flat("beta", dims=("geo", "age"))
        dev = beta - centre
        dev = dev - dev.mean(axis=1, keepdims=True)
        tau_b = pm.HalfNormal("tau_b", 0.3)
        pm.Potential("b_smooth", -0.5 * pt.sum(pt.sqr(pt.diff(dev, axis=1) / tau_b))
                     - Cg * (Ag - 1) * pt.log(tau_b))
        b = pm.Deterministic("b", pt.special.softmax(beta, axis=1), dims=("geo", "age"))
        if pooled:
            drift_eu = pm.Normal("drift_eu", -0.4, 0.4)
            s_drift = pm.HalfNormal("s_drift", 0.3)
            z_drift = pm.Normal("z_drift", 0, 1, dims="geo")
            drift = pm.Deterministic("drift", drift_eu + s_drift * z_drift, dims="geo")
            mu_lsk = pm.Normal("mu_lsk", np.log(0.3), 1.0)
            s_lsk = pm.HalfNormal("s_lsk", 0.5)
            z_sk = pm.Normal("z_sk", 0, 1, dims="geo")
            sigma_k = pm.Deterministic("sigma_k", pt.exp(mu_lsk + s_lsk * z_sk), dims="geo")
        else:
            drift = pm.Normal("drift", -0.4, 0.4, dims="geo")
            sigma_k = pm.LogNormal("sigma_k", np.log(0.3), 1.0, dims="geo")
        inc = pm.Normal("inc", drift[:, None], sigma_k[:, None], dims=("geo", "step"))
        kraw = pt.concatenate([pt.zeros((Cg, 1)), pt.cumsum(inc, axis=1)], axis=1)
        k = pm.Deterministic("k", kraw - kraw.mean(axis=1, keepdims=True), dims=("geo", "year"))
        eta = a[:, None, :] + b[:, None, :] * k[:, :, None] + lE
        mu = pt.exp(eta)
        if lik in ("nb", "nb_pymc"):
            if pooled:
                mu_lphi = pm.Normal("mu_lphi", np.log(1000.0), 1.5)
                s_lphi = pm.HalfNormal("s_lphi", 1.0)
                z_phi = pm.Normal("z_phi", 0, 1, dims="geo")
                phi = pm.Deterministic("phi", pt.exp(mu_lphi + s_lphi * z_phi), dims="geo")
            else:
                phi = pm.LogNormal("phi", np.log(1000.0), 1.5, dims="geo")
            ph = phi[:, None, None]
            if lik == "nb":  # NB2 log-likelihood written out, constants in the data dropped
                lp = (pt.gammaln(Dd + ph) - pt.gammaln(ph) + ph * pt.log(ph)
                      - (ph + Dd) * pt.log(ph + mu) + Dd * eta)
            else:            # the same through PyMC's own distribution, for the timing below
                lp = pm.logp(pm.NegativeBinomial.dist(mu=mu, alpha=ph), Dd)
        else:
            lp = Dd * eta - mu
        pm.Potential("lik", pt.sum(lp * ok))
    # start every chain from the SVD fit of each country
    init_a, init_b, init_k = [], [], []
    for c in range(Cg):
        a_c, b_c, k_c, _ = svd_lee_carter(Dg[c], Eg[c])
        a_c = np.where(np.isfinite(a_c), a_c, np.nanmean(np.log((Dg + 0.5) / Eg), axis=(0, 1)))
        init_a.append(a_c)
        init_b.append(np.log(np.clip(b_c, 1e-3, None)))
        init_k.append(k_c)
    lb = np.array(init_b)
    lb -= lb.mean(1, keepdims=True)
    init = {"a": np.array(init_a), "inc": np.diff(np.array(init_k), axis=1)}
    init.update({"beta_r": lb[:, :-1]} if zero_sum else {"beta": lb})
    if pooled:
        init.update(a_eu=np.array(init_a).mean(0), beta_eu=lb.mean(0))
    return model, init, ok


def fit(model, init, draws=1000, label=""):
    t0 = time.time()
    with model:
        idata = pm.sample(draws=draws, tune=400, chains=4, random_seed=RANDOM_SEED, initvals=init,
                          progressbar=False)
    TIMES[label] = time.time() - t0
    ss = idata.sample_stats
    print(f"[{label}] {TIMES[label]:.0f} s; {idata.posterior.attrs['tuning_steps']} tuning steps, "
          f"{idata.posterior.sizes['draw']} draws x 4 chains; "
          f"divergences {int(ss['diverging'].sum())}; "
          f"tree depth {int(ss['depth'].min())}-{int(ss['depth'].max())}")
    return idata


def check(idata, names, model):
    """r_hat and bulk ESS for the named quantities, then over every free variable, one at a time."""
    rows = {}

    def row(v):
        x = idata.posterior[v]
        n = x.size // (x.sizes["chain"] * x.sizes["draw"])
        return n, float(az.rhat(x).max()), float(az.ess(x).min())

    for v in names:
        rows[v] = row(v)
    free = [row(rv.name) for rv in model.free_RVs]
    n_free = sum(v.size for v in model.initial_point().values())  # on the sampling scale
    rows["ALL free variables"] = (n_free, max(f[1] for f in free), min(f[2] for f in free))
    print(pd.DataFrame(rows, index=["values", "max r_hat", "min ESS (bulk)"]).T.round(3)
          .to_string())


m_es, init_es, _ = lee_carter_model(D[[cE]][:, fit_years], E[[cE]][:, fit_years], ["ES"],
                                    YEARS[fit_years], pooled=False)
id_es = fit(m_es, init_es, label="Spain alone")
check(id_es, ["a", "b", "k", "drift", "sigma_k", "phi", "tau_b"], m_es)

po = id_es.posterior
k_b = po["k"].values.reshape(-1, fit_years.sum())
print(f"drift: SVD {drift_s:.2f}; posterior {float(po['drift'].mean()):.2f} "
      f"(89%: {np.quantile(po['drift'], 0.055):.2f} to {np.quantile(po['drift'], 0.945):.2f})")
print(f"sigma_k: SVD {sig_s:.2f}; posterior {float(po['sigma_k'].mean()):.2f}; "
      f"phi {float(po['phi'].median()):.0f}")


def project_k(k_last, drift, sigma_k, H, rng_):
    """Random walk with drift from the last fitted year: (S, C, H)."""
    eps = rng_.standard_normal(k_last.shape + (H,))
    return (k_last[..., None] + drift[..., None] * np.arange(1, H + 1)
            + sigma_k[..., None] * np.cumsum(eps, axis=-1))


def thin(idata, names, n, rng_):
    """n draws of each variable as arrays with the draws first (float32)."""
    S = idata.posterior.sizes["chain"] * idata.posterior.sizes["draw"]
    pick = rng_.choice(S, n, replace=False) if n < S else np.arange(S)
    out = {}
    for v in names:
        x = idata.posterior[v].values
        out[v] = x.reshape((S,) + x.shape[2:])[pick].astype(np.float32)
    return out


P_es = thin(id_es, ["a", "b", "k", "drift", "sigma_k"], 1000, np.random.default_rng(1))
ks_es = project_k(P_es["k"][:, :, -1], P_es["drift"], P_es["sigma_k"], 5, np.random.default_rng(2))
e0_es = e0_abridged(np.exp(P_es["a"][:, :, None, :]
                           + P_es["b"][:, :, None, :] * ks_es[..., None]))[:, 0]
bq = np.quantile(e0_es, [0.05, 0.5, 0.95], axis=0)
print("Bayesian projection of Spain, 5% / 50% / 95% (and the classic band's width):")
print(pd.DataFrame({"5%": bq[0], "50%": bq[1], "95%": bq[2], "width": bq[2] - bq[0],
                    "classic width": e0_cl[2] - e0_cl[0],
                    "observed": E0_OBS[cE, [YI[y] for y in PROJ]]}, index=PROJ)
      .round(2).to_string())
fit19_es = float(np.median(e0_abridged(np.exp(P_es["a"] + P_es["b"] * P_es["k"][:, :, -1:]))))
fit19_cl = float(e0_abridged(np.exp(a_s + b_s * k_s[-1])))
print(f"2019: observed {E0_OBS[cE, YI[2019]]:.2f}, fitted {fit19_es:.2f} (Bayesian), "
      f"{fit19_cl:.2f} (SVD)")

fig, axes = plt.subplots(1, 2, figsize=(12, 3.8))
ax = axes[0]
lo_, hi_ = np.quantile(k_b, [0.055, 0.945], axis=0)
ax.fill_between(YEARS[fit_years], lo_, hi_, color=HI, alpha=0.3, lw=0, label="Bayesian, 89%")
ax.plot(YEARS[fit_years], k_s, "o", color=INK, ms=3, label="SVD")
ax.set(title="Spain's period index k", xlabel="year")
ax.legend(frameon=False)
ax = axes[1]
xs = np.r_[2019, PROJ]
for lvl, al in [(0.9, 0.18), (0.6, 0.30), (0.3, 0.45)]:
    q_ = np.quantile(e0_es, [(1 - lvl) / 2, (1 + lvl) / 2], axis=0)
    ax.fill_between(xs, np.r_[fit19_es, q_[0]], np.r_[fit19_es, q_[1]], color=HI, alpha=al, lw=0,
                    label="Bayesian 90 / 60 / 30% fan" if lvl == 0.6 else None)
ax.plot(xs, np.r_[fit19_cl, e0_cl[0]], color=INK, lw=1, ls="--")
ax.plot(xs, np.r_[fit19_cl, e0_cl[2]], color=INK, lw=1, ls="--", label="classic 90% band")
ax.plot(YEARS[YI[2010]:], E0_OBS[cE, YI[2010]:], "o", color=INK, ms=3.5, label="observed")
ax.set(title="Spain's life expectancy and the path it was on", xlabel="year", ylabel="years")
ax.legend(frameon=False, loc="upper left");

po = id_main.posterior
drift_m = po["drift"].mean(("chain", "draw")).values
k_svd = [svd_lee_carter(D[c, fit_years], E[c, fit_years])[2] for c in range(C)]
drift_svd = np.array([(k[-1] - k[0]) / (fit_years.sum() - 1) for k in k_svd])
print(f"European drift {float(po['drift_eu'].mean()):.2f} (sd between countries s_drift "
      f"{float(po['s_drift'].mean()):.3f}); tau_b {float(po['tau_b'].mean()):.2f}")
print("phi (NB size) from", {g: round(float(po['phi'].sel(geo=g).mean())) for g in
                             (GEOS[int(np.argmin(po['phi'].mean(('chain', 'draw')).values))],
                              GEOS[int(np.argmax(po['phi'].mean(('chain', 'draw')).values))])},
      f"median over countries {np.median(po['phi'].mean(('chain', 'draw')).values):.0f}")
pop19 = E[:, YI[2019]].sum(-1)
order = np.argsort(pop19)
print(pd.DataFrame({"pop 2019 (M)": pop19 / 1e6, "SVD drift": drift_svd, "pooled drift": drift_m,
                    "sigma_k": po["sigma_k"].mean(("chain", "draw")).values},
                   index=[NAMES[g] for g in GEOS]).iloc[order[:6]].round(2).to_string())

fig, ax = plt.subplots(figsize=(10, 5))
srt = np.argsort(drift_m)
for i, c in enumerate(srt):
    ax.annotate("", (drift_m[c], i), (drift_svd[c], i),
                arrowprops=dict(arrowstyle="->", color=LO if pop19[c] < 1.5e6 else MUTED, lw=1.2))
    ax.plot(drift_svd[c], i, "o", color=LO if pop19[c] < 1.5e6 else MUTED, ms=4)
ax.axvline(float(po["drift_eu"].mean()), color=HI, lw=1, ls="--")
ax.plot([], [], "o", color=LO, ms=4, label="under 1.5 million people")
ax.plot([], [], "o", color=MUTED, ms=4, label="larger countries")
ax.plot([], [], color=HI, lw=1, ls="--", label="European drift")
ax.legend(frameon=False, loc="lower right", fontsize=8.5)
move = np.abs(drift_m - drift_svd)
print("smallest moves:", {NAMES[GEOS[c]]: round(float(move[c]), 2) for c in np.argsort(move)[:5]})
print("moves of the four largest:", {NAMES[GEOS[c]]: round(float(move[c]), 2)
                                     for c in np.argsort(-pop19)[:4]})
ax.set_yticks(range(C), [NAMES[GEOS[c]] for c in srt], fontsize=8.5)
ax.set_xlabel("drift of k per year (more negative = faster improvement)")
ax.set_title("Each country's improvement rate: alone (dot) and pooled (arrow head)", loc="left");

# posterior predictive check on the fitted cells: does the 90% interval of deaths cover 90%?
P_main = thin(id_main, ["a", "b", "k", "drift", "sigma_k", "phi"], 1000, np.random.default_rng(4))
Df, Ef = D[:, fit_years], E[:, fit_years]


def ppc_coverage(P, lik, rng_):
    m_ = np.exp(P["a"][:, :, None, :] + P["b"][:, :, None, :] * P["k"][:, :, :, None])
    m_ = m_.astype(np.float64)
    mu_ = m_ * np.nan_to_num(Ef)[None]
    expected = mu_.mean(0)
    if lik == "nb":
        ph = P["phi"][:, :, None, None].astype(np.float64)
        mu_ = rng_.gamma(ph, mu_ / ph)
    y = rng_.poisson(mu_)
    lo_, hi_ = np.quantile(y, [0.05, 0.95], axis=0)
    inside = np.where(ok_main, (Df >= lo_) & (Df <= hi_), np.nan)
    by_size = {}
    for a, b in ((0, 10), (10, 1000), (1000, np.inf)):  # small cells are covered by discreteness
        cell = (expected >= a) & (expected < b) & ok_main
        by_size[f"{a}-{b}"] = (round(float(cell.sum() / ok_main.sum()), 2),
                               round(float(np.nanmean(inside[cell])), 3))
    return np.nanmean(inside), by_size, m_


cov_nb, size_nb, m_fit = ppc_coverage(P_main, "nb", np.random.default_rng(5))
e0_fit = e0_abridged(m_fit)  # (S, C, T)
E0_FIT19 = np.median(e0_fit[:, :, -1], 0)
gap19 = E0_FIT19 - E0_OBS[:, YI[2019]]
rmse = np.sqrt(np.nanmean((np.median(e0_fit, 0) - E0_OBS[:, fit_years]) ** 2, axis=1))
print(f"90% predictive intervals of deaths cover {cov_nb:.1%} of the {int(ok_main.sum()):,} cells")
print("by expected deaths in the cell (share of cells, coverage):", size_nb)
print(f"fitted minus observed e0 in 2019: {gap19.min():+.2f} to {gap19.max():+.2f} years; "
      f"root-mean-square error of fitted e0 over 2000-2019: {rmse.min():.2f}-{rmse.max():.2f} "
      "years per country, largest in "
      f"{', '.join(NAMES[GEOS[c]] for c in np.argsort(rmse)[-3:][::-1])}")
del m_fit