In [1]:
%matplotlib inline
import pymc3 as pm
import numpy as np
import theano.tensor as tt
import matplotlib.pylab as plt

np.random.seed(20090425)
In [2]:
# Size of dataset
n = 800
# Number of groups
k = 25

pop_mean = 8
pop_sd = 3
obs_sd = 10

# Simulate group means
m = np.random.normal(loc=pop_mean, scale=pop_sd, size=k)
In [3]:
m[:10]
Out[3]:
array([ 10.11910178,   2.52654191,   4.51891947,  10.61821576,
         4.42735171,   7.75110791,   1.02261171,  12.1946459 ,
         6.09602232,  14.20286786])
In [4]:
group_index = np.random.randint(0, k, size=n)
y = np.random.normal(loc=m[group_index], scale=obs_sd)

Centered model

In [5]:
with pm.Model() as model:
    
    μ = pm.Normal('μ', 0, sd=100)
    σ = pm.HalfCauchy('σ', 5)
    
    θ = pm.Normal('θ', μ, sd=σ, shape=k)
    
    y_obs = pm.Normal('y_obs', θ[group_index], sd=obs_sd, observed=y)
In [6]:
n_iterations = 2000
n_to_keep = 1000
In [7]:
with model:
    
    tr = pm.sample(n_iterations, n_init=50000, njobs=2)
Auto-assigning NUTS sampler...
Initializing NUTS using advi...
Average ELBO = -2,983.9: 100%|██████████| 50000/50000 [00:05<00:00, 8989.05it/s]
Finished [100%]: Average ELBO = -2,983.2
100%|██████████| 2000/2000 [00:10<00:00, 183.96it/s]
In [8]:
pm.traceplot(tr[-n_to_keep:], varnames=['μ', 'σ'])
Out[8]:
array([[<matplotlib.axes._subplots.AxesSubplot object at 0x12153b7f0>,
        <matplotlib.axes._subplots.AxesSubplot object at 0x12150de48>],
       [<matplotlib.axes._subplots.AxesSubplot object at 0x121518e10>,
        <matplotlib.axes._subplots.AxesSubplot object at 0x12138b080>]], dtype=object)

Non-centered model

In [9]:
with pm.Model() as model_reparam:
    
    μ = pm.Normal('μ', 0, sd=100)
    σ = pm.HalfCauchy('σ', 5)
    ν = pm.Normal('ν', 0, sd=1, shape=k)
    
    θ = μ + ν*σ
    
    y_obs = pm.Normal('y_obs', θ[group_index], sd=obs_sd, observed=y)
In [10]:
with model_reparam:
    tr_reparam = pm.sample(n_iterations, n_init=50000, njobs=2)
Auto-assigning NUTS sampler...
Initializing NUTS using advi...
Average ELBO = -2,984.4: 100%|██████████| 50000/50000 [00:04<00:00, 10300.03it/s]
Finished [100%]: Average ELBO = -2,984.1
100%|██████████| 2000/2000 [00:12<00:00, 157.40it/s]
In [11]:
plots = pm.traceplot(tr_reparam[-n_to_keep:], varnames=['μ', 'σ'])