bayes_sleepstudy.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

%matplotlib inline

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 + np.exp(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:]

# model without mixed / hierarchical effects
with pm.Model() as lm:

    intercept = pm.Normal("intercept", mu=200, sd=5)
    slope = pm.Normal("slope", mu=10, sd=5)

    error = pm.HalfCauchy("eps", beta=1)

    predicted = intercept + slope * df.days.values
    y_like = pm.Normal("y_like", mu=predicted, sd=error, observed=df.reaction.values)

with lm:
    lm_trace = pm.sample(5000, tune=1000)[1000:]