Introduction to Bayesian Modeling with PyMC3

Bayes Theorem

Prepare Notebook

import numpy as np
import pandas as pd
import scipy.special as sp
import scipy.stats as ss
import pymc3 as pm
import arviz as az
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style(
    style='darkgrid', 
    rc={'axes.facecolor': '.9', 'grid.color': '.8'}
)
sns.set_palette(palette='deep')
sns_c = sns.color_palette(palette='deep')

plt.rcParams['figure.figsize'] = [10, 7]

Generate Sample Data

Let us generate some sample data (for we of course know the true parameter \(\lambda\)):

# We set a seed so that the results are reproducible.
SEED = 5
np.random.seed(SEED)
# number of samples.
n = 100
# true parameter.
lam_true = 2
# sample array.
y = np.random.poisson(lam=lam_true, size=n)
y

array([2, 4, 1, 0, 2, 2, 2, 2, 1, 1, 3, 2, 0, 1, 3, 3, 4, 2, 0, 0, 3, 6,
       1, 2, 1, 2, 5, 2, 3, 0, 1, 3, 1, 4, 1, 2, 4, 0, 6, 4, 1, 2, 2, 0,
       1, 2, 4, 4, 1, 3, 0, 3, 3, 2, 4, 2, 2, 1, 1, 2, 5, 2, 3, 0, 1, 1,
       1, 3, 4, 1, 3, 4, 2, 1, 2, 4, 2, 2, 1, 0, 2, 2, 3, 0, 3, 3, 4, 2,
       2, 1, 2, 1, 3, 0, 1, 0, 3, 3, 1, 2])

Let us plot the sample distribution:

# Histogram of the sample.
fig, ax = plt.subplots()
sns.countplot(x=y, color=sns_c[0], label=f'mean = {y.mean(): 0.3f}', ax=ax)
ax.legend(loc='upper right')
ax.set(title='Sample Distribution', xlabel='y');

Prior: Gamma Distribution

Let us consider a gamma prior distribution for the parameter \(\lambda \sim \Gamma(a,b)\). Recall that the density function for the gamma distribution is

\[ f(\lambda)=\frac{b^a}{\Gamma(a)}\lambda^{a-1} e^{-b\lambda} \]

where \(a>0\) is the shape parameter and \(b>0\) is the rate parameter.

The expected value and variance of the gamma distribution is

\[ E(\lambda)=\frac{a}{b} \quad \text{and} \quad Var(\lambda)=\frac{a}{b^2} \]

# Parameters of the prior gamma distribution.
a = 3 # shape
b = 1 # rate = 1/scale

x = np.linspace(start=0,stop=10, num=300)

fig, ax = plt.subplots()
ax.plot(x, ss.gamma.pdf(x, a=a, scale=1/b), color=sns_c[3])
ax.axvline(x= a /b, color=sns_c[4], linestyle='--', label='gamma mean')
ax.axvline(x= y.mean(), color=sns_c[1], linestyle='--', label='sample mean')
ax.legend()
ax.set(title=f'Gamma Density Function for a={a} and b={b}');

# Define the prior distribution.
prior = lambda x: ss.gamma.pdf(x, a=a, scale=1/b)

Likelihood

As the observations are independent the likelihood function is

\[ f(y|\lambda)=\prod_{i=1}^{n} \frac{e^{-\lambda}\lambda^{y_i}}{y_i!} =\frac{e^{-n\lambda}\lambda^{\sum_{i=1}^n y_i}}{\prod_{i=1}^{n}y_i!} \]

# Define the likelihood function.
def likelihood(lam, y):
    
    factorials = np.apply_along_axis(
        lambda x: sp.gamma(x + 1),
        axis=0,
        arr=y
    )
    
    numerator = np.exp(- lam * y.size) * (lam ** (y.sum()))
    
    denominator = np.multiply.reduce(factorials)
    
    return numerator / denominator  

Posterior distribution for \(\lambda\) up to a constant

As we are just interested in the structure of the posterior distribution, up to a constant, we see

\[\begin{align} f(\lambda|y)\propto & \: \text{likelihood} \times \text{prior}\\ \propto & \quad f(y|\lambda)f(\lambda)\\ \propto & \quad e^{-n\lambda}\lambda^{\sum_{i=1}^n y_i} \lambda^{a-1} e^{-b\lambda}\\ \propto & \quad \lambda^{\left(\sum_{i=1}^n y_i+a\right)-1} e^{-(n+b)\lambda}\\ \end{align}\]

# Define the posterior distribution.
# (up to a constant)
def posterior_up_to_constant(lam, y):
    return likelihood(lam=lam, y=y) * prior(lam)

Now we plot of the prior and (scaled) posterior distributionfor the parameter lambda. We multiply the posterior distribution function by the amplitude factor \(2.5\times 10^{74}\) to make it visually comparable with the prior gamma distribution.

fig, ax = plt.subplots()
ax.plot(x, ss.gamma.pdf(x, a=a, scale=1/b), color=sns_c[3], label='prior')
ax.plot(
  x, 2.0e74 * posterior_up_to_constant(x,y), color=sns_c[2], label='posterior_up_to_constant'
)
ax.legend()
ax.set(xlabel='$\lambda$');

True posterior distribution for \(\lambda\)

In fact, as \(f(\lambda|y) \propto\: \lambda^{\left(\sum_{i=1}^n y_i+a\right)-1} e^{-(n+b)\lambda}\), one verifies that the posterior distribution is again a gamma

\[\begin{align} f(\lambda|y) = \Gamma\left(\sum_{i=1}^n y_i+a, n+b\right) \end{align}\]

This means that the gamma and Poisson distribution form a conjugate pair.

def posterior(lam, y):
    shape = a + y.sum()
    rate = b + y.size
    return ss.gamma.pdf(lam, shape, scale=1/rate)

fig, ax = plt.subplots()
ax.plot(x, ss.gamma.pdf(x, a=a, scale=1/b), color=sns_c[3], label='prior')
ax.plot(x, posterior(x,y), color=sns_c[2], label='posterior')
ax.legend()
ax.set(title='Prior and Posterior Distributions', xlabel='$\lambda$');

We indeed see how the posterior distribution is concentrated around the true parameter \(\lambda=2\).

Note that the posterior mean is

\[\begin{align} \frac{\sum_{i=1}^n y_i+a}{n+b} = \frac{b}{n+b}\frac{a}{b}+\frac{n}{n+b}\frac{\sum_{i=1}^n y_i}{n} \end{align}\]

That is, it is a weighted average of the prior mean \(a/b\) and the sample average \(\bar{y}\). As \(n\) increases,

\[\begin{align} \lim_{n\rightarrow +\infty}\frac{b}{n+b}\frac{a}{b}+\frac{n}{n+b}\frac{\sum_{i=1}^n y_i}{n} = \bar{y}. \end{align}\]

# Posterior gamma parameters.
shape = a + y.sum()
rate = b + y.size

# Posterior mean.
print(f'Posterior Mean = {shape / rate: 0.3f}')

Posterior Mean =  2.069

Metropolis–Hastings algorithm

Let \(\phi\) be a function that is proportional to the desired probability distribution \(f\).

Initialization:

Pick \(x_{0}\) to be the first sample, and choose an arbitrary probability density

\[\begin{equation} g(x_{n+1}| x_{n}) \end{equation}\]

that suggests a candidate for the next sample value \(x_{n+1}\). Assume \(g\) is symmetric.

For each iteration:

Generate a candidate \(x\) for the next sample by picking from the distribution \(g(x|x_n)\). Calculate the acceptance ratio

\[\begin{equation} \alpha := \frac{f(x)}{f(x_n)} = \frac{\phi(x)}{\phi(x_n)} \end{equation}\]

If \(\alpha \geq 1\), automatically accept the candidate by setting

\[\begin{equation} x_{n+1} = x. \end{equation}\]

Otherwise, accept the candidate with probability \(\alpha\). If the candidate is rejected, set

\[\begin{equation} x_{n+1} = x_{n}. \end{equation}\]

Why does this algorithm solve the initial problem? The full explanation is beyond the scope of this post (some references are provided at the end). It relies in the in the following result.

Ergodic Theorem for Markov Chains

Theorem (Ergodic Theorem for Markov Chains) If \(\{x^{(1)} , x^{(2)} , . . .\}\) is an irreducible, aperiodic and recurrent Markov chain, then there is a unique probability distribution \(\pi\) such that as \(N\longrightarrow\infty\),

• \(P(x^{(N)} ∈ A) \longrightarrow \pi(A)\).
• \(\displaystyle{\frac{1}{N}\sum_{n=1}^{N} g(x^{(n)})) \longrightarrow \int g(x)\pi(x) dx }\).

Recall:

• A Markov chain is said to be irreducible if it is possible to get to any state from any state.

• A state \(n\) has period \(k\) if any return to state \(n\) must occur in multiples of \(k\) time steps.

• If \(k=1\), then the state is said to be aperiodic.

• A state \(n\) is said to be transient if, given that we start in state \(n\), there is a non-zero probability that we will never return to \(i\).

• A state \(n\) is recurrent if it is not transient.

Remark: PyMC3 has many modern samplers which are in general better that MCMC, like Hamiltonian Monte Carlo (HMC) and No-U-Turn Sampler. For more details I recommend checking Statistical Rethinking - Chapter 9 and Michael Betancourt’s blog.

Data

Assume there is a big factory producing chocolate cookies around the world. The cookies follow a unique recipe. You want to study the chocolate chips distribution for cookies produced in \(5\) different locations. Let us read the data into a pandas dataframe.

cookies = pd.read_csv('../data/cookies.dat', sep = ' ')

cookies.head()

	chips	location
0	12	1
1	12	1
2	6	1
3	13	1
4	12	1

In data set we have \(5\) different locations with \(30\) observations each.

cookies.value_counts('location')

location
1    30
2    30
3    30
4    30
5    30
dtype: int64

Let us start by plotting the sample data distribution:

# Histogram distribution of chocolate chips for all cookies.
fig, ax = plt.subplots()
sns.countplot(x='chips', data=cookies, color='black', ax=ax);
ax.set(title='Chips Count(All Locations)');

Now let us plot the distribution per location:

g = sns.catplot(
    x='chips',
    data=cookies,
    col='location',
    hue='location',
    kind='count',
    dodge=False,
    col_wrap=3
)
g.fig.suptitle('Chips Count per Location', y=1.02);

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 8), constrained_layout=True
)
sns.kdeplot(
    x='chips',
    data=cookies,
    hue='location',
    palette=sns_c[:5],ax=ax[0]
)
sns.boxplot(
    x='location',
    y='chips',
    data=cookies,
    hue='location',
    palette=sns_c[:5],
    dodge=False,
    ax=ax[1]
)
ax[1].legend(title='location', loc='center left', bbox_to_anchor=(1, 0.5))
fig.suptitle('Chips Distribution per Location');

In view of these plots and the problem statement, one would consider the following aspects to model the data:

• On the one hand side you would imagine that the distribution across the locations is similar, as they all come from a unique recipe. This is why you may want to model all locations together (pooling).

• On the other hand, in reality, as the locations are not exactly the same you might expect some differences between each location. This is why you may want to model all locations separately (no pooling).

Hence, we would like to have a mix of the two settings described above (partial pooling). Hierarchical models are a great way of doing this. How? Well, we can assume \(5\) different parameters for each location, but all of them coming from a common global (hyper) prior distribution.

The Model: Hierarchical Approach

• Hierarchical Model:

We model the chocolate chip counts by a Poisson distribution with parameter \(\lambda\). Motivated by the example above, we choose a gamma prior.

\[\begin{align} \text{chips} \sim \text{Pois}(\lambda) \quad\quad\quad \lambda \sim \Gamma(\alpha, \beta) \end{align}\]

• Parametrization:

We parametrize the shape and scale of the gamma prior with the mean \(\mu\) and variance \(\sigma^2\).

\[\begin{align} \alpha =\frac{\mu^2}{\sigma^2} \quad\quad\quad \beta=\frac{\mu}{\sigma^2} \end{align}\]

• Prior Distributions:

We further impose prior for these parameters

\[\begin{align} \mu \sim \Gamma(2,1/5) \quad\quad\quad \sigma \sim \text{Exp}(1) \end{align}\]

Let us generate the plots of the prior distributions:

x = np.linspace(start=0, stop=50, num=500)

fig, ax = plt.subplots(
    nrows=1, ncols=2, figsize=(12, 6), constrained_layout=True
)
ax[0].plot(x, ss.gamma.pdf(x, a=2, scale=5), color=sns_c[3])
ax[0].set(
    title=f'Prior Distribution for $\mu$ \n Gamma Density Function for a={2} and b={1/5}'
)
ax[1].plot(x, ss.expon.pdf(x,1), color=sns_c[3])
ax[1].set(
    title='Prior Distribution for sigma \n Exponential Density Function for $\lambda=1$',
    xlim=(1, 10)
);

We will do some prior predictive sampling to understand the choice of these priors.

Let us write the model in PyMC3. Note how the syntax mimics the mathematical formulation.

# Define useful variables and coordinates in order to 
# specify shapes of the distributions.
chips = cookies['chips'].values
location_idx, locations = cookies['location'].factorize(sort=True)

COORDS = {
    'obs': cookies.index,
    'location': locations
}

with pm.Model(coords=COORDS) as cookies_model:
    # Hyperpriors:
    # Prior distribution for mu.
    mu = pm.Gamma('mu', alpha=2.0, beta=1.0/5)
    # Prior distribution for sigma2.
    sigma = pm.Exponential('sigma', 1.0)
    # Parametrization for the shape parameter.
    alpha =  pm.Deterministic('alpha', mu**2 / sigma**2)
    # Parametrization for the scale parameter.
    beta = pm.Deterministic('beta', mu / sigma**2)
    # Prior distribution for lambda.
    lam = pm.Gamma(
        'lam', 
        alpha=alpha, 
        beta=beta, 
        dims='location'
    )
    # Likehood function.
    rate = lam[location_idx]
    likelihood = pm.Poisson(
        'likelihood',
        mu=rate,
        observed=chips,
        dims='obs'
    )
    # Sample prior distribution.
    prior_predictive = pm.sample_prior_predictive()

pm.model_to_graphviz(cookies_model)

Let us what the model predicts using the prior before seeing the data:

fig, ax = plt.subplots(figsize=(12, 8))
sns.countplot(
    x=prior_predictive['likelihood'].flatten(), 
    color=sns_c[5],
    ax=ax
)
ax.set(
    title='Prior Predictive Sampling (all $\lambda$s)', 
    xlim=(None, 50)
);

Now let us pass the data trough the model and sample from the posterior distribution.

with cookies_model:

    cookies_trace = pm.sample(
        tune=1000,
        draws=4000,
        chains=4,
        cores=-1,
        return_inferencedata=True,
        target_accept=0.95
    )

    posterior_predictive = az.from_pymc3(
        posterior_predictive=pm.sample_posterior_predictive(trace=cookies_trace)
    )

Let us plot the traces:

az.plot_trace(cookies_trace);

From the traceplot we see that the chains have converged. Let us compute the mean and hdi for each posterior component.

az.plot_posterior(cookies_trace);

We can also have a detailed summary of the posterior distribution for each parameter:

az.summary(cookies_trace)

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
mu	9.114	0.959	7.469	11.095	0.013	0.009	5614.0	4547.0	1.0
sigma	2.072	0.710	0.969	3.412	0.011	0.008	4570.0	4473.0	1.0
alpha	25.948	17.252	2.789	55.881	0.269	0.232	4600.0	4362.0	1.0
beta	2.857	1.892	0.377	6.206	0.029	0.025	4494.0	4525.0	1.0
lam[0]	9.281	0.540	8.290	10.310	0.006	0.004	7964.0	5892.0	1.0
lam[1]	6.240	0.462	5.415	7.152	0.006	0.004	6221.0	4753.0	1.0
lam[2]	9.532	0.551	8.514	10.574	0.006	0.004	7970.0	6172.0	1.0
lam[3]	8.942	0.520	8.016	9.975	0.006	0.004	7568.0	5857.0	1.0
lam[4]	11.751	0.620	10.602	12.914	0.007	0.005	6835.0	5691.0	1.0

The r_hat are all very close to \(1.0\) which is what we want. There are many other diagnostics available in arviz.

We can also test for correlation between samples in the chains. We are aiming for zero auto-correlation to get “random” samples from the posterior distribution.

# Auto-correlation of the parameter sigma for the 3 chains.
az.plot_autocorr(data=cookies_trace, combined=True, max_lag=10);

From these plots we see that the auto-correlation is not problematic. Moreover, one can get an estimate of the effective sample size of the chains from the columns ess_ in the summary table above, see arviz.ess.

Finally, we can compute the Watanabe–Akaike information (see arviz.waic) and the Pareto-smoothed importance sampling leave-one-out cross-validation (see arviz.loo), which are metrics to estimate out-of-sample performance.

stats_df = pd.concat([
    pd.DataFrame(az.loo(data=cookies_trace)),
    pd.DataFrame(az.waic(data=cookies_trace))
])
stats_df.columns = ['values']

stats_df 

	values
loo	-394.831515
loo_se	10.674036
p_loo	5.756556
n_samples	8000
n_data_points	150
warning	False
loo_scale	log
waic	-394.82533
waic_se	10.673582
p_waic	5.750371
n_samples	8000
n_data_points	150
warning	True
waic_scale	log

Another useful diagnostic for HCM algorithms is an energy plot, see arviz.plot_energy. These algorithms can me seen as a dynamical system of a physical particle where the total energy is conserved (see The Geometric Foundations of Hamiltonian Monte Carlo). In this plot we want the marginal and energy transition plots to overlap as much as possible.

fig, ax = plt.subplots()

az.plot_energy(
    data=cookies_trace,
    fill_color=('C0', 'C3'),
    fill_alpha=(0.5, 0.5),
    ax=ax
)
ax.set(
    title='Energy transition distribution and marginal energy distribution'
);

Posterior Predictive Distribution

Finally, we are going to illustrate how to use the simulation results to derive predictions.

Bayesian Statistics

• Bayesian Analysis with Python (Second edition),

• Statistical Rethinking

• Coursera: Bayesian Statistics: From Concept to Data Analysis

• Coursera: Bayesian Statistics: Techniques and Models

• A First Course in Bayesian Statistical Methods, Peter D. Hoff

• An Introduction to Bayesian Analysis: Theory and Methods, Ghosh, Jayanta K., Delampady, Mohan, Samanta, Tapas

• Michael Betancourt’s blog

• The Geometric Foundations of Hamiltonian Monte Carlo

PyMC3

• Documentation

• Probabilistic Programming in Python using PyMC, John Salvatier, Thomas Wiecki, Christopher Fonnesbeck