pymc3_mixedmodel_sleepdata.py

import pandas as pd
import numpy as np
import os
import pymc3 as pm
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

df = pd.read_csv("sleep_study.csv")
df.columns = df.columns.str.lower()
dayscaler = StandardScaler()

df["days"] = dayscaler.fit_transform(df.iloc[:, 1].values.reshape(-1, 1))

n_days = df.days.nunique()
n_subjects = df.subject.nunique()

subject_names = df.subject.unique()
subject_ix = np.array(list(range(len(subject_names))))
subject_idx = df.subject.unique()
subject_code = df.subject.map(dict(zip(subject_idx, subject_ix))).values

with pm.Model() as sleepmodel:

    # specify priors for pooled intercept
    intercept_i_mu = pm.Normal("intercept_i_mu", mu=200, sd=20)
    intercept_i_sig = pm.HalfCauchy("intercept_i_sig", beta=1)

    # specify priors for pooled slope
    slope_i_mu = pm.Normal("slope_i_mu", mu=10, sd=3)
    slope_i_sig = pm.HalfCauchy("slope_i_sig", beta=3)

    # intercept for each subject
    intercept_ic_mu = pm.Normal("intercept_ic_mu",
                                              mu=0,
                                              sd=5,
                                              shape=n_subjects)
    inter = pm.Deterministic("inter", intercept_i_mu + intercept_ic_mu * intercept_i_sig)
    
    slope_ic_mu = pm.Normal("slope_ic_mu",
                                          mu=0,
                                          sd=5,
                                          shape=n_subjects)
    slope = pm.Deterministic("slope", slope_i_mu + slope_ic_mu * slope_i_sig)
    # Model error
    eps = pm.HalfCauchy('eps', beta=1)

    # Expected value
    pred = inter[subject_code] + slope[subject_code] * df.days.values

    # Data likelihood
    y_like = pm.Normal('y_like', mu=pred, sd=eps, observed=df.reaction.values)

with sleepmodel:
    hierarchical_trace = pm.sample(10000, tune=1000)[1000:]

pm.traceplot(hierarchical_trace)