%matplotlib inline
import pymc3 as pm
import numpy as np
import theano.tensor as tt
import matplotlib.pylab as plt
np.random.seed(20090425)
# 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)
m[:10]
array([ 10.11910178, 2.52654191, 4.51891947, 10.61821576, 4.42735171, 7.75110791, 1.02261171, 12.1946459 , 6.09602232, 14.20286786])
group_index = np.random.randint(0, k, size=n)
y = np.random.normal(loc=m[group_index], scale=obs_sd)
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)
n_iterations = 2000
n_to_keep = 1000
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]
pm.traceplot(tr[-n_to_keep:], varnames=['μ', 'σ'])
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)
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)
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]
plots = pm.traceplot(tr_reparam[-n_to_keep:], varnames=['μ', 'σ'])