Meta-analysis: learning from a database of studies that never shares a patient

import logging

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
from scipy import stats

from pymc_challenges import data

RANDOM_SEED = 42
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # this notebook runs ~50 fits: no banner per fit
BLUE, ORANGE, AQUA, GREY = "#2a78d6", "#eb6834", "#1baf7a", "#8a8a86"
print(f"PyMC {pm.__version__}, ArviZ {az.__version__}")


def fit(model, **kw):
    """One place for the sampler settings; every fit reports its divergences."""
    with model:
        idata = pm.sample(target_accept=0.95, random_seed=RANDOM_SEED, progressbar=False, **kw)
    print(f"  {model.name or 'model'}: divergences = {int(idata.sample_stats['diverging'].sum())}")
    return idata


def stacked(da):
    return da.stack(sample=("chain", "draw")).transpose("sample", ...).to_numpy()

data.describe("bcg")
bcg = data.load("bcg")
bcg["n_vacc"] = bcg.tpos + bcg.tneg
bcg["n_ctrl"] = bcg.cpos + bcg.cneg
bcg["rate_vacc"] = 1000 * bcg.tpos / bcg.n_vacc
bcg["rate_ctrl"] = 1000 * bcg.cpos / bcg.n_ctrl
bcg[["author", "year", "ablat", "alloc", "tpos", "n_vacc", "cpos", "n_ctrl", "rate_vacc", "rate_ctrl"]].round(1)

order = np.argsort(bcg.ablat.to_numpy())
fig, ax = plt.subplots(figsize=(9, 5.5))
for k, i in enumerate(order):
    r = bcg.iloc[i]
    ax.plot([r.rate_ctrl, r.rate_vacc], [k, k], color=GREY, lw=1 + np.log10(r.n_vacc + r.n_ctrl), zorder=1)
    ax.scatter(r.rate_ctrl, k, color=GREY, s=50, zorder=2)
    ax.scatter(r.rate_vacc, k, color=BLUE, s=50, zorder=3)
ax.set(yticks=range(len(bcg)), xlabel="tuberculosis cases per 1000 participants",
       title="Control (grey) vs vaccinated (blue) TB rate in each trial, ordered by latitude")
ax.set_yticklabels([f"{bcg.author[i]} ({bcg.year[i]}, {bcg.ablat[i]}°)" for i in order], fontsize=9)
ax.scatter([], [], color=GREY, label="control arm")
ax.scatter([], [], color=BLUE, label="BCG arm")
ax.legend(loc="lower right");

a, b, c, d = (bcg[col].to_numpy(float) for col in ["tpos", "tneg", "cpos", "cneg"])
i = 0
records = np.r_[np.ones(int(a[i])), np.zeros(int(b[i]))]  # the participant-level 0/1 outcomes, BCG arm

p_test = 0.03
lp_individual = pm.logp(pm.Bernoulli.dist(p=p_test), records).sum().eval()
lp_table = pm.logp(pm.Binomial.dist(n=a[i] + b[i], p=p_test), a[i]).eval()
log_binom_coef = stats.binom(a[i] + b[i], p_test).logpmf(a[i]) - (a[i] * np.log(p_test) + b[i] * np.log(1 - p_test))
print(f"individual records: {lp_individual:.3f}   table: {lp_table:.3f}   "
      f"difference = log binomial coefficient = {log_binom_coef:.3f} (a constant, so the posteriors agree)")

def summarise(idata, names, prefix=""):
    q = [0.05, 0.5, 0.95]
    return {n: np.quantile(stacked(idata.posterior[prefix + n]), q) for n in names}


fig, axes = plt.subplots(1, 2, figsize=(11, 2.6))
for ax, var, xl in [(axes[0], "mu", "average effect $\\mu$ (log odds ratio)"), (axes[1], "tau", "heterogeneity $\\tau$")]:
    for k, (name, idata, pre) in enumerate([("two-stage (normal)", idata_2, "two_stage::"), ("one-stage (binomial)", idata_1, "one_stage::")]):
        lo, med, hi = summarise(idata, [var], pre)[var]
        ax.plot([lo, hi], [k, k], color=BLUE, lw=2)
        ax.scatter(med, k, color="k", zorder=3)
        ax.text(hi + 0.02, k, f"{med:.2f} [{lo:.2f}, {hi:.2f}]", va="center", fontsize=9)
    ax.set(yticks=[0, 1], yticklabels=["two-stage", "one-stage"], xlabel=xl, ylim=(-0.6, 1.6))
    ax.margins(x=0.35)
fig.suptitle("Posterior medians and 90% intervals", y=1.05);

def forest(ax, idata, names, y_obs, se_obs, prefix, order=None, title=""):
    TH = stacked(idata.posterior[prefix + "theta"])
    MU = stacked(idata.posterior[prefix + "mu"])
    NEW = stacked(idata.posterior[prefix + "theta_new"])
    order = np.arange(len(names)) if order is None else order
    yy = np.arange(len(names))[::-1] + 2
    for row, i in zip(yy, order):
        ax.plot([y_obs[i] - 1.96 * se_obs[i], y_obs[i] + 1.96 * se_obs[i]], [row + 0.15, row + 0.15], color="k", lw=1)
        ax.scatter(y_obs[i], row + 0.15, color="k", s=8 + 200 / se_obs[i]**2 / (1 / se_obs**2).max(), marker="s", zorder=3)
        lo, q1, q3, hi = np.quantile(TH[:, i], [0.05, 0.25, 0.75, 0.95])
        ax.plot([lo, hi], [row - 0.15, row - 0.15], color=BLUE, lw=1)
        ax.plot([q1, q3], [row - 0.15, row - 0.15], color=BLUE, lw=3)
    lo, med, hi = np.quantile(MU, [0.05, 0.5, 0.95])
    ax.fill([lo, med, hi, med], [1, 1.25, 1, 0.75], color=BLUE)
    plo, phi = np.quantile(NEW, [0.05, 0.95])
    ax.plot([plo, phi], [0, 0], color=ORANGE, lw=2.5)
    ax.axvline(0, color=GREY, lw=1, ls="--")
    ax.set(yticks=np.r_[yy, 1, 0], xlabel="log odds ratio (BCG vs control); left = vaccine better", title=title)
    ax.set_yticklabels(list(np.asarray(names)[order]) + ["pooled $\\mu$ (90%)", "next trial (90% prediction)"], fontsize=9)
    ax.set_xticks(np.log([0.1, 0.2, 0.5, 1, 2]))
    ax.set_xticklabels(["OR 0.1", "0.2", "0.5", "1", "2"])
    return ax


fig, ax = plt.subplots(figsize=(9, 6.5))
forest(ax, idata_1, coords["trial"], y, se, "one_stage::", order=order[::-1],
       title="BCG trials, high latitude at the top: observed (black, area = precision) and posterior (blue, 50/90%)")
ax.text(0.98, 0.98, "black: trial estimate with 95% CI\nblue: shrunken posterior\ndiamond: pooled mean\norange: prediction for a new trial",
        transform=ax.transAxes, fontsize=8, va="top", ha="right");

lat_c = (bcg.ablat - bcg.ablat.mean()).to_numpy(float)
metareg = one_stage_model(a, a + b, c, c + d, coords["trial"],
                          extra=lambda m: pm.Normal("beta", 0, 0.1) * lat_c)
idata_r = fit(metareg)
BETA = stacked(idata_r.posterior["one_stage::beta"])
MU_r = stacked(idata_r.posterior["one_stage::mu"])
TAU_r = stacked(idata_r.posterior["one_stage::tau"])
print(f"beta: {np.median(BETA):.3f} per degree of latitude, 90% [{np.quantile(BETA, 0.05):.3f}, {np.quantile(BETA, 0.95):.3f}]")
print(f"tau: {np.median(TAU):.2f} -> {np.median(TAU_r):.2f} with latitude in the model")

lat_grid = np.linspace(10, 60, 100)
line = MU_r[:, None] + BETA[:, None] * (lat_grid - bcg.ablat.mean())
new = line + TAU_r[:, None] * rng.normal(size=(len(MU_r), 1))
fig, axes = plt.subplots(1, 2, figsize=(13, 4.5), gridspec_kw={"width_ratios": [2, 1]})
ax = axes[0]
ax.fill_between(lat_grid, *np.quantile(new, [0.05, 0.95], axis=0), color=ORANGE, alpha=0.2, lw=0, label="90% prediction, new trial")
ax.fill_between(lat_grid, *np.quantile(line, [0.05, 0.95], axis=0), color=BLUE, alpha=0.3, lw=0, label="90% interval of the line")
ax.plot(lat_grid, np.median(line, axis=0), color=BLUE, lw=2)
ax.scatter(bcg.ablat, y, s=600 * w / w.max() + 10, facecolor="none", edgecolor="k", lw=1.2, label="trial (area = precision)")
ax.axhline(0, color=GREY, ls="--", lw=1)
ax.set(xlabel="absolute latitude (degrees)", ylabel="log odds ratio", title="Vaccine effect against latitude")
ax.legend(fontsize=9, loc="lower left")
axes[1].hist(TAU, bins=50, color=GREY, alpha=0.7, density=True, label="without latitude")
axes[1].hist(TAU_r, bins=50, color=BLUE, alpha=0.7, density=True, label="with latitude")
axes[1].set(xlabel="residual $\\tau$", title="Latitude absorbs most of the heterogeneity")
axes[1].legend(fontsize=9);

def mu_summary(idata):
    m = stacked(idata.posterior["one_stage::mu"])
    return np.quantile(m, [0.05, 0.5, 0.95]), (m < 0).mean()


loo_rows = []
for k in range(K):
    keep = np.arange(K) != k
    m_k = one_stage_model(a[keep], (a + b)[keep], c[keep], (c + d)[keep], coords["trial"][keep])
    with m_k:
        id_k = pm.sample(target_accept=0.95, random_seed=RANDOM_SEED, progressbar=False)
    loo_rows.append((bcg.author[k], *mu_summary(id_k)))

by_year = np.argsort(bcg.year.to_numpy(), kind="stable")
cum_rows = []
for n_trials in range(2, K + 1):
    keep = by_year[:n_trials]
    m_k = one_stage_model(a[keep], (a + b)[keep], c[keep], (c + d)[keep], coords["trial"][keep])
    with m_k:
        id_k = pm.sample(target_accept=0.95, random_seed=RANDOM_SEED, progressbar=False)
    cum_rows.append((bcg.year.to_numpy()[keep[-1]], bcg.author[keep[-1]], *mu_summary(id_k)))

full_lo, full_med, full_hi = np.quantile(MU, [0.05, 0.5, 0.95])
fig, axes = plt.subplots(1, 2, figsize=(14, 5.5))
ax = axes[0]
for row, (name, (lo, med, hi), p_neg) in enumerate(loo_rows[::-1]):
    ax.plot([lo, hi], [row, row], color=BLUE, lw=2)
    ax.scatter(med, row, color="k", zorder=3, s=20)
ax.axvspan(full_lo, full_hi, color=GREY, alpha=0.2, lw=0)
ax.axvline(full_med, color=GREY, lw=1)
ax.set(yticks=range(K), xlabel="pooled $\\mu$ with this trial removed (90%)", title="Leave-one-out influence")
ax.set_yticklabels([f"without {r[0]}" for r in loo_rows[::-1]], fontsize=9)

ax = axes[1]
for row, (year, name, (lo, med, hi), p_neg) in enumerate(cum_rows):
    ax.plot([lo, hi], [row, row], color=BLUE, lw=2)
    ax.scatter(med, row, color="k", zorder=3, s=20)
    ax.text(hi + 0.05, row, f"P($\\mu$<0) = {p_neg:.2f}", va="center", fontsize=8)
ax.axvline(0, color=GREY, lw=1, ls="--")
ax.set(yticks=range(len(cum_rows)), xlabel="pooled $\\mu$ after each trial was published (90%)",
       title="Cumulative meta-analysis by year", xlim=(-2.2, 1.2))
ax.set_yticklabels([f"{r[0]}: + {r[1]}" for r in cum_rows], fontsize=9)
ax.invert_yaxis();

data.describe("magnesium")
mg = data.load("magnesium")
a2, n1, c2, n2 = (mg[col].to_numpy(float) for col in ["ai", "n1i", "ci", "n2i"])
cc = 0.5 * ((a2 == 0) | (c2 == 0))               # continuity correction only where a cell is zero
y2 = np.log(((a2 + cc) / (n1 - a2 + cc)) / ((c2 + cc) / (n2 - c2 + cc)))
se2 = np.sqrt(1 / (a2 + cc) + 1 / (n1 - a2 + cc) + 1 / (c2 + cc) + 1 / (n2 - c2 + cc))
mg_names = mg.study.to_numpy(dtype=object)


def normal_re(y_, se_, names, slope=None, name="re"):
    with pm.Model(coords={"trial": names}, name=name) as m:
        mu = pm.Normal("mu", 0, 1.5)
        tau = pm.HalfNormal("tau", 1)
        z = pm.Normal("z", 0, 1, dims="trial")
        shift = pm.Normal("gamma", 0, 2) * slope if slope is not None else 0.0
        pm.Normal("y", mu + shift + tau * z, se_, observed=y_, dims="trial")
    return m


id_all = fit(normal_re(y2, se2, mg_names, name="all16"))
id_small = fit(normal_re(y2[:-1], se2[:-1], mg_names[:-1], name="without_ISIS4"))
id_egger = fit(normal_re(y2, se2, mg_names, slope=se2, name="egger"))
GAMMA = stacked(id_egger.posterior["egger::gamma"])
print(f"small-study slope gamma: median {np.median(GAMMA):.2f}, P(gamma < 0) = {(GAMMA < 0).mean():.3f}")

THR = 10
suppressed_t, suppressed_c = a < THR, c < THR
print(f"cells published as '<{THR}': {suppressed_t.sum()} BCG-arm and {suppressed_c.sum()} control-arm counts "
      f"({', '.join(bcg.author[suppressed_t | suppressed_c])})")

with pm.Model(coords=coords, name="suppressed") as suppressed:
    mu = pm.Normal("mu", 0, 1.5)
    tau = pm.HalfNormal("tau", 1)
    z = pm.Normal("z", 0, 1, dims="trial")
    theta = pm.Deterministic("theta", mu + tau * z, dims="trial")
    alpha = pm.Normal("alpha", -4, 2, dims="trial")
    # lower=THR-1 where suppressed: an observed value of THR-1 then means "at most THR-1"
    pm.Censored("cases_c", pm.Binomial.dist(n=c + d, p=pm.math.invlogit(alpha)),
                lower=np.where(suppressed_c, THR - 1, -1), upper=None,
                observed=np.where(suppressed_c, THR - 1, c), dims="trial")
    pm.Censored("cases_t", pm.Binomial.dist(n=a + b, p=pm.math.invlogit(alpha + theta)),
                lower=np.where(suppressed_t, THR - 1, -1), upper=None,
                observed=np.where(suppressed_t, THR - 1, a), dims="trial")
    pm.Deterministic("theta_new", mu + tau * pm.Normal("z_new", 0, 1))
idata_s = fit(suppressed)

TH_full = stacked(idata_1.posterior["one_stage::theta"])
TH_sup = stacked(idata_s.posterior["suppressed::theta"])
fig, ax = plt.subplots(figsize=(9, 5.5))
yy = np.arange(K)[::-1]
for row, i in zip(yy, order):
    for S, off, col in [(TH_full, 0.18, BLUE), (TH_sup, -0.18, ORANGE)]:
        lo, q1, q3, hi = np.quantile(S[:, i], [0.05, 0.25, 0.75, 0.95])
        ax.plot([lo, hi], [row + off, row + off], color=col, lw=1)
        ax.plot([q1, q3], [row + off, row + off], color=col, lw=3)
    if suppressed_t[i] or suppressed_c[i]:
        ax.text(1.05, row, f"'<{THR}'", color=ORANGE, fontsize=9, va="center")
ax.axvline(0, color=GREY, lw=1, ls="--")
ax.set(yticks=yy, xlabel="log odds ratio", xlim=(-2.6, 1.3),
       title="Trial posteriors from the full counts (blue) and with small cells suppressed (orange)")
ax.set_yticklabels(list(coords["trial"][order]), fontsize=9)
for pre, idata, lab in [("one_stage::", idata_1, "full counts"), ("suppressed::", idata_s, "suppressed")]:
    m = stacked(idata.posterior[pre + "mu"])
    print(f"{lab:12s} mu = {np.median(m):.2f} [{np.quantile(m, 0.05):.2f}, {np.quantile(m, 0.95):.2f}]   "
          f"tau = {np.median(stacked(idata.posterior[pre + 'tau'])):.2f}")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
ax = axes[0]
x = np.arange(len(epsilons))
ax.axhspan(full_lo, full_hi, color=GREY, alpha=0.2, lw=0, label="no privacy: full counts (90%)")
ax.axhline(full_med, color=GREY, lw=1)
for key, col, off, lab in [("aware", BLUE, -0.1, "noise-aware model"), ("naive", ORANGE, 0.1, "naive model (ignores the noise)")]:
    lo, med, hi = np.array([r[key] for r in dp_results]).T
    ax.errorbar(x + off, med, yerr=[med - lo, hi - med], fmt="o", color=col, capsize=3, label=lab)
ax.axhline(0, color="k", lw=0.8, ls=":")
ax.set(xticks=x, xticklabels=[f"$\\varepsilon$ = {e}\n$\\sigma$ = {s:.0f}" for e, s in zip(epsilons, sigmas)],
       ylabel="pooled $\\mu$ (90%)")
ax.set_title("Privacy-utility: stronger privacy to the right", fontsize=11)
ax.legend(fontsize=9, loc="upper left")

ax = axes[1]
n_cases = a + c                                   # what the noise competes with: the number of events
cmap = plt.get_cmap("Blues")
label_these = np.r_[np.argsort(n_cases)[:2], np.argsort(n_cases)[-3:]]
for i in np.argsort(n_cases):
    sd_path = [r["aware_theta_sd"][i] for r in dp_results]
    ax.plot(x, sd_path, marker="o", ms=4, lw=1.5, color=cmap(0.3 + 0.7 * np.log(n_cases[i]) / np.log(n_cases.max())))
    if i in label_these:
        ax.text(x[-1] + 0.1, sd_path[-1], f"{coords['trial'][i]}: {int(n_cases[i])} cases", fontsize=8, va="center")
ax.set(xticks=x, xticklabels=[f"$\\varepsilon$ = {e}" for e in epsilons], ylabel="posterior sd of the trial effect $\\theta_k$",
       xlim=(-0.3, len(x) + 1.6))
ax.set_title("Which trials still carry information (darker = more cases)", fontsize=11);

data.describe("konstantopoulos")
ko = data.load("konstantopoulos")
districts = np.sort(ko.district.unique())
d_idx = pd.Index(districts).get_indexer(ko.district)

with pm.Model(coords={"district": districts.astype(str), "study": ko.study.to_numpy()}, name="three_level") as three_level:
    mu = pm.Normal("mu", 0, 1)
    tau_d = pm.HalfNormal("tau_district", 0.5)
    tau_s = pm.HalfNormal("tau_study", 0.5)
    delta = pm.Deterministic("delta", mu + tau_d * pm.Normal("zd", 0, 1, dims="district"), dims="district")
    theta = pm.Deterministic("theta", delta[d_idx] + tau_s * pm.Normal("zs", 0, 1, dims="study"), dims="study")
    pm.Normal("y", theta, np.sqrt(ko.vi.to_numpy()), observed=ko.yi.to_numpy(), dims="study")
    pm.Deterministic("icc", tau_d**2 / (tau_d**2 + tau_s**2))
idata_3 = fit(three_level)
az.summary(idata_3, var_names=["three_level::mu", "three_level::tau_district", "three_level::tau_study", "three_level::icc"],
           ci_kind="hdi", ci_prob=0.9, round_to=2)

DELTA_ = stacked(idata_3.posterior["three_level::delta"])
TH3 = stacked(idata_3.posterior["three_level::theta"])
ICC = stacked(idata_3.posterior["three_level::icc"])
cmap = plt.get_cmap("tab20")
fig, axes = plt.subplots(1, 2, figsize=(14, 6), gridspec_kw={"width_ratios": [2.2, 1]})
ax = axes[0]
row = 0
for j, dname in enumerate(districts):
    members = np.where(d_idx == j)[0]
    col = cmap(j / len(districts))
    lo, hi = np.quantile(DELTA_[:, j], [0.05, 0.95])
    ax.fill_betweenx([row - 0.5, row + len(members) - 0.5], lo, hi, color=col, alpha=0.2, lw=0)
    ax.axvline(0, color=GREY, lw=1, ls="--")
    for s in members:
        yi, sei = ko.yi[s], np.sqrt(ko.vi[s])
        ax.plot([yi - 1.96 * sei, yi + 1.96 * sei], [row, row], color="k", lw=0.8, alpha=0.6)
        ax.plot(np.quantile(TH3[:, s], [0.05, 0.95]), [row, row], color=col, lw=2.5)
        row += 1
    ax.text(1.25, row - len(members) / 2 - 0.5, f"district {dname}", va="center", fontsize=8, color=col)
ax.set(yticks=[], xlim=(-1.2, 1.6), xlabel="standardised mean difference")
ax.set_title("56 studies in 11 districts: observed 95% CI (black), posterior 90% (colour), district mean (shaded)", fontsize=10)
ax.invert_yaxis()
axes[1].hist(ICC, bins=50, color=BLUE, alpha=0.8, density=True)
axes[1].set(xlabel="share of heterogeneity between districts", xlim=(0, 1))
axes[1].set_title(f"ICC: median {np.median(ICC):.2f}, 90% [{np.quantile(ICC, 0.05):.2f}, {np.quantile(ICC, 0.95):.2f}]", fontsize=11);