!pip install arviz

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm
import pandas as pd
import theano
import seaborn as sns

sns.set_style("whitegrid")
np.random.seed(123)

url = "https://github.com/twiecki/WhileMyMCMCGentlySamples/blob/master/content/downloads/notebooks/radon.csv?raw=true"
data = pd.read_csv(url)
# data = pd.read_csv('../data/radon.csv')
data["log_radon"] = data["log_radon"].astype(theano.config.floatX)
county_names = data.county.unique()
county_idx = data.county_code.values

n_counties = len(data.county.unique())

"""
Explanation: <a href="https://colab.research.google.com/github/probml/pyprobml/blob/master/book1/bayes_stats/linreg_hierarchical_non_centered_pymc3.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Hierarchical non-centered Bayesian Linear Regression in PyMC3
The text and code for this notebook are taken directly from this blog post
 by Thomas Wiecki. Original notebook
End of explanation
"""


with pm.Model() as hierarchical_model_centered:
    # Hyperpriors for group nodes
    mu_a = pm.Normal("mu_a", mu=0.0, sd=100**2)
    sigma_a = pm.HalfCauchy("sigma_a", 5)
    mu_b = pm.Normal("mu_b", mu=0.0, sd=100**2)
    sigma_b = pm.HalfCauchy("sigma_b", 5)

    # Intercept for each county, distributed around group mean mu_a
    # Above we just set mu and sd to a fixed value while here we
    # plug in a common group distribution for all a and b (which are
    # vectors of length n_counties).
    a = pm.Normal("a", mu=mu_a, sd=sigma_a, shape=n_counties)

    # Intercept for each county, distributed around group mean mu_a
    b = pm.Normal("b", mu=mu_b, sd=sigma_b, shape=n_counties)

    # Model error
    eps = pm.HalfCauchy("eps", 5)

    # Linear regression
    radon_est = a[county_idx] + b[county_idx] * data.floor.values

    # Data likelihood
    radon_like = pm.Normal("radon_like", mu=radon_est, sd=eps, observed=data.log_radon)

# Inference button (TM)!
with hierarchical_model_centered:
    hierarchical_centered_trace = pm.sample(draws=5000, tune=1000)[1000:]

pm.traceplot(hierarchical_centered_trace);

"""
Explanation: The intuitive specification
Usually, hierachical models are specified in a centered way. In a regression model, individual slopes would be centered around a group mean with a certain group variance, which controls the shrinkage:
End of explanation
"""


print("Rhat(sigma_b) = {}".format(pm.diagnostics.gelman_rubin(hierarchical_centered_trace)["sigma_b"]))

"""
Explanation: I have seen plenty of traces with terrible convergences but this one might look fine to the unassuming eye. Perhaps sigma_b has some problems, so let's look at the Rhat:
End of explanation
"""


fig, axs = plt.subplots(nrows=2)
axs[0].plot(hierarchical_centered_trace.get_values("sigma_b", chains=1), alpha=0.5)
axs[0].set(ylabel="sigma_b")
axs[1].plot(hierarchical_centered_trace.get_values("b", chains=1), alpha=0.5)
axs[1].set(ylabel="b");

"""
Explanation: Not too bad -- well below 1.01. I used to think this wasn't a big deal but Michael Betancourt in his StanCon 2017 talk makes a strong point that it is actually very problematic. To understand what's going on, let's take a closer look at the slopes b and their group variance (i.e. how far they are allowed to move from the mean) sigma_b. I'm just plotting a single chain now.
End of explanation
"""


x = pd.Series(hierarchical_centered_trace["b"][:, 75], name="slope b_75")
y = pd.Series(hierarchical_centered_trace["sigma_b"], name="slope group variance sigma_b")

sns.jointplot(x, y, ylim=(0, 0.7));

"""
Explanation: sigma_b seems to drift into this area of very small values and get stuck there for a while. This is a common pattern and the sampler is trying to tell you that there is a region in space that it can't quite explore efficiently.  While stuck down there, the slopes b_i become all squished together. We've entered The Funnel of Hell (it's just called the funnel, I added the last part for dramatic effect).
The Funnel of Hell (and how to escape it)
Let's look at the joint posterior of a single slope b (I randomly chose the 75th one) and the slope group variance sigma_b.
End of explanation
"""


with pm.Model() as hierarchical_model_non_centered:
    # Hyperpriors for group nodes
    mu_a = pm.Normal("mu_a", mu=0.0, sd=100**2)
    sigma_a = pm.HalfCauchy("sigma_a", 5)
    mu_b = pm.Normal("mu_b", mu=0.0, sd=100**2)
    sigma_b = pm.HalfCauchy("sigma_b", 5)

    # Before:
    # a = pm.Normal('a', mu=mu_a, sd=sigma_a, shape=n_counties)
    # Transformed:
    a_offset = pm.Normal("a_offset", mu=0, sd=1, shape=n_counties)
    a = pm.Deterministic("a", mu_a + a_offset * sigma_a)

    # Before:
    # b = pm.Normal('b', mu=mu_b, sd=sigma_b, shape=n_counties)
    # Now:
    b_offset = pm.Normal("b_offset", mu=0, sd=1, shape=n_counties)
    b = pm.Deterministic("b", mu_b + b_offset * sigma_b)

    # Model error
    eps = pm.HalfCauchy("eps", 5)

    radon_est = a[county_idx] + b[county_idx] * data.floor.values

    # Data likelihood
    radon_like = pm.Normal("radon_like", mu=radon_est, sd=eps, observed=data.log_radon)

"""
Explanation: This makes sense, as the slope group variance goes to zero (or, said differently, we apply maximum shrinkage), individual slopes are not allowed to deviate from the slope group mean, so they all collapose to the group mean.
While this property of the posterior in itself is not problematic, it makes the job extremely difficult for our sampler. Imagine a Metropolis-Hastings exploring this space with a medium step-size (we're using NUTS here but the intuition works the same): in the wider top region we can comfortably make larger jumps to explore the space efficiently. However, once we move to the narrow bottom region we can change b_75 and sigma_b only by tiny amounts. This causes the sampler to become trapped in that region of space. Most of the proposals will be rejected because our step-size is too large for this narrow part of the space and exploration will be very inefficient.
You might wonder if we could somehow choose the step-size based on the denseness (or curvature) of the space. Indeed that's possible and it's called Riemannian HMC. It works very well but is quite costly to run. Here, we will explore a different, simpler method.
Finally, note that this problem does not exist for the intercept parameters a. Because we can determine individual intercepts a_i with enough confidence, sigma_a is not small enough to be problematic. Thus, the funnel of hell can be a problem in hierarchical models, but it does not have to be. (Thanks to John Hall for pointing this out).
Reparameterization
If we can't easily make the sampler step-size adjust to the region of space, maybe we can adjust the region of space to make it simpler for the sampler? This is indeed possible and quite simple with a small reparameterization trick, we will call this the non-centered version.
End of explanation
"""


# Inference button (TM)!
with hierarchical_model_non_centered:
    hierarchical_non_centered_trace = pm.sample(draws=5000, tune=1000)[1000:]

pm.traceplot(hierarchical_non_centered_trace, varnames=["sigma_b"]);

"""
Explanation: Pay attention to the definitions of a_offset, a, b_offset, and b and compare them to before (commented out). What's going on here? It's pretty neat actually. Instead of saying that our individual slopes b are normally distributed around a group mean (i.e. modeling their absolute values directly), we can say that they are offset from a group mean by a certain value (b_offset; i.e. modeling their values relative to that mean). Now we still have to consider how far from that mean we actually allow things to deviate (i.e. how much shrinkage we apply). This is where sigma_b makes a comeback. We can simply multiply the offset by this scaling factor to get the same effect as before, just under a different parameterization. For a more formal introduction, see e.g. Betancourt & Girolami (2013).
Critically, b_offset and sigma_b are now mostly independent. This will become more clear soon. Let's first look at if this transform helped our sampling:
End of explanation
"""


fig, axs = plt.subplots(ncols=2, sharex=True, sharey=True)

x = pd.Series(hierarchical_centered_trace["b"][:, 75], name="slope b_75")
y = pd.Series(hierarchical_centered_trace["sigma_b"], name="slope group variance sigma_b")

axs[0].plot(x, y, ".")
axs[0].set(title="Centered", ylabel="sigma_b", xlabel="b_75")

x = pd.Series(hierarchical_non_centered_trace["b"][:, 75], name="slope b_75")
y = pd.Series(hierarchical_non_centered_trace["sigma_b"], name="slope group variance sigma_b")

axs[1].plot(x, y, ".")
axs[1].set(title="Non-centered", xlabel="b_75");

"""
Explanation: That looks much better as also confirmed by the joint plot:
End of explanation
"""


pm.kdeplot(
    np.stack(
        [
            hierarchical_centered_trace["sigma_b"],
            hierarchical_non_centered_trace["sigma_b"],
        ]
    ).T
)
plt.axvline(hierarchical_centered_trace["sigma_b"].mean(), color="b", linestyle="--")
plt.axvline(hierarchical_non_centered_trace["sigma_b"].mean(), color="g", linestyle="--")
plt.legend(["Centered", "Non-cenetered", "Centered posterior mean", "Non-centered posterior mean"])
plt.xlabel("sigma_b")
plt.ylabel("Probability Density");

"""
Explanation: To really drive this home, let's also compare the sigma_b marginal posteriors of the two models:
End of explanation
"""


x = pd.Series(hierarchical_non_centered_trace["b_offset"][:, 75], name="slope b_offset_75")
y = pd.Series(hierarchical_non_centered_trace["sigma_b"], name="slope group variance sigma_b")

sns.jointplot(x, y, ylim=(0, 0.7))

"""
Explanation: That's crazy -- there's a large region of very small sigma_b values that the sampler could not even explore before. In other words, our previous inferences ("Centered") were severely biased towards higher values of sigma_b. Indeed, if you look at the previous blog post the sampler never even got stuck in that low region causing me to believe everything was fine. These issues are hard to detect and very subtle, but they are meaningful as demonstrated by the sizable difference in posterior mean.
But what does this concretely mean for our analysis? Over-estimating sigma_b means that we have a biased (=false) belief that we can tell individual slopes apart better than we actually can. There is less information in the individual slopes than what we estimated.
Why does the reparameterized model work better?
To more clearly understand why this model works better, let's look at the joint distribution of b_offset:
End of explanation
"""


with hierarchical_model_centered:
    mode = pm.find_MAP()

mode["b"]

np.exp(mode["sigma_b_log_"])

"""
Explanation: This is the space the sampler sees; you can see how the funnel is flattened out. We can freely change the (relative) slope offset parameters even if the slope group variance is tiny as it just acts as a scaling parameter.
Note that the funnel is still there -- it's a perfectly valid property of the model -- but the sampler has a much easier time exploring it in this different parameterization.
Why hierarchical models are Bayesian
Finally, I want to take the opportunity to make another point that is not directly related to hierarchical models but can be demonstrated quite well here.
Usually when talking about the perils of Bayesian statistics we talk about priors, uncertainty, and flexibility when coding models using Probabilistic Programming. However, an even more important property is rarely mentioned because it is much harder to communicate. Ross Taylor touched on this point in his tweet:
<blockquote class="twitter-tweet" data-lang="en"><p lang="en" dir="ltr">It&#39;s interesting that many summarize Bayes as being about priors; but real power is its focus on integrals/expectations over maxima/modes</p>&mdash; Ross Taylor (@rosstaylor90) <a href="https://twitter.com/rosstaylor90/status/827263854002401281">February 2, 2017</a></blockquote>
<script async src="//platform.twitter.com/widgets.js" charset="utf-8"></script>

Michael Betancourt makes a similar point when he says "Expectations are the only thing that make sense."
But what's wrong with maxima/modes? Aren't those really close to the posterior mean (i.e. the expectation)? Unfortunately, that's only the case for the simple models we teach to build up intuitions. In complex models, like the hierarchical one, the MAP can be far away and not be interesting or meaningful at all.
Let's compare the posterior mode (i.e. the MAP) to the posterior mean of our hierachical linear regression model:
End of explanation
"""


hierarchical_non_centered_trace["b"].mean(axis=0)

hierarchical_non_centered_trace["sigma_b"].mean(axis=0)

"""
Explanation: As you can see, the slopes are all identical and the group slope variance is effectively zero. The reason is again related to the funnel. The MAP only cares about the probability density which is highest at the bottom of the funnel. 
But if you could only choose one point in parameter space to summarize the posterior above, would this be the one you'd pick? Probably not.
Let's instead look at the Expected Value (i.e. posterior mean) which is computed by integrating probability density and volume to provide probabilty mass -- the thing we really care about. Under the hood, that's the integration performed by the MCMC sampler.
End of explanation
"""


