#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('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]


# 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)


# In[8]:


pm.traceplot(tr[-n_to_keep:], varnames=['μ', 'σ'])


# ## 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)


# In[11]:


plots = pm.traceplot(tr_reparam[-n_to_keep:], varnames=['μ', 'σ'])