Hierarchical Models

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet jupyterlab-vim)"
!jupyter labextension enable

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet graphviz)"

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet dataframe_image)"

%load_ext autoreload
%autoreload 2


import arviz as az
import pandas as pd
import pymc as pm
import numpy as np
import seaborn as sns

dir_name = "./L07_data"

!ls $dir_name

import msml610.tutorials.msml610_utils as ut

ut.config_notebook()

tips = pd.read_csv(dir_name + "/tips.csv")
tips

sns.boxplot(x="day", y="tip", data=tips)

# Extract the tips.
tip = tips["tip"].values
print(tip[:10])

# Create a vector going from day to group idx.
idx = pd.Categorical(tips["day"]).codes
print("idx=", idx)

# Count the groups.
groups = np.unique(idx)
n_groups = len(groups)
print("groups=", n_groups, groups)

# The model is the same as before but it can be easily vectorized.
# There is no need to write a for-loop.
with pm.Model() as comparing_groups:
    # mu is a vector of 4 elems.
    mu = pm.Normal("mu", mu=0, sigma=10, shape=n_groups)
    # sigma is a vector of 4 elems.
    sigma = pm.HalfNormal("sigma", sigma=10, shape=n_groups)
    # y is a vector of 4 normals each with mean and sigma for the group.
    y = pm.Normal("y", mu=mu[idx], sigma=sigma[idx], observed=tip)
    idata_cg = pm.sample(5000)

cs_data = pd.read_csv(dir_name + "/chemical_shifts_theo_exp.csv")
cs_data["diff"] = cs_data["theo"] - cs_data["exp"]
display(cs_data)

diff = cs_data.theo.values - cs_data.exp.values
print("diff=", diff)

# Array of categorical values.
cat_encode = pd.Categorical(cs_data["aa"])
print("cat_encode=", cat_encode)
idx = cat_encode.codes
print("idx=", len(idx), idx)
coords = {"aa": cat_encode.categories}
print("coords=", coords)

# Non-hierarchical model.
with pm.Model(coords=coords) as cs_nh:
    # One separate prior for each group.
    mu = pm.Normal("mu", mu=0, sigma=10, dims="aa")
    sigma = pm.HalfNormal("sigma", sigma=10, dims="aa")
    # Likelihood.
    y = pm.Normal("y", mu=mu[idx], sigma=sigma[idx], observed=diff)
    idata_cs_nh = pm.sample()

pm.model_to_graphviz(cs_nh)

with pm.Model(coords=coords) as cs_h:
    # Hyper-priors.
    mu_mu = pm.Normal("mu_mu", mu=0, sigma=10)
    mu_sigma = pm.HalfNormal("mu_sigma", sigma=10)

    # Priors.
    mu = pm.Normal("mu", mu=mu_mu, sigma=mu_sigma, dims="aa")
    sigma = pm.HalfNormal("sigma", sigma=10, dims="aa")

    # Likelihood (same as before).
    y = pm.Normal("y", mu=mu[idx], sigma=sigma[idx], observed=diff)
    idata_cs_h = pm.sample()

pm.model_to_graphviz(cs_h)

# We have two models and we want to compare the estimates.
# - There are 20 groups and each model has 4 estimates.
# - We plot the 94% credible intervals.
# - The vertical line is the global mean according to the hierarchical model.
# - The blue (hierarchical) means are pulled towards the mean, wrt the orange (non-hierarchical) ones.
axes = az.plot_forest(
    [idata_cs_h, idata_cs_nh],
    model_names=["h", "n_h"],
    var_names="mu",
    combined=True,
    colors="cycle",
)

y_lims = axes[0].get_ylim()
axes[0].vlines(idata_cs_h.posterior["mu_mu"].mean(), *y_lims, color="navy")
axes[0].vlines(idata_cs_nh.posterior["mu"].mean(), *y_lims, color="orange")