<jupyter_start><jupyter_text>Bayesian Hierarchical ModelingThis jupyter notebook accompanies the Bayesian Hierarchical Modeling lecture(s) delivered by <NAME> as part of <NAME>'s [Computational Data Analysis class](http://dwh.gg/FlatironCDA). As part of the lecture(s) you will be asked to complete a number of tasks, some of which will involve direct coding into the notebook; these sections are marked by task. This notebook requires numpy, matplotlib, scipy, [corner](https://github.com/sfeeney/bhm_lecture.git), [pystan](https://pystan.readthedocs.io/en/latest/getting_started.html) and pickle to run (the last two are required solely for the final task).The model we're going to be inferring is below.We start with imports...<jupyter_code>from __future__ import print_function

# make sure everything we need is installed if running on Google Colab
def is_colab():
    try:
        cfg = get_ipython().config
        if cfg['IPKernelApp']['kernel_class'] == 'google.colab._kernel.Kernel':
            return True
        else:
            return False
    except NameError:
        return False
if is_colab():
    !pip install --quiet numpy matplotlib scipy corner pystan

import numpy as np
import numpy.random as npr
import matplotlib.pyplot as mp
%matplotlib inline<jupyter_output><empty_output><jupyter_text>... and immediately move to... Task 2In which I ask you to write a Python function to generate a simulated Cepheid sample using the period-luminosity relation $m_{ij} = \mu_i + M^* + s\,\log p_{ij} + \epsilon(\sigma_{\rm int})$. For simplicity, assume Gaussian priors on everything, Gaussian intrinsic scatter and Gaussian measurement uncertainties. Assume only the first host has a distance modulus estimate.<jupyter_code># setup
n_gal = 2
n_star = 200
n_samples = 50000

# PL relation parameters
abs_bar = -26.0   # mean of standard absolute magnitude prior
abs_sig = 4.0     # std dev of standard absolute magnitude prior
s_bar = -1.0      # mean of slope prior
s_sig = 1.0       # std dev of slope prior
mu_bar = 30.0     # mean of distance modulus prior
mu_sig = 5.0      # std dev of distance modulus prior
m_sig_int = 0.05  # intrinsic scatter, assumed known

# uncertainties
mu_hat_sig = 0.01 # distance modulus measurement uncertainty
m_hat_sig = 0.02  # apparent magnitude measurement uncertainty

def simulate(n_gal, n_star, abs_bar, abs_sig, s_bar, s_sig, mu_bar, mu_sig, mu_hat_sig, m_sig_int, m_hat_sig):
    
    # draw CPL parameters from Gaussian prior with means abs_bar and s_bar and standard deviations
    # abs_sig and s_sig
    #abs_true = abs_bar
    #s_true = s_bar
    abs_true = abs_bar + npr.randn() * abs_sig
    s_true = s_bar + npr.randn() * s_sig
    
    # draw n_gal distance moduli from Gaussian prior with mean mu_bar and standard deviation mu_sig
    # i've chosen to sort here so the closest galaxy is the one with the measured distance modulus
    mu_true = np.sort(mu_bar + npr.randn(n_gal) * mu_sig)
     
    # measure ONLY ONE galaxy's distance modulus noisily. the noise here is assumed Gaussian with
    # zero mean and standard deviation mu_hat_sig
    mu_hat = mu_true[0] + npr.randn() * mu_hat_sig
    
    # draw log periods. these are assumed to be perfectly observed in this model, so they 
    # are simply a set of pre-specified numbers. i have chosen to generate new values with 
    # each simulation, drawn such that log-periods are uniformly drawn in the range 1-2 (i.e., 
    # 10 to 100 days). you can have these for free!
    lp_true = 1.0 + npr.rand(n_gal, n_star)
    
    # draw true apparent magnitudes. these are distributed around the Cepheid period-luminosity 
    # relation with Gaussian intrinsic scatter (mean 0, standard deviation m_sig_int)
    m_true = np.zeros((n_gal, n_star))
    for i in range(n_gal):
        m_true[i, :] = mu_true[i] + abs_true + s_true * lp_true[i, :] + npr.randn(n_star) * m_sig_int
    
    # measure the apparent magnitudes noisily, all with the same measurement uncertainty m_hat_sig
    m_hat = m_true + npr.randn(n_gal, n_star) * m_hat_sig
    
    # return!
    return (abs_true, s_true, mu_true, lp_true, m_true, mu_hat, m_hat)<jupyter_output><empty_output><jupyter_text>Let's check that the simulation generates something sane. A simple test that the magnitude measurements errors are correctly generated.<jupyter_code># simulate
abs_true, s_true, mu_true, lp_true, m_true, mu_hat, m_hat = \
    simulate(n_gal, n_star, abs_bar, abs_sig, s_bar, s_sig, mu_bar, mu_sig, mu_hat_sig, m_sig_int, m_hat_sig)

# plot difference between true and observed apparent magnitudes. this should be the 
# noise, which is Gaussian distributed with mean zero and std dev m_hat_sig
outs = mp.hist((m_true - m_hat).flatten())
dm_grid = np.linspace(np.min(outs[1]), np.max(outs[1]))
mp.plot(dm_grid, np.exp(-0.5 * (dm_grid/m_hat_sig) ** 2) * np.max(outs[0]))
mp.xlabel(r'$m_{ij} - \hat{m}_{ij}$')
mp.ylabel(r'$N \left(m_{ij} - \hat{m}_{ij}\right)$')<jupyter_output><empty_output><jupyter_text>And another test that the intrinsic scatter is added as expected.<jupyter_code># plot difference between true apparent magnitudes and expected apparent 
# magnitude given a perfect (i.e., intrinsic-scatter-free) period-luminosity 
# relation. this should be the intrinsic scatter, which is Gaussian-
# distributed with mean zero and std dev m_sig_int
eps = np.zeros((n_gal, n_star))
for i in range(n_gal):
    eps[i, :] = mu_true[i] + abs_true + s_true * lp_true[i, :] - m_true[i, :]
outs = mp.hist(eps.flatten())
dm_grid = np.linspace(np.min(outs[1]), np.max(outs[1]))
mp.plot(dm_grid, np.exp(-0.5 * (dm_grid/m_sig_int) ** 2) * np.max(outs[0]))
mp.xlabel(r'$m_{ij} - \hat{m}_{ij}$')
mp.ylabel(r'$N \left(m_{ij} - \hat{m}_{ij}\right)$')<jupyter_output><empty_output><jupyter_text>Generalized Least Squares DemoCoding up the [GLS estimator](https://en.wikipedia.org/wiki/Generalized_least_squares) is a little involved, so I've done it for you below. Note that, rather unhelpfully, I've done so in a different order than in the notes. When I get a chance I will re-write. For now, you can simply evaluate the cells and bask in the glory of the fastest inference you will ever do!<jupyter_code>def gls_fit(n_gal, n_star, mu_hat, mu_hat_sig, m_hat, m_sig_int, m_hat_sig, \
            lp_true, priors=None):

    # setup
    # n_obs is one anchor constraint and one magnitude per Cepheid.
    # n_par is one mu per Cepheid host and 2 CPL params. if priors 
    # are used, we add on n_gal + 2 observations: one prior constraint 
    # on each host distance modulus and CPL parameter
    n_obs = n_gal * n_star + 1
    n_par = n_gal + 2
    if priors is not None:
        n_obs += n_gal + 2
    data = np.zeros(n_obs)
    design = np.zeros((n_obs, n_par))
    cov_inv = np.zeros((n_obs, n_obs))
    
    # anchor
    data[0] = mu_hat
    design[0, 0] = 1.0
    cov_inv[0, 0] = 1.0 / mu_hat_sig ** 2

    # Cepheids
    k = 1
    for i in range(0, n_gal):
        for j in range(0, n_star):

            data[k] = m_hat[i, j]
            design[k, i] = 1.0
            design[k, n_gal] = 1.0
            design[k, n_gal + 1] = lp_true[i, j]
            cov_inv[k, k] = 1.0 / (m_hat_sig ** 2 + m_sig_int ** 2)
            k += 1
    
    # and, finally, priors if desired
    if priors is not None:
        abs_bar, abs_sig, s_bar, s_sig, mu_bar, mu_sig = priors
        for i in range(n_gal):
            data[k] = mu_bar
            design[k, i] = 1.0
            cov_inv[k, k] = 1.0 / mu_sig ** 2
            k += 1
        data[k] = abs_bar
        design[k, n_gal] = 1.0
        cov_inv[k, k] = 1.0 / abs_sig ** 2
        k += 1
        data[k] = s_bar
        design[k, n_gal + 1] = 1.0
        cov_inv[k, k] = 1.0 / s_sig ** 2
        k += 1
        
    # fit and return
    destci = np.dot(design.transpose(), cov_inv)
    pars_cov = np.linalg.inv(np.dot(destci, design))
    pars = np.dot(np.dot(pars_cov, destci), data)
    res = data - np.dot(design, pars)
    dof = n_obs - n_par
    chisq_dof = np.dot(res.transpose(), np.dot(cov_inv, res))
    return pars, pars_cov, chisq_dof
gls_pars, gls_pars_cov, gls_chisq = gls_fit(n_gal, n_star, mu_hat, mu_hat_sig, m_hat, \
                                            m_sig_int, m_hat_sig, lp_true, \
                                            priors=[abs_bar, abs_sig, s_bar, s_sig, mu_bar, mu_sig])<jupyter_output><empty_output><jupyter_text>In order to plot the outputs of the GLS fit we could draw a large number of samples from the resulting multivariate Gaussian posterior and pass them to something like [`corner`](https://corner.readthedocs.io/en/latest/); however, as we have analytic results we might as well use those directly. I've coded up something totally hacky here in order to do so. Information on how to draw confidence ellipses can be found in [<NAME>'s note](https://arxiv.org/pdf/0906.4123.pdf).<jupyter_code># this is a hacky function designed to transform the analytic GLS outputs
# into a corner.py style triangle plot, containing 1D and 2D marginalized
# posteriors
import scipy.stats as sps
import matplotlib.patches as mpp
def schmorner(par_mean, par_cov, par_true, par_label):
    
    # setup
    par_std = np.sqrt(np.diag(par_cov))
    x_min = par_mean[0] - 3.5 * par_std[0]
    x_max = par_mean[0] + 3.5 * par_std[0]
    y_min = par_mean[1] - 3.5 * par_std[1]
    y_max = par_mean[1] + 3.5 * par_std[1]
    fig, axes = mp.subplots(2, 2)
    
    # 1D marge
    x = np.linspace(x_min, x_max, 100)
    axes[0, 0].plot(x, sps.norm.pdf(x, par_mean[0], par_std[0]), 'k')
    axes[0, 0].axvline(par_true[0])
    axes[1, 0].axvline(par_true[0])
    axes[0, 0].set_xticklabels([])
    axes[0, 0].set_yticklabels([])
    axes[0, 0].set_xlim(x_min, x_max)
    axes[0, 0].set_title(par_label[0])
    axes[0, 0].set_title(par_label[0] + r'$=' + '{:6.2f}'.format(par_mean[0]) + \
                         r'\pm' + '{:4.2f}'.format(par_std[0]) + r'$')
    y = np.linspace(y_min, y_max, 100)
    axes[1, 1].plot(y, sps.norm.pdf(y, par_mean[1], par_std[1]), 'k')
    axes[1, 0].axhline(par_true[1])
    axes[1, 1].axvline(par_true[1])
    axes[1, 1].tick_params(labelleft=False)
    axes[1, 1].set_xlim(y_min, y_max)
    for tick in axes[1, 1].get_xticklabels():
        tick.set_rotation(45)
    axes[1, 1].set_title(par_label[1] + r'$=' + '{:5.2f}'.format(par_mean[1]) + \
                         r'\pm' + '{:4.2f}'.format(par_std[1]) + r'$')

    # 2D marge
    vals, vecs = np.linalg.eig(par_cov)
    theta = np.degrees(np.arctan2(*vecs[::-1, 0]))
    w, h = 2 * np.sqrt(vals)
    ell = mpp.Ellipse(xy=par_mean, width=w, height=h,
                      angle=theta, color='k')
    ell.set_facecolor("none")
    axes[1, 0].add_artist(ell)
    ell = mpp.Ellipse(xy=par_mean, width=2*w, height=2*h,
                      angle=theta, color='k')
    ell.set_facecolor("none")
    axes[1, 0].add_artist(ell)
    axes[1, 0].set_xlim(x_min, x_max)
    axes[1, 0].set_ylim(y_min, y_max)
    for tick in axes[1, 0].get_xticklabels():
        tick.set_rotation(45)
    for tick in axes[1, 0].get_yticklabels():
        tick.set_rotation(45)
    axes[1, 0].set_xlabel(par_label[0])
    axes[1, 0].set_ylabel(par_label[1])
    fig.delaxes(axes[0, 1])
    fig.subplots_adjust(hspace=0, wspace=0)
    
test = schmorner(gls_pars[n_gal:], gls_pars_cov[n_gal:, n_gal:], \
                 [abs_true, s_true], [r'$M$', r'$s$'])
#
#lazy = npr.multivariate_normal(gls_pars[n_gal:], gls_pars_cov[n_gal:, n_gal:], n_samples)
#fig = corner.corner(samples.T, labels=[r"$M$", r"$s$"],
#                    show_titles=True, truths=[abs_bar, s_bar])<jupyter_output><empty_output><jupyter_text>Task 3BBelow I've written the majority of a Gibbs sampler to infer the hyper-parameters of the Cepheid PL relation from our simulated sample. One component is missing: drawing from the conditional distribution of the standard absolute magnitude, $M^*$. Please fill it in, using the results of whiteboard/paper Task 3A.<jupyter_code>def gibbs_sample(n_samples, n_gal, n_star, abs_bar, abs_sig, \
                 s_bar, s_sig, mu_bar, mu_sig, mu_hat_sig, \
                 m_sig_int, m_hat_sig, mu_hat, lp_true, m_hat):
    
    # storage
    abs_samples = np.zeros(n_samples)
    s_samples = np.zeros(n_samples)
    mu_samples = np.zeros((n_gal, n_samples))
    m_samples = np.zeros((n_gal, n_star, n_samples))
    
    # initialize sampler
    abs_samples[0] = abs_bar + npr.randn() * abs_sig
    s_samples[0] = s_bar + npr.randn() * s_sig
    mu_samples[:, 0] = mu_bar + npr.randn(n_gal) * mu_bar
    for i in range(n_gal):
        m_samples[i, :, 0] = mu_samples[i, 0] + abs_samples[0] + s_samples[0] * lp_true[i, :]
    
    # sample!
    for k in range(1, n_samples):
        
        # sample abs mag
        abs_sig_pl = m_sig_int / np.sqrt(n_gal * n_star)
        abs_bar_pl = 0.0
        for j in range(n_gal):
            abs_bar_pl += np.sum(m_samples[j, :, k - 1] - mu_samples[j, k - 1] - s_samples[k - 1] * lp_true[j, :])
        abs_bar_pl /= (n_gal * n_star)
        abs_std = np.sqrt((abs_sig * abs_sig_pl) ** 2 / (abs_sig ** 2 + abs_sig_pl ** 2))
        abs_mean = (abs_sig ** 2 * abs_bar_pl + abs_sig_pl ** 2 * abs_bar) / \
                   (abs_sig ** 2 + abs_sig_pl ** 2)
        abs_samples[k] = abs_mean + npr.randn() * abs_std
        
        # sample slope
        s_sig_pl = m_sig_int / np.sqrt(np.sum(lp_true ** 2))
        s_bar_pl = 0.0
        for j in range(n_gal):
            s_bar_pl += np.sum((m_samples[j, :, k - 1] - mu_samples[j, k - 1] - abs_samples[k]) * lp_true[j, :])
        s_bar_pl /= np.sum(lp_true ** 2)
        s_std = np.sqrt((s_sig * s_sig_pl) ** 2 / (s_sig ** 2 + s_sig_pl ** 2))
        s_mean = (s_sig ** 2 * s_bar_pl + s_sig_pl ** 2 * s_bar) / \
                 (s_sig ** 2 + s_sig_pl ** 2)
        s_samples[k] = s_mean + npr.randn() * s_std
        
        # sample apparent magnitudes
        for j in range(n_gal):
            m_mean_pl = mu_samples[j, k - 1] + abs_samples[k] + s_samples[k] * lp_true[j, :]
            m_std = np.sqrt(m_sig_int ** 2 * m_hat_sig ** 2 / (m_sig_int ** 2 + m_hat_sig ** 2))
            m_mean = (m_sig_int ** 2 * m_hat[j, :] + m_hat_sig ** 2 * m_mean_pl) / (m_sig_int ** 2 + m_hat_sig ** 2)
            m_samples[j, :, k] = m_mean + npr.randn(n_star) * m_std
            
        # sample distance moduli
        mu_sig_pl = m_sig_int / np.sqrt(n_star)
        mu_bar_pl = np.mean(m_samples[0, :, k] - abs_samples[k] - s_samples[k] * lp_true[0, :])
        mu_var = 1.0 / (1.0 / mu_sig ** 2 + 1.0 / mu_hat_sig ** 2 + 1.0 / mu_sig_pl ** 2)
        mu_mean = (mu_bar / mu_sig ** 2 + mu_hat / mu_hat_sig ** 2 + mu_bar_pl / mu_sig_pl ** 2) * mu_var
        mu_samples[0, k] = mu_mean + npr.randn() * np.sqrt(mu_var)
        for j in range(1, n_gal):
            mu_sig_pl = m_sig_int / np.sqrt(n_star)
            mu_bar_pl = np.mean(m_samples[j, :, k] - abs_samples[k] - s_samples[k] * lp_true[j, :])
            mu_std = (mu_sig * mu_sig_pl) ** 2 / (mu_sig ** 2 + mu_sig_pl ** 2)
            mu_mean = (mu_sig ** 2 * mu_bar_pl + mu_sig_pl ** 2 * mu_bar) / \
                      (mu_sig ** 2 + mu_sig_pl ** 2)
            mu_samples[j, k] = mu_mean + npr.randn() * mu_std
        
    return (abs_samples, s_samples, mu_samples, m_samples)<jupyter_output><empty_output><jupyter_text>Now let's sample, setting aside the first half of the samples as warmup.<jupyter_code>all_samples = gibbs_sample(n_samples, n_gal, n_star, abs_bar, abs_sig, \
                           s_bar, s_sig, mu_bar, mu_sig, mu_hat_sig, \
                           m_sig_int, m_hat_sig, mu_hat, lp_true, m_hat)
n_warmup = int(n_samples / 2)
g_samples = [samples[n_warmup:] for samples in all_samples]<jupyter_output><empty_output><jupyter_text>Let's make sure that the absolute magnitude is being inferred as expected. First, generate a trace plot of the absolute magnitude samples (the first entry in `g_samples`), overlaying the ground truth. Then print out the mean and standard deviation of the marginalized absolute magnitude posterior. Recall that marginalizing is as simple as throwing away the samples of all other parameters.<jupyter_code>mp.plot(g_samples[0])
mp.axhline(abs_true)
mp.xlabel('sample')
mp.ylabel(r'$M^*$')
print('Truth {:6.2f}; inferred {:6.2f} +/- {:4.2f}'.format(abs_true, np.mean(g_samples[0]), np.std(g_samples[0])))<jupyter_output>Truth -30.95; inferred -30.97 +/- 0.02<jupyter_text>Now let's generate some marginalized parameter posteriors (by simply discarding all samples of the latent parameters) using DFM's [`corner`](https://corner.readthedocs.io/en/latest/) package. Note the near identical nature of this plot to the `schmorner` plot we generated above.<jupyter_code>import corner
samples = np.stack((g_samples[0], g_samples[1]))
fig = corner.corner(samples.T, labels=[r"$M^*$", r"$s$"],
                    show_titles=True, truths=[abs_true, s_true])<jupyter_output><empty_output><jupyter_text>Task 4The final task is to write a [Stan model](https://pystan.readthedocs.io/en/latest/getting_started.html) to infer the parameters of the period-luminosity relation. I've coded up the other two blocks required (`data` and `parameters`), so all that is required is for you to write the joint posterior (factorized into its individual components) in Stan's sampling-statement-based syntax. Essentially all you need are Gaussian sampling statements (`abs_true ~ normal(abs_bar, abs_sig);`) and for loops (`for(i in 1: n_gal){...}`).When you evaluate this cell, Stan will translate your model into `c++` code and compile it. We will then pickle the compiled model so you can re-use it rapidly without recompiling. To do so, please set `recompile = False` in the notebook.<jupyter_code>import sys
import pystan as ps
import pickle

stan_code = """
data {
    int<lower=0> n_gal;
    int<lower=0> n_star;
    real mu_hat;
    real mu_hat_sig;
    real m_hat[n_gal, n_star];
    real m_hat_sig;
    real m_sig_int;
    real lp_true[n_gal, n_star];
    real abs_bar;
    real abs_sig;
    real s_bar;
    real s_sig;
    real mu_bar;
    real mu_sig;
}
parameters {
    real mu_true[n_gal];
    real m_true[n_gal, n_star];
    real abs_true;
    real s_true;
}
model {
    // priors
    abs_true ~ normal(abs_bar, abs_sig);
    s_true ~ normal(s_bar, s_sig);
    mu_true ~ normal(mu_bar, mu_sig);
    
    // whatevers
    for(i in 1: n_gal){
        for(j in 1: n_star){
            m_true[i, j] ~ normal(mu_true[i] + abs_true + s_true * lp_true[i, j], m_sig_int);
        }
    }
    
    // likelihoods
    mu_hat ~ normal(mu_true[1], mu_hat_sig);
    for(i in 1: n_gal){
        for(j in 1: n_star){
            m_hat[i, j] ~ normal(m_true[i, j], m_hat_sig);
        }
    }
}
"""
n_samples_stan = 5000
recompile = True
pkl_fname = 'bhms_stan_model_v{:d}p{:d}p{:d}.pkl'.format(sys.version_info[0], \
                                                         sys.version_info[1], \
                                                         sys.version_info[2])
if recompile:
    stan_model = ps.StanModel(model_code=stan_code)
    with open(pkl_fname, 'wb') as f:
        pickle.dump(stan_model, f)
else:
    try:
        with open(pkl_fname, 'rb') as f:
            stan_model = pickle.load(f)
    except EnvironmentError:
        print('ERROR: pickled Stan model (' + pkl_fname + ') not found. ' + \
              'Please set recompile = True')<jupyter_output>INFO:pystan:COMPILING THE C++ CODE FOR MODEL anon_model_ab41b39e55c2f57c74acf30e86ea4ea5 NOW.<jupyter_text>Now let's sample...<jupyter_code>stan_data = {'n_gal': n_gal, 'n_star': n_star, 'mu_hat': mu_hat, 'mu_hat_sig': mu_hat_sig, \
             'm_hat': m_hat, 'm_hat_sig': m_hat_sig, 'm_sig_int': m_sig_int, 'lp_true': lp_true, \
             'abs_bar': abs_bar, 'abs_sig': abs_sig, 's_bar': s_bar, 's_sig': s_sig, \
             'mu_bar': mu_bar, 'mu_sig': mu_sig}
fit = stan_model.sampling(data=stan_data, iter=n_samples_stan, chains=4)<jupyter_output><empty_output><jupyter_text>... print out Stan's posterior summary (note this is for _all_ parameters)...<jupyter_code>samples = fit.extract(permuted=True)
print(fit)<jupyter_output><empty_output><jupyter_text>... and plot the marginalized posterior of the PL parameters, as with the Gibbs sampler.<jupyter_code>c_samples = np.stack((samples['abs_true'], samples['s_true']))
fig = corner.corner(c_samples.T, labels=[r"$M^*$", r"$s$"],
                    show_titles=True, truths=[abs_true, s_true])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12415173558560758, 'fraction_numerical': 0.01718498266855441, 'mean_word_length': 3.093525179856115, 'pattern_counts': {'":': 0, '<': 60, '<?xml version=': 0, '>': 60, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '1162', 'license': 'MIT', 'n_tokens_mistral': 7842, 'n_tokens_neox': 7457, 'n_words': 1993, 'path': 'bhms.ipynb', 'repo_name': 'sfeeney/bhm_lecture', 'size': '20496'}