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

# # Chapter 16: Bayesian statistics

# Robert Johansson
# 
# Source code listings for [Numerical Python - Scientific Computing and Data Science Applications with Numpy, SciPy and Matplotlib](https://link.springer.com/book/10.1007/979-8-8688-0413-7) (ISBN 979-8-8688-0412-0).

# In[1]:


import pymc as mc 


# In[2]:


import numpy as np


# In[3]:


get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format='retina'")
import matplotlib.pyplot as plt


# In[4]:


import seaborn as sns
sns.set()


# In[5]:


from scipy import stats


# In[6]:


import statsmodels.api as sm


# In[7]:


import statsmodels.formula.api as smf


# In[8]:


import pandas as pd


# ## Changelog:
# 
# * 160828: The keyword argument `vars` to the functions `mc.traceplot` and `mc.forestplot` was changed to `varnames`.
# * 231125: The keyword `varnames` was changed to `var_names`

# # Simple example: Normal distributed random variable

# In[9]:


# try this
# dir(mc.distributions)


# In[10]:


np.random.seed(100)


# In[11]:


mu = 4.0


# In[12]:


sigma = 2.0


# In[13]:


model = mc.Model()


# In[14]:


with model:
    mc.Normal('X', mu, tau=1/sigma**2)


# In[15]:


model.continuous_value_vars


# In[16]:


start = dict(X=2)


# In[17]:


with model:
    step = mc.Metropolis()
    trace = mc.sample(10000, step=step, initvals=start)


# In[18]:


x = np.linspace(-4, 12, 1000)


# In[19]:


y = stats.norm(mu, sigma).pdf(x)


# In[20]:


# import arviz as az


# In[21]:


# az.extract(trace)


# In[22]:


def get_values(trace, variable):
    return trace.posterior.stack(sample=['chain', 'draw'])[variable].values


# In[23]:


X = get_values(trace, "X")


# In[24]:


X


# In[25]:


fig, ax = plt.subplots(figsize=(8, 3))

ax.plot(x, y, 'r', lw=2)
sns.histplot(X, ax=ax, kde=True, stat='density')
ax.set_xlim(-4, 12)
ax.set_xlabel("x")
ax.set_ylabel("Probability distribution")
fig.tight_layout()
fig.savefig("ch16-normal-distribution-sampled.pdf")
fig.savefig("ch16-normal-distribution-sampled.png", dpi=600)


# In[26]:


fig, axes = plt.subplots(1, 2, figsize=(12, 4), squeeze=False)
mc.plot_trace(trace, axes=axes)
axes[0,0].plot(x, y, 'r', lw=0.5)
fig.tight_layout()
fig.savefig("ch16-normal-sampling-trace.png", dpi=600)
fig.savefig("ch16-normal-sampling-trace.pdf")


# ## Dependent random variables

# In[27]:


model = mc.Model()


# In[28]:


#mc.HalfNormal??


# In[29]:


with model:
    mean = mc.Normal('mean', 3.0)
    sigma = mc.HalfNormal('sigma', sigma=1.0)
    X = mc.Normal('X', mean, tau=sigma)


# In[30]:


model.continuous_value_vars


# In[31]:


with model:
    start = mc.find_MAP()


# In[32]:


start


# In[33]:


with model:
    step = mc.Metropolis()
    trace = mc.sample(10000, initvals=start, step=step)


# In[34]:


get_values(trace, "sigma").mean()


# In[35]:


X = get_values(trace, "X")


# In[36]:


X.mean()


# In[37]:


X.std()


# In[38]:


fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=['mean', 'sigma', 'X'], axes=axes)
fig.tight_layout()
fig.savefig("ch16-dependent-rv-sample-trace.png", dpi=600)
fig.savefig("ch16-dependent-rv-sample-trace.pdf")


# ## Posterior distributions

# In[39]:


mu = 2.5


# In[40]:


s = 1.5


# In[41]:


data = stats.norm(mu, s).rvs(100)


# In[42]:


with mc.Model() as model:
    
    mean = mc.Normal('mean', 4.0, tau=1.0) # true 2.5
    sigma = mc.HalfNormal('sigma', tau=3.0 * np.sqrt(np.pi/2)) # true 1.5

    X = mc.Normal('X', mean, tau=1/sigma**2, observed=data)


# In[43]:


model.continuous_value_vars


# In[44]:


with model:
    start = mc.find_MAP()
    step = mc.Metropolis()
    trace = mc.sample(10000, initvals=start, step=step)
    #step = mc.NUTS()
    #trace = mc.sample(10000, initvals=start, step=step)


# In[45]:


start


# In[46]:


model.continuous_value_vars


# In[47]:


fig, axes = plt.subplots(2, 2, figsize=(8, 4), squeeze=False)
mc.plot_trace(trace, var_names=['mean', 'sigma'], axes=axes)
fig.tight_layout()
fig.savefig("ch16-posterior-sample-trace.png", dpi=600)
fig.savefig("ch16-posterior-sample-trace.pdf")


# In[49]:


mu, get_values(trace, "mean").mean() # trace.posterior.stack(sample=['chain', 'draw'])["mean"]


# In[50]:


s, get_values(trace, "sigma").mean() # trace.posterior.stack(sample=['chain', 'draw'])["sigma"].values.mean()


# In[51]:


gs = mc.plot_forest(trace, var_names=['mean', 'sigma'])
plt.savefig("ch16-forestplot.pdf")
plt.savefig("ch16-forestplot.png", dpi=600)


# In[52]:


# help(mc.summary)


# In[53]:


mc.summary(trace, var_names=['mean', 'sigma'])


# ## Linear regression

# In[54]:


dataset = sm.datasets.get_rdataset("Davis", "carData", cache=True)


# In[55]:


data = dataset.data[dataset.data.sex == 'M']


# In[56]:


data = data[data.weight < 110]


# In[57]:


data.head(3)


# In[58]:


model = smf.ols("height ~ weight", data=data)


# In[59]:


result = model.fit()


# In[60]:


print(result.summary())


# In[61]:


x = np.linspace(50, 110, 25)


# In[62]:


y = result.predict({"weight": x})


# In[63]:


fig, ax = plt.subplots(1, 1, figsize=(8, 3))
ax.plot(data.weight, data.height, 'o')
ax.plot(x, y, color="blue")
ax.set_xlabel("weight")
ax.set_ylabel("height")
fig.tight_layout()
fig.savefig("ch16-linear-ols-fit.pdf")
fig.savefig("ch16-linear-ols-fit.png", dpi=600)


# In[64]:


with mc.Model() as model:
    sigma = mc.Uniform('sigma', 0, 10)
    intercept = mc.Normal('intercept', 125, sigma=30)
    beta = mc.Normal('beta', 0, sigma=5)
    
    height_mu = intercept + beta * data.weight

    # likelihood function
    mc.Normal('height', mu=height_mu, sigma=sigma, observed=data.height)

    # predict
    predict_height = mc.Normal('predict_height', mu=intercept + beta * x, sigma=sigma, shape=len(x)) 


# In[65]:


model.continuous_value_vars


# In[66]:


# help(mc.NUTS)


# In[67]:


with model:
    # start = mc.find_MAP()
    step = mc.NUTS()
    trace = mc.sample(10000, step=step) # , initvals=start)


# In[68]:


model.continuous_value_vars


# In[69]:


fig, axes = plt.subplots(2, 2, figsize=(8, 4), squeeze=False)
mc.plot_trace(trace, var_names=['intercept', 'beta'], axes=axes)
fig.tight_layout()
fig.savefig("ch16-linear-model-sample-trace.pdf")
fig.savefig("ch16-linear-model-sample-trace.png", dpi=600)


# In[70]:


intercept = get_values(trace, "intercept").mean()
intercept


# In[71]:


beta = get_values(trace, "beta").mean() # trace.posterior.stack(sample=['chain', 'draw'])["beta"].values.mean()
beta


# In[72]:


result.params


# In[73]:


result.predict({"weight": 90})


# In[74]:


weight_index = np.where(x == 90)[0][0]


# In[75]:


# trace.posterior.stack(sample=['chain', 'draw'])["predict_height"].values[weight_index, :].mean()
get_values(trace, "predict_height")[weight_index, :].mean()


# In[76]:


#trace.get_values("predict_height")[:, weight_index].mean()


# In[77]:


#trace.posterior.stack(sample=['chain', 'draw'])["predict_height"].values.shape


# In[78]:


fig, ax = plt.subplots(figsize=(8, 3))

sns.histplot(get_values(trace, "predict_height")[weight_index, :], ax=ax, kde=True, stat='density')
ax.set_xlim(150, 210)
ax.set_xlabel("height")
ax.set_ylabel("Probability distribution")
fig.tight_layout()
fig.savefig("ch16-linear-model-predict-cut.pdf")
fig.savefig("ch16-linear-model-predict-cut.png", dpi=600)


# In[79]:


# trace.posterior.stack(sample=['chain', 'draw'])["intercept"].values


# In[80]:


fig, ax = plt.subplots(1, 1, figsize=(12, 4))

for n in range(500, 2000, 1):
    intercept = get_values(trace, "intercept")[n]
    beta = get_values(trace, "beta")[n]
    ax.plot(x, intercept + beta * x, color='red', lw=0.25, alpha=0.05)

intercept = get_values(trace, "intercept").mean()
beta = get_values(trace, "beta").mean()
ax.plot(x, intercept + beta * x, color='k', label="Mean Bayesian prediction")

ax.plot(data.weight, data.height, 'o')
ax.plot(x, y, '--', color="blue", label="OLS prediction")
ax.set_xlabel("weight")
ax.set_ylabel("height")
ax.legend(loc=0)
fig.tight_layout()
fig.savefig("ch16-linear-model-fit.pdf")
fig.savefig("ch16-linear-model-fit.png", dpi=600)


# In[81]:


import bambi


# In[82]:


model = bambi.Model('height ~ weight', data)


# In[83]:


trace = model.fit(draws=2000) #  chains=4


# In[84]:


fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=['Intercept', 'weight', 'height_sigma'], axes=axes)
fig.tight_layout()
fig.savefig("ch16-glm-sample-trace.pdf")
fig.savefig("ch16-glm-sample-trace.png", dpi=600)


# ## Multilevel model

# In[85]:


dataset = sm.datasets.get_rdataset("Davis", "carData", cache=True)


# In[86]:


data = dataset.data.copy()
data = data[data.weight < 110]


# In[87]:


data["sex"] = data["sex"].apply(lambda x: 1 if x == "F" else 0)


# In[88]:


data.head()


# In[89]:


with mc.Model() as model:

    # heirarchical model: hyper priors
    #intercept_mu = mc.Normal("intercept_mu", 125)
    #intercept_sigma = 30.0 #mc.Uniform('intercept_sigma', lower=0, upper=50)
    #beta_mu = mc.Normal("beta_mu", 0.0)
    #beta_sigma = 5.0 #mc.Uniform('beta_sigma', lower=0, upper=10)
    
    # multilevel model: prior parameters
    intercept_mu, intercept_sigma = 125, 30
    beta_mu, beta_sigma = 0.0, 5.0
    
    # priors
    intercept = mc.Normal('intercept', intercept_mu, sigma=intercept_sigma, shape=2)
    beta = mc.Normal('beta', beta_mu, sigma=beta_sigma, shape=2)
    error = mc.Uniform('error', 0, 10)

    # model equation
    sex_idx = data.sex.values
    height_mu = intercept[sex_idx] + beta[sex_idx] * data.weight

    mc.Normal('height', mu=height_mu, sigma=error, observed=data.height)


# In[90]:


model.continuous_value_vars


# In[92]:


with model:
    start = mc.find_MAP()
    # hessian = mc.find_hessian(start)
    step = mc.NUTS()
    trace = mc.sample(5000, step=step, initvals=start)


# In[93]:


fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=['intercept', 'beta', 'error'], axes=axes)
fig.tight_layout()
fig.savefig("ch16-multilevel-sample-trace.pdf")
fig.savefig("ch16-multilevel-sample-trace.png", dpi=600)


# In[94]:


intercept_m, intercept_f = get_values(trace, 'intercept').mean(axis=1)


# In[95]:


intercept = get_values(trace, 'intercept').mean()


# In[96]:


get_values(trace, 'beta')


# In[97]:


beta_m, beta_f = get_values(trace, 'beta').mean(axis=1)


# In[98]:


beta = get_values(trace, 'beta').mean()


# In[99]:


fig, ax = plt.subplots(1, 1, figsize=(8, 3))

mask_m = data.sex == 0
mask_f = data.sex == 1

ax.plot(data.weight[mask_m], data.height[mask_m], 'o', color="steelblue", label="male", alpha=0.5)
ax.plot(data.weight[mask_f], data.height[mask_f], 'o', color="green", label="female", alpha=0.5)

x = np.linspace(35, 110, 50)
ax.plot(x, intercept_m + x * beta_m, color="steelblue", label="model male group")
ax.plot(x, intercept_f + x * beta_f, color="green", label="model female group")
ax.plot(x, intercept + x * beta, color="black", label="model both groups")

ax.set_xlabel("weight")
ax.set_ylabel("height")
ax.legend(loc=0)
fig.tight_layout()
fig.savefig("ch16-multilevel-linear-model-fit.pdf")
fig.savefig("ch16-multilevel-linear-model-fit.png", dpi=600)


# In[100]:


get_values(trace, 'error').mean()


# # Version

# In[101]:


get_ipython().run_line_magic('reload_ext', 'version_information')


# In[102]:


get_ipython().run_line_magic('version_information', 'numpy, pandas, matplotlib, statsmodels, pymc3, theano')