<jupyter_start><jupyter_text>Grid algorithms for hierarchical models[Bayesian Inference with PyMC](https://allendowney.github.io/BayesianInferencePyMC)Copyright 2021 <NAME>License: [Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/) It is widely believed that grid algorithms are only practical for models with 1-3 parameters, or maybe 4-5 if you are careful.[I've said so myself](https://allendowney.github.io/ThinkBayes2/chap19.html).But recently I used a grid algorithm to solve the [emitter-detector problem](https://www.allendowney.com/blog/2021/09/05/emitter-detector-redux/), and along the way I noticed something about the structure of the problem: although the model has two parameters, the data only depend on one of them.That makes it possible to evaluate the likelihood function and update the model very efficiently.Many hierarchical models have a similar structure: the data depend on a small number of parameters, which depend on a small number of hyperparameters.I wondered whether the same method would generalize to more complex models, and it does.As an example, in this notebook I'll use a logitnormal-binomial hierarchical model to solve a problem with two hyperparameters and 13 parameters.The grid algorithm is not just practical; it's substantially faster than MCMC.<jupyter_code># If we're running on Colab, install libraries
import sys
IN_COLAB = 'google.colab' in sys.modules

if IN_COLAB:
    !pip install pymc3
    !pip install arviz
    !pip install empiricaldist
# PyMC generates a FutureWarning we don't need to deal with

import warnings
warnings.filterwarnings("ignore", category=FutureWarning)<jupyter_output><empty_output><jupyter_text>The following are some utility functions I'll use.<jupyter_code>import matplotlib.pyplot as plt

def legend(**options):
    """Make a legend only if there are labels."""
    handles, labels = plt.gca().get_legend_handles_labels()
    if len(labels):
        plt.legend(**options)
def decorate(**options):
    plt.gca().set(**options)
    legend()
    plt.tight_layout()
from empiricaldist import Cdf

def compare_cdf(pmf, sample):
    pmf.make_cdf().step(label='grid')
    Cdf.from_seq(sample).plot(label='mcmc')
    print(pmf.mean(), sample.mean())
    decorate()
from empiricaldist import Pmf

def make_pmf(ps, qs, name):
    pmf = Pmf(ps, qs)
    pmf.normalize()
    pmf.index.name = name
    return pmf<jupyter_output><empty_output><jupyter_text>Heart Attack DataThe problem I'll solve is based on [Chapter 10 of *Probability and Bayesian Modeling*](https://bayesball.github.io/BOOK/bayesian-hierarchical-modeling.htmlexample-deaths-after-heart-attack); it uses data on death rates due to heart attack for patients treated at various hospitals in New York City.We can use Pandas to read the data into a `DataFrame`.<jupyter_code>import os

filename = 'DeathHeartAttackManhattan.csv'
if not os.path.exists(filename):
    !wget https://github.com/AllenDowney/BayesianInferencePyMC/raw/main/DeathHeartAttackManhattan.csv
import pandas as pd

df = pd.read_csv(filename)
df<jupyter_output><empty_output><jupyter_text>The columns we need are `Cases`, which is the number of patients treated at each hospital, and `Deaths`, which is the number of those patients who died.<jupyter_code>data_ns = df['Cases'].values
data_ks = df['Deaths'].values<jupyter_output><empty_output><jupyter_text>Solution with PyMCHere's a hierarchical model that estimates the death rate for each hospital and simultaneously estimates the distribution of rates across hospitals.<jupyter_code>import pymc3 as pm

def make_model():
    with pm.Model() as model:
        mu = pm.Normal('mu', 0, 2)
        sigma = pm.HalfNormal('sigma', sigma=1)
        xs = pm.LogitNormal('xs', mu=mu, sigma=sigma, shape=len(data_ns))
        ks = pm.Binomial('ks', n=data_ns, p=xs, observed=data_ks)
    return model
%time model = make_model()
pm.model_to_graphviz(model)
with model:
    pred = pm.sample_prior_predictive(1000)
    %time trace = pm.sample(500, target_accept=0.97)<jupyter_output>Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [xs, sigma, mu]<jupyter_text>To be fair, PyMC doesn't like this parameterization much (although I'm not sure why). One most runs, there are a moderate number of divergences. Even so, the results are good enough. Here are the posterior distributions of the hyperparameters.<jupyter_code>import arviz as az

with model:
    az.plot_posterior(trace, var_names=['mu', 'sigma'])<jupyter_output><empty_output><jupyter_text>And we can extract the posterior distributions of the xs.<jupyter_code>trace_xs = trace['xs'].transpose()
trace_xs.shape<jupyter_output><empty_output><jupyter_text>As an example, here's the posterior distribution of x for the first hospital.<jupyter_code>with model:
    az.plot_posterior(trace_xs[0])<jupyter_output><empty_output><jupyter_text>The grid priorsNow let's solve the same problem using a grid algorithm.I'll use the same priors for the hyperparameters, approximated by a grid with about 100 elements in each dimension.<jupyter_code>import numpy as np
from scipy.stats import norm

mus = np.linspace(-6, 6, 101)
ps = norm.pdf(mus, 0, 2)
prior_mu = make_pmf(ps, mus, 'mu')

prior_mu.plot()
decorate(title='Prior distribution of mu')
from scipy.stats import logistic

sigmas = np.linspace(0.03, 3.6, 90)
ps = norm.pdf(sigmas, 0, 1)
prior_sigma = make_pmf(ps, sigmas, 'sigma')

prior_sigma.plot()
decorate(title='Prior distribution of sigma')<jupyter_output><empty_output><jupyter_text>The following cells confirm that these priors are consistent with the prior samples from PyMC.<jupyter_code>compare_cdf(prior_mu, pred['mu'])
decorate(title='Prior distribution of mu')
compare_cdf(prior_sigma, pred['sigma'])
decorate(title='Prior distribution of sigma')<jupyter_output>0.8033718951689776 0.8244605687886865<jupyter_text>The joint distribution of hyperparametersI'll use `make_joint` to make an array that represents the joint prior distribution of the hyperparameters.<jupyter_code>def make_joint(prior_x, prior_y):
    X, Y = np.meshgrid(prior_x.ps, prior_y.ps, indexing='ij')
    hyper = X * Y
    return hyper
prior_hyper = make_joint(prior_mu, prior_sigma)
prior_hyper.shape<jupyter_output><empty_output><jupyter_text>Here's what it looks like.<jupyter_code>import pandas as pd
from utils import plot_contour

plot_contour(pd.DataFrame(prior_hyper, index=mus, columns=sigmas))
decorate(title="Joint prior of mu and sigma")<jupyter_output><empty_output><jupyter_text>Joint prior of hyperparameters and xNow we're ready to lay out the grid for x, which is the proportion we'll estimate for each hospital.<jupyter_code>xs = np.linspace(0.01, 0.99, 295)<jupyter_output><empty_output><jupyter_text>For each pair of hyperparameters, we'll compute the distribution of `x`.<jupyter_code>from scipy.special import logit

M, S, X = np.meshgrid(mus, sigmas, xs, indexing='ij')
LO = logit(X)
LO.sum()
from scipy.stats import norm

%time normpdf = norm.pdf(LO, M, S)
normpdf.sum()<jupyter_output>CPU times: user 69.6 ms, sys: 16.5 ms, total: 86.1 ms
Wall time: 84.9 ms<jupyter_text>We can speed this up by computing skipping the terms that don't depend on x<jupyter_code>%%time

z = (LO-M) / S
normpdf = np.exp(-z**2/2)<jupyter_output>CPU times: user 26 ms, sys: 10.6 ms, total: 36.6 ms
Wall time: 35.1 ms<jupyter_text>The result is a 3-D array with axes for mu, sigma, and x.Now we need to normalize each distribution of `x`.<jupyter_code>totals = normpdf.sum(axis=2)
totals.shape<jupyter_output><empty_output><jupyter_text>To normalize, we have to use a safe version of `divide` where `0/0` is `0`.<jupyter_code>def divide(x, y):
    out = np.zeros_like(x)
    return np.divide(x, y, out=out, where=(y!=0))    
shape = totals.shape + (1,)
normpdf = divide(normpdf, totals.reshape(shape))
normpdf.shape<jupyter_output><empty_output><jupyter_text>The result is an array that contains the distribution of `x` for each pair of hyperparameters.Now, to get the prior distribution, we multiply through by the joint distribution of the hyperparameters.<jupyter_code>def make_prior(hyper):

    # reshape hyper so we can multiply along axis 0
    shape = hyper.shape + (1,)
    prior = normpdf * hyper.reshape(shape)

    return prior
%time prior = make_prior(prior_hyper)
prior.sum()<jupyter_output>CPU times: user 5.57 ms, sys: 0 ns, total: 5.57 ms
Wall time: 4.87 ms<jupyter_text>The result is a 3-D array that represents the joint prior distribution of `mu`, `sigma`, and `x`.To check that it is correct, I'll extract the marginal distributions and compare them to the priors.<jupyter_code>def marginal(joint, axis):
    axes = [i for i in range(3) if i != axis]
    return joint.sum(axis=tuple(axes))
prior_mu.plot()
marginal_mu = Pmf(marginal(prior, 0), mus)
marginal_mu.plot()
decorate(title='Checking the marginal distribution of mu')
prior_sigma.plot()
marginal_sigma = Pmf(marginal(prior, 1), sigmas)
marginal_sigma.plot()
decorate(title='Checking the marginal distribution of sigma')<jupyter_output><empty_output><jupyter_text>We didn't compute the prior distribution of `x` explicitly; it follows from the distribution of the hyperparameters. But we can extract the prior marginal of `x` from the joint prior.<jupyter_code>marginal_x = Pmf(marginal(prior, 2), xs)
marginal_x.plot()
decorate(title='Checking the marginal distribution of x',
         ylim=[0, np.max(marginal_x) * 1.05])<jupyter_output><empty_output><jupyter_text>And compare it to the prior sample from PyMC.<jupyter_code>pred_xs = pred['xs'].transpose()
pred_xs.shape
compare_cdf(marginal_x, pred_xs[0])
decorate(title='Prior distribution of x')<jupyter_output>0.49996889063901967 0.4879934000104224<jupyter_text>The prior distribution of `x` I get from the grid is a bit different from what I get from PyMC. I'm not sure why, but it doesn't seem to affect the results much.In addition to the marginals, we'll also find it useful to extract the joint marginal distribution of the hyperparameters.<jupyter_code>def get_hyper(joint):
    return joint.sum(axis=2)
hyper = get_hyper(prior)
plot_contour(pd.DataFrame(hyper, 
                          index=mus, 
                          columns=sigmas))
decorate(title="Joint prior of mu and sigma")<jupyter_output><empty_output><jupyter_text>The UpdateThe likelihood of the data only depends on `x`, so we can compute it like this.<jupyter_code>from scipy.stats import binom

data_k = data_ks[0]
data_n = data_ns[0]

like_x = binom.pmf(data_k, data_n, xs)
like_x.shape
plt.plot(xs, like_x)
decorate(title='Likelihood of the data')<jupyter_output><empty_output><jupyter_text>And here's the update.<jupyter_code>def update(prior, data):
    n, k = data
    like_x = binom.pmf(k, n, xs)
    posterior = prior * like_x
    posterior /= posterior.sum()
    return posterior
data = data_n, data_k
%time posterior = update(prior, data)<jupyter_output>CPU times: user 11.6 ms, sys: 11.9 ms, total: 23.5 ms
Wall time: 7.66 ms<jupyter_text>Serial updatesAt this point we can do an update based on a single hospital, but how do we update based on all of the hospitals?As a step toward the right answer, I'll start with a wrong answer, which is to do the updates one at a time.After each update, we extract the posterior distribution of the hyperparameters and use it to create the prior for the next update.At the end, the posterior distribution of hyperparameters is correct, and the marginal posterior of `x` for the *last* hospital is correct, but the other marginals are wrong because they do not take into account data from subsequent hospitals.<jupyter_code>def multiple_updates(prior, ns, ks):
    for data in zip(ns, ks):
        print(data)
        posterior = update(prior, data)
        hyper = get_hyper(posterior)
        prior = make_prior(hyper)
    return posterior
%time posterior = multiple_updates(prior, data_ns, data_ks)<jupyter_output>(129, 4)
(35, 1)
(228, 18)
(84, 7)
(291, 24)
(270, 16)
(46, 6)
(293, 19)
(241, 15)
(105, 13)
(353, 25)
(250, 11)
(41, 4)
CPU times: user 185 ms, sys: 35.4 ms, total: 220 ms
Wall time: 172 ms<jupyter_text>Here are the posterior distributions of the hyperparameters, compared to the results from PyMC.<jupyter_code>marginal_mu = Pmf(marginal(posterior, 0), mus)
compare_cdf(marginal_mu, trace['mu'])
marginal_sigma = Pmf(marginal(posterior, 1), sigmas)
compare_cdf(marginal_sigma, trace['sigma'])
marginal_x = Pmf(marginal(posterior, 2), xs)
compare_cdf(marginal_x, trace_xs[-1])<jupyter_output>0.07330826956150183 0.07297933578329886<jupyter_text>Parallel updatesDoing updates one at time is not quite right, but it gives us an insight.Suppose we start with a uniform distribution for the hyperparameters and do an update with data from one hospital. If we extract the posterior joint distribution of the hyperparameters, what we get is the likelihood function associated with one dataset.The following function computes these likelihood functions and saves them in an array called `hyper_likelihood`.<jupyter_code>def compute_hyper_likelihood(ns, ks):
    shape = ns.shape + mus.shape + sigmas.shape
    hyper_likelihood = np.empty(shape)
    
    for i, data in enumerate(zip(ns, ks)):
        print(data)
        n, k = data
        like_x = binom.pmf(k, n, xs)
        posterior = normpdf * like_x
        hyper_likelihood[i] = get_hyper(posterior)
    return hyper_likelihood
%time hyper_likelihood = compute_hyper_likelihood(data_ns, data_ks)<jupyter_output>(129, 4)
(35, 1)
(228, 18)
(84, 7)
(291, 24)
(270, 16)
(46, 6)
(293, 19)
(241, 15)
(105, 13)
(353, 25)
(250, 11)
(41, 4)
CPU times: user 82 ms, sys: 55.2 ms, total: 137 ms
Wall time: 75.5 ms<jupyter_text>We can multiply this out to get the product of the likelihoods.<jupyter_code>%time hyper_likelihood_all = hyper_likelihood.prod(axis=0)
hyper_likelihood_all.sum()<jupyter_output>CPU times: user 279 µs, sys: 0 ns, total: 279 µs
Wall time: 158 µs<jupyter_text>This is useful because it provides an efficient way to compute the marginal posterior distribution of `x` for any hospital.Here's an example.<jupyter_code>i = 3
data = data_ns[i], data_ks[i]
data<jupyter_output><empty_output><jupyter_text>Suppose we did the updates serially and saved this hospital for last.The prior distribution for the final update would reflect the updates from all previous hospitals, which we can compute by dividing out `hyper_likelihood[i]`.<jupyter_code>%time hyper_i = divide(prior_hyper * hyper_likelihood_all, hyper_likelihood[i])
hyper_i.sum()<jupyter_output>CPU times: user 310 µs, sys: 147 µs, total: 457 µs
Wall time: 342 µs<jupyter_text>We can use `hyper_i` to make the prior for the last update.<jupyter_code>prior_i = make_prior(hyper_i)<jupyter_output><empty_output><jupyter_text>And then do the update.<jupyter_code>posterior_i = update(prior_i, data)<jupyter_output><empty_output><jupyter_text>And we can confirm that the results are similar to the results from PyMC.<jupyter_code>marginal_mu = Pmf(marginal(posterior_i, 0), mus)
marginal_sigma = Pmf(marginal(posterior_i, 1), sigmas)
marginal_x = Pmf(marginal(posterior_i, 2), xs)
compare_cdf(marginal_mu, trace['mu'])
compare_cdf(marginal_sigma, trace['sigma'])
compare_cdf(marginal_x, trace_xs[i])<jupyter_output>0.07245354421667904 0.07224440565018131<jupyter_text>Compute all marginalsThe following function computes the marginals for all hospitals and stores the results in an array.<jupyter_code>def compute_all_marginals(ns, ks):
    shape = len(ns), len(xs)
    marginal_xs = np.zeros(shape)
    numerator = prior_hyper * hyper_likelihood_all
    
    for i, data in enumerate(zip(ns, ks)):
        hyper_i = divide(numerator, hyper_likelihood[i])
        prior_i = make_prior(hyper_i) 
        posterior_i = update(prior_i, data)
        marginal_xs[i] = marginal(posterior_i, 2)
        
    return marginal_xs
%time marginal_xs = compute_all_marginals(data_ns, data_ks)<jupyter_output>CPU times: user 184 ms, sys: 49.8 ms, total: 234 ms
Wall time: 173 ms<jupyter_text>Here's what the results look like, compared to the results from PyMC.<jupyter_code>for i, ps in enumerate(marginal_xs):
    pmf = Pmf(ps, xs)
    plt.figure()
    compare_cdf(pmf, trace_xs[i])
    decorate(title=f'Posterior marginal of x for Hospital {i}',
             xlabel='Death rate',
             ylabel='CDF',
             xlim=[trace_xs[i].min(), trace_xs[i].max()])<jupyter_output>0.06123636407822421 0.0617519291444324
0.06653003152551518 0.06643868288267936
0.07267383211481376 0.07250041300148316
0.07245354421667904 0.07224440565018131
0.07430385699796423 0.07433369435815212
0.06606326919655045 0.06646020352443961
0.07774639529896528 0.07776805141855801
0.06788483681522386 0.06807113157490664
0.06723306224279789 0.06735326167909643
0.08183332535205982 0.08115900598539395
0.07003760661997555 0.0704088595242495
0.06136130741477605 0.06159674913422137
0.07330826956150185 0.07297933578329886<jupyter_text>And here are the percentage differences between the results from the grid algorithm and PyMC. Most of them are less than 1%.<jupyter_code>for i, ps in enumerate(marginal_xs):
    pmf = Pmf(ps, xs)
    diff = abs(pmf.mean() - trace_xs[i].mean()) / pmf.mean()
    print(diff * 100)<jupyter_output>0.841926319383687
0.13730437329010417
0.23862662568368032
0.28865194761527047
0.04015586995533174
0.6008396688759207
0.027854821447936134
0.274427646029194
0.17878024931315142
0.8240155997142278
0.5300765148763152
0.38369736461746806
0.44869941709241024
{'fraction_non_alphanumeric': 0.09696693791612132, 'fraction_numerical': 0.06612416775735504, 'mean_word_length': 5.376632801161103, 'pattern_counts': {'":': 0, '<': 130, '<?xml version=': 0, '>': 130, 'https://': 6, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '36', 'dir': 'jupyter-structured-clean-dedup', 'id': '471743', 'license': 'MIT', 'n_tokens_mistral': 6789, 'n_tokens_neox': 5896, 'n_words': 1806, 'path': 'examples/hospital.ipynb', 'repo_name': 'EricRobertCampbell/ThinkBayes2', 'size': '17582'}