Showing uncertainty: one posterior, twenty displays

# Everything in g from here on: F is the mean curve, Y a new reading, both (sample, t_grid).
F = stacked(crash_pred.predictions["f"]) * ACCEL_SD
Y = stacked(crash_pred.predictions["y"]) * ACCEL_SD
accel = mcycle.accel.to_numpy()
print(F.shape, Y.shape)

data.describe("radon")
data.describe("us_counties_geojson")
radon = data.load("radon")
counties = np.sort(radon.county.unique())
county_idx = pd.Index(counties).get_indexer(radon.county)
log_u = np.log(radon.groupby("county").Uppm.first().loc[counties].to_numpy())
n_homes = radon.groupby("county").size().loc[counties].to_numpy()
county_fips = radon.groupby("county").fips.first().loc[counties].astype(int).to_numpy()

with pm.Model(coords={"county": counties, "obs": radon.index}) as radon_model:
    county_id = pm.Data("county_id", county_idx, dims="obs")
    floor_d = pm.Data("floor", radon.floor.to_numpy().astype(float), dims="obs")
    log_u_d = pm.Data("log_u", log_u, dims="county")
    gamma_0 = pm.Normal("gamma_0", 1, 2)
    gamma_1 = pm.Normal("gamma_1", 0, 1)
    sigma_alpha = pm.HalfNormal("sigma_alpha", 1)
    z = pm.Normal("z", 0, 1, dims="county")
    alpha = pm.Deterministic("alpha", gamma_0 + gamma_1 * log_u_d + sigma_alpha * z, dims="county")
    beta = pm.Normal("beta", 0, 1)
    sigma = pm.HalfNormal("sigma", 1)
    pm.Normal("y", alpha[county_id] + beta * floor_d, sigma,
              observed=radon.log_radon.to_numpy(), dims="obs")
    radon_idata = pm.sample(random_seed=RANDOM_SEED, progressbar=False)

print("divergences:", int(radon_idata.sample_stats["diverging"].sum()))
az.summary(radon_idata, var_names=["gamma_0", "gamma_1", "sigma_alpha", "beta", "sigma"], round_to=2)

# County geometric-mean basement radon in pCi/L, (sample, county). The EPA action level is 4.
GM = np.exp(stacked(radon_idata.posterior["alpha"]))
SIGMA = stacked(radon_idata.posterior["sigma"])
gm_mean = GM.mean(0)
gm_lo, gm_hi = np.quantile(GM, [0.05, 0.95], axis=0)
p_over4 = (GM > 4).mean(0)
raw_gm = np.exp(radon[radon.floor == 0].groupby("county").log_radon.mean()).reindex(counties)
print(f"county GM: {gm_mean.min():.1f} to {gm_mean.max():.1f} pCi/L; "
      f"90% interval widths {np.min(gm_hi - gm_lo):.1f} to {np.max(gm_hi - gm_lo):.1f}")

data.describe("epl_2425")
epl = data.load("epl_2425")
teams = np.sort(epl.HomeTeam.unique())
home_i = pd.Index(teams).get_indexer(epl.HomeTeam)
away_i = pd.Index(teams).get_indexer(epl.AwayTeam)

with pm.Model(coords={"team": teams, "match": epl.index}) as epl_model:
    intercept = pm.Normal("intercept", 0.2, 0.5)
    home_adv = pm.Normal("home_adv", 0.2, 0.3)
    sd_att = pm.HalfNormal("sd_att", 0.5)
    sd_def = pm.HalfNormal("sd_def", 0.5)
    att = pm.ZeroSumNormal("att", sigma=sd_att, dims="team")
    dfn = pm.ZeroSumNormal("def", sigma=sd_def, dims="team")
    pm.Deterministic("strength", att + dfn, dims="team")
    lam_h = pm.Deterministic("lam_h", pm.math.exp(intercept + home_adv + att[home_i] - dfn[away_i]), dims="match")
    lam_a = pm.Deterministic("lam_a", pm.math.exp(intercept + att[away_i] - dfn[home_i]), dims="match")
    pm.Poisson("home_goals", lam_h, observed=epl.FTHG.to_numpy(), dims="match")
    pm.Poisson("away_goals", lam_a, observed=epl.FTAG.to_numpy(), dims="match")
    epl_idata = pm.sample(random_seed=RANDOM_SEED, progressbar=False)

print("divergences:", int(epl_idata.sample_stats["diverging"].sum()))
az.summary(epl_idata, var_names=["intercept", "home_adv", "sd_att", "sd_def"], round_to=2)

def league_table(gh, ga):
    """Points, goal difference and goals for each team from home/away goals of shape (..., match)."""
    H = np.eye(len(teams))[home_i]  # (match, team) one-hot
    A = np.eye(len(teams))[away_i]
    hp = 3 * (gh > ga) + (gh == ga)
    ap = 3 * (ga > gh) + (gh == ga)
    pts = hp @ H + ap @ A
    gd = (gh - ga) @ H + (ga - gh) @ A
    gf = gh @ H + ga @ A
    return pts, gd, gf


def ranks_from(pts, gd, gf):
    """Final position (1 = champion) for each team, ties broken by goal difference then goals."""
    order = np.lexsort((-gf, -gd, -pts), axis=-1)  # last key is primary
    r = np.empty_like(order)
    np.put_along_axis(r, order, np.arange(1, len(teams) + 1) + 0 * order, axis=-1)
    return r


actual_pts, actual_gd, actual_gf = league_table(epl.FTHG.to_numpy(), epl.FTAG.to_numpy())
actual_pos = ranks_from(actual_pts, actual_gd, actual_gf)
finish_order = teams[np.argsort(actual_pos)]  # champion first
STRENGTH = stacked(epl_idata.posterior["strength"])
print("2024/25 final table (top 6):", ", ".join(f"{tm} {int(actual_pts[i])}" for tm, i in
      zip(finish_order[:6], np.argsort(actual_pos)[:6])))

def band(ax, x, samples, prob, color, alpha, label=None, **kw):
    lo, hi = np.quantile(samples, [(1 - prob) / 2, (1 + prob) / 2], axis=0)
    ax.fill_between(x, lo, hi, color=color, alpha=alpha, lw=0, label=label, **kw)
    return lo, hi


fig, axes = plt.subplots(1, 2, figsize=(12, 4.2), sharey=True)
for ax, S, what in [(axes[0], F, "mean curve $f(t)$"), (axes[1], Y, "a new reading $y(t)$")]:
    ax.scatter(t, accel, s=12, color="k", alpha=0.5, zorder=3, label="data")
    band(ax, t_grid, S, 0.94, BLUE, 0.25, label=f"94% interval of {what}")
    ax.plot(t_grid, S.mean(0), color=BLUE, lw=2, label="posterior mean")
    ax.set(xlabel="time after impact (ms)", title=what)
axes[0].set_ylabel("head acceleration (g)")
axes[0].legend(loc="lower right", fontsize=9)

# How much of the data does each band contain? Interpolate the bands to the data times.
inside = {}
for name, S in [("f", F), ("y", Y)]:
    lo, hi = np.quantile(S, [0.03, 0.97], axis=0)
    inside[name] = np.mean((accel >= np.interp(t, t_grid, lo)) & (accel <= np.interp(t, t_grid, hi)))
print({k: f"{v:.0%} of the data inside the 94% band" for k, v in inside.items()})

fig, ax = plt.subplots(figsize=(8, 4.5))
ax.scatter(t, accel, s=12, color="k", alpha=0.5, zorder=3)
levels = [0.98, 0.9, 0.8, 0.6, 0.4, 0.2]
for i, p in enumerate(levels):
    band(ax, t_grid, Y, p, BLUE, 0.12 + 0.06 * i, label=f"{p:.0%}" if p in (0.98, 0.6, 0.2) else None)
ax.plot(t_grid, np.median(Y, axis=0), color="white", lw=1.5)
ax.set(xlabel="time after impact (ms)", ylabel="head acceleration (g)",
       title="Fan chart of the predictive distribution of a reading (98% ... 20%)")
ax.legend(title="central interval", loc="lower right", fontsize=9);

lo94, hi94 = np.quantile(F, [0.03, 0.97], axis=0)
idx = rng.choice(F.shape[0], 60, replace=False)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.2))
ax = axes[0]
ax.fill_between(t_grid, lo94, hi94, color=BLUE, alpha=0.15, lw=0, label="94% pointwise band of f")
ax.plot(t_grid, F[idx].T, color=BLUE, lw=0.6, alpha=0.4)
ax.plot(t_grid, lo94, color=ORANGE, lw=2, label="lower band edge (not a plausible curve)")
ax.scatter(t, accel, s=8, color="k", alpha=0.4, zorder=3)
ax.set(xlabel="time after impact (ms)", ylabel="head acceleration (g)", title="60 posterior draws of $f$")
ax.legend(loc="lower right", fontsize=9)

# Feature uncertainty: the depth and the timing of the trough, per draw.
trough_depth = F.min(axis=1)
trough_time = t_grid[F.argmin(axis=1)]
ax = axes[1]
ax.hist(trough_depth, bins=40, color=BLUE, alpha=0.8)
ax.axvline(lo94.min(), color=ORANGE, lw=2, label=f"deepest point of the band edge: {lo94.min():.0f} g")
ax.set(xlabel="minimum of $f$ over the record (g)", ylabel="posterior draws", title="How deep is the trough?")
ax.legend(fontsize=9)
print(f"P(trough deeper than the band's lowest point) = {(trough_depth < lo94.min()).mean():.3f}")
print(f"trough time: median {np.median(trough_time):.1f} ms, 90% interval "
      f"{np.quantile(trough_time, 0.05):.1f}-{np.quantile(trough_time, 0.95):.1f} ms")

n_frames = 40
frame_idx = rng.choice(Y.shape[0], n_frames, replace=False)

fig, ax = plt.subplots(figsize=(8, 4.5), dpi=72)  # 40 embedded frames: keep them small
ax.fill_between(t_grid, *np.quantile(Y, [0.05, 0.95], axis=0), color=BLUE, alpha=0.1, lw=0)
ax.scatter(t, accel, s=12, color="k", alpha=0.35, zorder=3)
(curve,) = ax.plot(t_grid, F[frame_idx[0]], color=BLUE, lw=2)
dots = ax.scatter(t_grid[::3], Y[frame_idx[0], ::3], s=18, color=ORANGE, zorder=4)
title = ax.set_title("posterior draw 1 of 40")
ax.set(xlabel="time after impact (ms)", ylabel="head acceleration (g)", ylim=(-220, 120))
plt.close(fig)  # keep the static figure out of the output; the animation is shown below


def update(i):
    curve.set_ydata(F[frame_idx[i]])
    dots.set_offsets(np.c_[t_grid[::3], Y[frame_idx[i], ::3]])
    title.set_text(f"posterior draw {i + 1} of 40")
    return curve, dots, title


hops = animation.FuncAnimation(fig, update, frames=n_frames, interval=400, blit=False)
HTML(hops.to_jshtml(default_mode="loop"))

def quantile_dotplot(ax, samples, n_dots=20, n_bins=24, color=BLUE, xlim=None, max_stack=12):
    """Wilkinson-style dotplot of the n_dots quantiles of `samples`, stacked in n_bins columns."""
    q = np.quantile(samples, (np.arange(n_dots) + 0.5) / n_dots)
    lo, hi = xlim if xlim is not None else (q.min(), q.max())
    edges = np.linspace(lo, hi, n_bins + 1)
    width = edges[1] - edges[0]
    cols = np.clip(np.digitize(q, edges) - 1, 0, n_bins - 1)
    heights = np.zeros(n_bins, int)
    for c in cols:
        ax.add_patch(Ellipse((edges[c] + width / 2, heights[c] + 0.5), width * 0.92, 0.92, color=color))
        heights[c] += 1
    ax.set(xlim=(lo, hi), ylim=(0, max_stack + 0.5), yticks=[])
    ax.set_aspect(width)  # one dot diameter = one bin width, so the dots are round
    sns.despine(ax=ax, left=True)
    return q


fig, axes = plt.subplots(1, 3, figsize=(13, 4))
for ax, t0 in zip(axes, [15, 20, 35]):
    j = np.argmin(np.abs(t_grid - t0))
    q = quantile_dotplot(ax, Y[:, j], xlim=(-250, 150), max_stack=11)
    ax.axvline(-100, color=ORANGE, lw=2)
    n_below = int((q < -100).sum())
    ax.set(xlabel="head acceleration (g)")
    ax.set_title(f"{t0} ms after impact\n{n_below}/20 dots below -100 g (P = {(Y[:, j] < -100).mean():.2f})", fontsize=10)
fig.suptitle("Quantile dotplots of the predictive distribution; every dot is 5%", y=1.02);

fig, ax = plt.subplots(figsize=(8, 3.6))
ax.plot(t_grid, (F < -100).mean(0), color=BLUE, lw=2, label="P(mean curve $f(t)$ < -100 g)")
ax.plot(t_grid, (Y < -100).mean(0), color=ORANGE, lw=2, label="P(a reading $y(t)$ < -100 g)")
ax.set(xlabel="time after impact (ms)", ylabel="posterior probability", ylim=(0, 1.02))
ax.legend(loc="upper right", fontsize=9);

qs = np.quantile(Y, [0.01, 0.05, 0.25, 0.5, 0.75, 0.95, 0.99], axis=0)
p_below = (Y < -100).mean(0)
fig = go.Figure()
for (lo_i, hi_i, name, a) in [(0, 6, "98%", 0.12), (1, 5, "90%", 0.18), (2, 4, "50%", 0.3)]:
    fig.add_trace(go.Scatter(x=t_grid, y=qs[lo_i], line=dict(width=0), showlegend=False, hoverinfo="skip"))
    fig.add_trace(go.Scatter(x=t_grid, y=qs[hi_i], fill="tonexty", fillcolor=f"rgba(42,120,214,{a})",
                             line=dict(width=0), name=f"{name} interval", hoverinfo="skip"))
fig.add_trace(go.Scatter(
    x=t_grid, y=qs[3], name="median", line=dict(color=BLUE, width=2),
    customdata=np.c_[qs[1], qs[5], p_below],
    hovertemplate="median %{y:.0f} g<br>90%%: %{customdata[0]:.0f} to %{customdata[1]:.0f} g"
                  "<br>P(< -100 g) = %{customdata[2]:.2f}<extra></extra>",
))
fig.add_trace(go.Scatter(x=t, y=accel, mode="markers", name="data",
                         marker=dict(color="black", size=5, opacity=0.5), hoverinfo="skip"))
fig.update_layout(hovermode="x unified", template="plotly_white", width=800, height=420,
                  xaxis_title="time after impact (ms)", yaxis_title="head acceleration (g)",
                  legend=dict(orientation="h", y=-0.2))
fig

order = np.argsort(gm_mean)
q25, q75 = np.quantile(GM, [0.25, 0.75], axis=0)
x = np.arange(len(counties))

fig, ax = plt.subplots(figsize=(13, 4.5))
ax.axhline(4, color=GREY, lw=1, ls="--")
ax.text(0.5, 4.1, "EPA action level 4 pCi/L", color=GREY, fontsize=9)
ax.vlines(x, gm_lo[order], gm_hi[order], color=ORANGE, lw=1)
ax.vlines(x, q25[order], q75[order], color=ORANGE, lw=3.5)
ax.scatter(x, raw_gm.to_numpy()[order], s=8 + 1.5 * n_homes[order], facecolor="none",
           edgecolor="k", lw=0.8, zorder=3, label="raw county mean (area = homes measured)")
ax.scatter(x, gm_mean[order], s=14, color="k", zorder=4, label="posterior mean, 50% and 90% intervals")
ax.set(xticks=x, xlim=(-1, len(counties)), ylim=(0, 14), ylabel="basement radon, county geometric mean (pCi/L)")
ax.set_xticklabels(counties[order], rotation=90, fontsize=6.5)
ax.legend(loc="upper left", fontsize=9);

pick = np.r_[np.argsort(n_homes)[-8:], np.argsort(n_homes)[:8]]
pick = pick[np.argsort(gm_mean[pick])]
grid = np.linspace(1, 10, 300)
cmap = plt.get_cmap("Oranges")
n_norm = plt.Normalize(0, np.log(n_homes.max()))

fig, ax = plt.subplots(figsize=(8, 7))
for k, i in enumerate(pick):
    dens = gaussian_kde(GM[:, i])(grid)
    dens = dens / dens.max() * 1.6
    ax.fill_between(grid, k, k + dens, color=cmap(0.3 + 0.7 * n_norm(np.log(n_homes[i]))), alpha=0.9, lw=0)
    ax.plot(grid, k + dens, color="white", lw=0.8)
    ax.text(0.9, k + 0.15, f"{counties[i].title()}  (n = {n_homes[i]})", ha="right", fontsize=9)
ax.axvline(4, color=GREY, lw=1, ls="--")
ax.set(xlim=(0, 10), yticks=[], xlabel="basement radon, county geometric mean (pCi/L)",
       title="Posterior of 16 counties (darker = more homes measured)")
sns.despine(ax=ax, left=True);

geo = json.load(open(data.path("us_counties_geojson")))
mn = {int(ft["id"]): ft for ft in geo["features"] if ft["id"].startswith("27")}
polys = {fips: np.array(ft["geometry"]["coordinates"][0]) for fips, ft in mn.items()}
missing_fips = sorted(set(mn) - set(county_fips))
print("counties without data:", [mn[f]["properties"]["NAME"] for f in missing_fips])


def draw_map(ax, values, cmap, norm=None, title="", cbar_label="", hatch_mask=None, hatch="////"):
    """Choropleth of one value per county (in `counties` order) on `ax`; returns the collection."""
    pc = PolyCollection([polys[f] for f in county_fips], array=np.asarray(values), cmap=cmap, norm=norm,
                        edgecolor="white", linewidth=0.4)
    ax.add_collection(pc)
    ax.add_collection(PolyCollection([polys[f] for f in missing_fips], facecolor="#dddddd",
                                     edgecolor="white", linewidth=0.4, hatch="....", label="no data"))
    if hatch_mask is not None:
        ax.add_collection(PolyCollection([polys[f] for f, m in zip(county_fips, hatch_mask) if m],
                                         facecolor="none", edgecolor="#333333", linewidth=0.3, hatch=hatch))
    ax.autoscale()
    ax.set_aspect(1 / np.cos(np.deg2rad(46)))  # roughly equal-area at Minnesota's latitude
    ax.axis("off")
    ax.set_title(title, fontsize=11)
    if cbar_label:
        ax.figure.colorbar(pc, ax=ax, shrink=0.55, pad=0.01, label=cbar_label)
    return pc

straddles = (gm_lo < 4) & (gm_hi > 4)
fig, axes = plt.subplots(1, 2, figsize=(12, 6.5))
draw_map(axes[0], gm_mean, "Oranges", plt.Normalize(2, 7), "Posterior mean of county radon",
         "pCi/L (basement geometric mean)", hatch_mask=straddles)
draw_map(axes[1], gm_hi - gm_lo, "Purples", plt.Normalize(0, 3.5), "Width of the 90% interval", "pCi/L")
axes[0].text(0.55, 0.12, "hatched: 90% interval\nstraddles 4 pCi/L\ndotted grey: no data",
             transform=axes[0].transAxes, fontsize=9);

fig, ax = plt.subplots(figsize=(6.5, 6.5))
draw_map(ax, p_over4, "Purples", plt.Normalize(0, 1),
         "P(county basement geometric mean > 4 pCi/L)", "posterior probability")
print(f"{(p_over4 > 0.9).sum()} counties above 0.9, {(p_over4 < 0.1).sum()} below 0.1, "
      f"{((p_over4 >= 0.1) & (p_over4 <= 0.9)).sum()} genuinely undecided (out of {len(counties)})")

draw_ids = rng.choice(GM.shape[0], 8, replace=False)
fig, axes = plt.subplots(2, 4, figsize=(13, 7))
for ax, d in zip(axes.flat, draw_ids):
    draw_map(ax, GM[d], "Oranges", plt.Normalize(2, 7), f"draw {d}")
fig.suptitle("Eight draws of the county radon map (2 to 7 pCi/L, same scale as the mean map)", y=1.02);

mn_geojson = {"type": "FeatureCollection", "features": list(mn.values())}
county_df = pd.DataFrame({
    "fips": [f"{f:05d}" for f in county_fips], "county": [c.title() for c in counties],
    "mean": gm_mean.round(2), "lo": gm_lo.round(2), "hi": gm_hi.round(2),
    "p_over4": p_over4.round(2), "homes": n_homes,
})
fig = go.Figure(go.Choropleth(
    geojson=mn_geojson, locations=county_df.fips, z=county_df["mean"], zmin=2, zmax=7,
    colorscale="Oranges", marker_line_color="white", marker_line_width=0.5,
    colorbar=dict(title="pCi/L", thickness=12),
    customdata=county_df[["county", "lo", "hi", "p_over4", "homes"]].to_numpy(),
    hovertemplate="<b>%{customdata[0]}</b><br>mean %{z:.1f} pCi/L, 90%% interval %{customdata[1]:.1f}-%{customdata[2]:.1f}"
                  "<br>P(> 4) = %{customdata[3]:.2f}<br>%{customdata[4]} homes measured<extra></extra>",
))
fig.update_geos(fitbounds="locations", visible=False)  # NB: fitbounds is ignored when scope="usa" is set
fig.update_layout(title="County radon: hover for the posterior summary", width=650, height=600,
                  margin=dict(l=0, r=0, t=40, b=0))
fig

anim_ids = rng.choice(GM.shape[0], 12, replace=False)
frames = [go.Frame(name=str(k), data=[go.Choropleth(z=GM[d], locations=county_df.fips)],
                   layout=go.Layout(title_text=f"posterior draw {d}")) for k, d in enumerate(anim_ids)]
fig = go.Figure(
    data=[go.Choropleth(geojson=mn_geojson, locations=county_df.fips, z=GM[anim_ids[0]], zmin=2, zmax=7,
                        colorscale="Oranges", marker_line_color="white", marker_line_width=0.5,
                        colorbar=dict(title="pCi/L", thickness=12), text=county_df.county,
                        hovertemplate="<b>%{text}</b>: %{z:.1f} pCi/L in this draw<extra></extra>")],
    frames=frames,
)
fig.update_geos(fitbounds="locations", visible=False)
fig.update_layout(
    title_text=f"posterior draw {anim_ids[0]}", width=650, height=620, margin=dict(l=0, r=0, t=40, b=0),
    updatemenus=[dict(type="buttons", x=0.05, y=0.05, buttons=[
        dict(label="play", method="animate",
             args=[None, dict(frame=dict(duration=600, redraw=True), fromcurrent=True, transition=dict(duration=0))]),
        dict(label="pause", method="animate", args=[[None], dict(frame=dict(duration=0, redraw=False), mode="immediate")]),
    ])],
    sliders=[dict(steps=[dict(method="animate", label=str(k),
                              args=[[str(k)], dict(frame=dict(duration=0, redraw=True), mode="immediate")])
                         for k in range(len(frames))], x=0.15, len=0.8, y=0.02)],
)
fig

def new_home_reading(county_name, n=4000):
    i = np.where(counties == county_name)[0][0]
    return np.exp(np.log(GM[:n, i]) + rng.normal(0, SIGMA[:n]))


# Two counties with about the same posterior mean but very different sample sizes.
few, many = np.where(n_homes <= 2)[0], np.where(n_homes >= 40)[0]
i_few, i_many = min(((i, j) for i in few for j in many), key=lambda ij: abs(gm_mean[ij[0]] - gm_mean[ij[1]]))
print(f"{counties[i_few].title()}: n = {n_homes[i_few]}, county GM {gm_mean[i_few]:.1f} "
      f"(90%: {gm_lo[i_few]:.1f}-{gm_hi[i_few]:.1f});  {counties[i_many].title()}: n = {n_homes[i_many]}, "
      f"county GM {gm_mean[i_many]:.1f} (90%: {gm_lo[i_many]:.1f}-{gm_hi[i_many]:.1f})")

fig, axes = plt.subplots(2, 2, figsize=(12, 7), gridspec_kw={"width_ratios": [2, 1]})
for row, i in zip(axes, [i_few, i_many]):
    reading = new_home_reading(counties[i])
    p = (reading > 4).mean()
    q = quantile_dotplot(row[0], np.log(reading), n_dots=50, n_bins=30, color=ORANGE,
                         xlim=(np.log(0.3), np.log(40)), max_stack=8)
    row[0].axvline(np.log(4), color="k", lw=2)
    ticks = [0.5, 1, 2, 4, 8, 16, 32]
    row[0].set(xticks=np.log(ticks), xticklabels=ticks, xlabel="basement reading of a new home (pCi/L, log axis)",
               title=f"{counties[i].title()} county (n = {n_homes[i]}): "
                     f"{int((q > np.log(4)).sum())} of 50 dots above 4 pCi/L")
    # icon array: 100 houses, filled in reading order so the count is visible at a glance
    n_over = int(round(100 * p))
    for k in range(100):
        row[1].add_patch(Rectangle((k % 10, 9 - k // 10), 0.85, 0.85,
                                   color=ORANGE if k < n_over else "#e4e4e0"))
    row[1].set(xlim=(0, 10), ylim=(0, 10), xticks=[], yticks=[], aspect="equal",
               title=f"{n_over} of 100 homes like this one\nwould exceed 4 pCi/L")

pos_idx = np.argsort(actual_pos)  # team indices, champion first
s_mean = STRENGTH.mean(0)
s50 = np.quantile(STRENGTH, [0.25, 0.75], axis=0)
s90 = np.quantile(STRENGTH, [0.05, 0.95], axis=0)
yy = np.arange(len(teams))[::-1]

fig, ax = plt.subplots(figsize=(8, 6.5))
ax.hlines(yy, s90[0, pos_idx], s90[1, pos_idx], color=BLUE, lw=1)
ax.hlines(yy, s50[0, pos_idx], s50[1, pos_idx], color=BLUE, lw=3.5)
ax.scatter(s_mean[pos_idx], yy, color="k", s=18, zorder=3)
ax.set(yticks=yy, xlabel="latent strength (attack + defence, log goal-rate units)",
       title="Posterior of team strength, ordered by the actual final position")
ax.set_yticklabels([f"{p}. {tm}  ({int(actual_pts[i])} pts)" for p, tm, i in
                    zip(range(1, 21), finish_order, pos_idx)], fontsize=9)
ax.axvline(0, color=GREY, lw=1, ls="--");

rank_draws = np.argsort(np.argsort(-STRENGTH, axis=1), axis=1) + 1  # (sample, team), 1 = strongest
rank_prob = np.stack([(rank_draws == k).mean(0) for k in range(1, 21)], axis=1)  # (team, rank)
rank_df = pd.DataFrame(rank_prob[pos_idx], index=finish_order, columns=range(1, 21))

fig, ax = plt.subplots(figsize=(11, 7))
sns.heatmap(rank_df, cmap="Blues", vmin=0, vmax=0.6, annot=rank_df.map(lambda v: f"{v:.2f}" if v >= 0.1 else ""),
            fmt="", annot_kws={"size": 7}, cbar_kws={"label": "P(team is the k-th strongest)", "shrink": 0.7},
            linewidths=0.5, linecolor="white", ax=ax)
for r in range(20):
    ax.add_patch(Rectangle((r, r), 1, 1, fill=False, edgecolor=ORANGE, lw=2))
ax.set(xlabel="rank of latent strength (1 = strongest)", ylabel="team, in actual finishing order",
       title="Rank-probability heatmap (orange box = actual final position)")
ax.tick_params(axis="y", rotation=0);

pair = (STRENGTH[:, :, None] > STRENGTH[:, None, :]).mean(0)  # (i, j) = P(i stronger than j)
pair = pair[np.ix_(pos_idx, pos_idx)]
np.fill_diagonal(pair, np.nan)
pair_df = pd.DataFrame(pair, index=finish_order, columns=finish_order)

fig, ax = plt.subplots(figsize=(10, 8.5))
sns.heatmap(pair_df, cmap="RdBu_r", center=0.5, vmin=0, vmax=1, square=True, linewidths=0.5, linecolor="white",
            cbar_kws={"label": "P(row team stronger than column team)", "shrink": 0.7}, ax=ax)
ax.set(xlabel="", ylabel="", title="Pairwise strength comparison (teams in actual finishing order)")
ax.tick_params(axis="x", rotation=90)
ax.tick_params(axis="y", rotation=0);

n_sim = 2000
sim = rng.choice(STRENGTH.shape[0], n_sim, replace=False)
lam_h = stacked(epl_idata.posterior["lam_h"])[sim]
lam_a = stacked(epl_idata.posterior["lam_a"])[sim]
gh, ga = rng.poisson(lam_h), rng.poisson(lam_a)
sim_pts, sim_gd, sim_gf = league_table(gh, ga)
sim_pos = ranks_from(sim_pts, sim_gd, sim_gf)  # (n_sim, team)

season = pd.DataFrame({
    "actual points": actual_pts.astype(int),
    "expected points": sim_pts.mean(0).round(1),
    "points 90% interval": [f"{lo:.0f}-{hi:.0f}" for lo, hi in np.quantile(sim_pts, [0.05, 0.95], axis=0).T],
    "luck (actual - expected)": (actual_pts - sim_pts.mean(0)).round(1),
    "P(title)": (sim_pos == 1).mean(0),
    "P(top 4)": (sim_pos <= 4).mean(0),
    "P(relegated)": (sim_pos >= 18).mean(0),
}, index=teams).loc[finish_order]
season.index.name = "team (actual order)"