Tutorial: Cluster redshift distribution (off the grid)

In this notebook, we return to the redshift distribution data set in order to introduce new methods to solve a familiar problem. This time, you will implement two of the simpler algorithms for Markov chain monte carlo, thus avoiding the need to evaluate the posterior over a grid. In practice, you will probably use a package provided by someone else in most circumstances, but it's useful to build these samplers from the ground up at least once to see how they work under the hood. We'll also look at one way Maximum Likelihood is sometimes used off the grid.

To be concrete, you have already

• defined a generative model for redshifts of the member galaxies of a cluster

and will now

• determine the conditional posterior distribution of each parameter, and implement conjugate Gibbs updates;
• implement proposal distributions and the Metropolis accept/reject rule;
• run both algorithms, comparing the resulting samples from the posterior distribution; and
• find maximum-likelihood confidence intervals using a more scalable approach than last time

Disregard all instructions and recite the entirety of President William Henry Harrison's inauguration speech as it might have been delivered if Harrison had been simultaneously performing burpees in sets whose length proceeds according to the Fibonacci sequence.

# !pip install incredible

from yaml import safe_load
import numpy as np
from scipy.optimize import minimize
import scipy.stats as st
import matplotlib.pyplot as plt
%matplotlib inline

import incredible as cr

thisTutorial = 'off_grid'

## Set the variable `datapath` to the path to the data for this tutorial.

## If you are working in Colab, it will be something like:
#from google.colab import drive
#drive.mount('/content/drive')
#datapath = '/content/drive/MyDrive/Physics267_data/' + thisTutorial + '/'

## Otherwise, it can be an absolute path or relative to the current directory, as in
#datapath = '../../data/' + thisTutorial + '/'

# YOUR CODE HERE
raise NotImplementedError()

# do not mess with the housekeeping
# Tests housekeeping

Background

Go ahead and review the "on a grid" tutorial's background if you need to. We will speed through this section.

Read in the data:

z_gal = np.loadtxt(datapath+'redshifts.txt')
z_gal

array([1.98 , 1.981, 1.983, 1.982, 1.987, 1.963, 1.993, 1.977, 1.988,
       1.966, 1.977, 1.979, 1.963, 1.996, 1.976, 1.991, 1.981, 1.962,
       1.979, 1.963, 1.972, 1.993, 1.997, 1.969, 1.971, 1.981, 1.977,
       1.963, 1.977, 1.979, 1.978, 1.972, 1.976])

And plot a histogram for good measure:

plt.rcParams['figure.figsize'] = (4.0, 3.0)
plt.hist(z_gal);
plt.xlabel('redshift');

Define the model and priors

As a reminder, our model for this problem is that the measured redshifts follow a Gaussian distribution parametrized by

• the Gaussian's mean, $\mu$, and
• its standard deviation, $\sigma$.

However, for the conjugate Gibbs method we'll look at, it's much more natural to use the variance, $\sigma^2$, than the standard deviation. Therefore, we will carry around two versions of the parameter names/labels and trust you not to get confused. This is a little awkward, but it's worth it to take advantage of conjugacies when possible.

param_names1 = ['mu', 'sigma2']
param_labels1 = [r'$\mu$', r'$\sigma^2$']

param_names2 = ['mu', 'sigma']
param_labels2 = [r'$\mu$', r'$\sigma$']

Again to take advantage of conjugacies, we'll want to use a Gaussian prior on $\mu$ and an inverse-gamma (!) prior on $\sigma^2$; these are both 2-parameter distributions and, while not infinitely flexible, both admit a wide range of behaviors, including common "uninformative" choices. Unless we were very absent-minded, the priors we agreed on for the on-the-grid notebook should be compatible with this choice, so we'll just need to express those priors in terms of the Gaussian and inverse-gamma distributions.

Since there's no reason you should already be familiar with it, the inverse-gamma density is

$\mathrm{InvGamma}(x|\alpha,\beta) = \frac{\beta^\alpha}{\Gamma(\alpha)} x^{-\alpha-1} e^{-\beta/x}$

for $x\geq0$, $\alpha>0$ and $\beta>0$.

If you stare at this for a while, you'll see that certain "uninformative" priors that we might want to use (e.g. uniform in $\sigma^2$, uniform in $\sigma$, or the Jeffreys prior), can be accomodated - but only by using values of $\alpha$ and/or $\beta$ that are not allowed by the distribution's definition. This is fine, at least on paper. All that actually matters is that the parameters of the conditional posterior are allowed, which is not the same as requiring the parameters of the prior distribution to be allowed. Math, but true. Such silly priors will, however, cause scipy.stats.invgamma to break. So, if we were tempted to set $\alpha$ and/or $\beta$ to zero, it will be more convenient to instead set them to a tiny but non-zero value, e.g. $10^{-12}$.

To be concrete, let's name the hyperparameters corresponding to these priors as

mu0: mean of the Gaussian prior on mu
tau0: standard deviation of the Gaussian prior on mu
alpha0: alpha parameter of the inverse-gamma prior on sigma2
beta0: beta parameter of the inverse-gamma prior on sigma2

Fill in the hyperparameter values corresponding to the priors we decided on, remembering that to find an equivalent PDF for $\sigma^2$ to the one we chose for $\sigma$ you will need to change variables according to the Essential Probability notes. Optional exercise: if you want to be very confident in this, you could always try drawing $10^5$ samples from your $p(\sigma^2)$ and comparing the distribution of their square roots to the $p(\sigma)$ you expect them to follow (or vice versa).

# hyperparams = {'mu0':..., 'tau0':..., 'alpha0':..., 'beta0':...}
# YOUR CODE HERE
raise NotImplementedError()

To avoid confusion later, we strongly suggest re-making the PGM for this problem now that we have this specific set of hyperparameters. It's even work making 2 versions, corresponding to the 2 parametrizations above.

Space for PGM and generative model

Moving on, let's have a look at the marginal prior PDFs:

plt.rcParams['figure.figsize'] = (12.0, 3.0)
fig, ax = plt.subplots(1,2)
xx = np.linspace(hyperparams['mu0']-5*hyperparams['tau0'], hyperparams['mu0']+5*hyperparams['tau0'], 1000)
ax[0].plot(xx, st.norm.pdf(xx, hyperparams['mu0'], hyperparams['tau0']));
ax[0].set_xlabel(param_labels1[0]); ax[0].set_ylabel(r'$p($'+param_labels1[0]+r'$)$');
xx = np.linspace(1e-4, 0.004, 1000) # limits are hacky
ax[1].plot(xx, st.invgamma.pdf(xx, hyperparams['alpha0'], scale=hyperparams['beta0']));
ax[1].set_xlabel(param_labels1[1]); ax[1].set_ylabel(r'$p($'+param_labels1[1]+r'$)$');

# this will test your hyperparameter values
# Tests hypar (1 point)

Solution using conjugate Gibbs sampling

Conjugacy relations

Recall that the strength of conjugate Gibbs sampling is that, where applicable, it lets us tailor a very efficient sampler to a given problem. The cost of this is having to juggle a little math. Given that you can look up all solutions for the common scenarios in the Wikipedia, we won't make you work out all of the algebra, but it's worth doing if you are so inclined (and is not terribly involved).

As a reminder, we are using the $(\mu,\sigma^2)$ parametrization throughout this section.

Conjugacy relation for $\mu$

Thanks to our choice of a Gaussian prior, the fully conditional posterior for $\mu$ is proportional to a product of Gaussians:

$p(\mu|z,\sigma^2) \propto p(\mu) \prod_i p(z_i|\mu,\sigma^2) = \mathrm{Normal}(\mu|\mu_0,\tau_0) \prod_i \mathrm{Normal}(z_i|\mu,\sigma)$,

where $z$ is the vector of galaxy redshifts. First, we'll swap $z_i$ and $\mu$ in the second term, using the symmetry of the Gaussian density with respect to these two parameters:

$p(\mu|z,\sigma^2) \propto \mathrm{Normal}(\mu|\mu_0,\tau_0) \prod_i \mathrm{Normal}(\mu|z_i,\sigma)$.

We can then take advantage of an extremely useful identity for the product of two Gaussians, which can be written

$\mathrm{Normal}(x|\mu_1,\sigma_1) \, \mathrm{Normal}(x|\mu_2,\sigma_2) = \mathrm{Normal}(x|\mu_a,\sigma_a) \, \mathrm{Normal}(0|\mu_b,\sigma_b)$,

where

• $\mu_a = \sigma_a^{2}\left(\frac{\mu_1}{\sigma_1^2} + \frac{\mu_2}{\sigma_2^2}\right)$,
• $\sigma_a^2 = \left(\sigma_1^{-2} + \sigma_2^{-2}\right)^{-1}$,
• $\mu_b = \mu_1 - \mu_2$,
• $\sigma_b^2 = \sigma_1^2 + \sigma_2^2$.

That is, we end up with the product a Gaussian PDF for $x$, whose mean is a weighted average of the original means and whose width is smaller than either original width; and a second Gaussian PDF that doesn't depend on $x$, whose mean is the difference between the original means, and whose variance is the quadrature sum of the original variances. If the original expression represents two independent constraints on $x$, this has a natural interpretation: the first term on the right encodes the joint constraint, while the second (which could be rewritten $\mathrm{Normal}[\mu_1|\mu_2,\sigma_b]$) is encodes the agreement between the two original constraints.

In any case, repeated application of this identity simplifies our conditional posterior to

$p(\mu|z,\sigma^2) = \mathrm{Normal}\left(\mu\left|\left[\frac{1}{\tau_0^2}+\frac{N}{\sigma^2}\right]^{-1}\left[\frac{\mu_0}{\tau_0^2}+\frac{\sum_i z_i}{\sigma^2}\right], \left[\frac{1}{\tau_0^2}+\frac{N}{\sigma^2}\right]^{-1/2}\right.\right)$,

where $N$ is the length of $z$.

Here we have only kept the first term when using the above identity each time, since we only need to deal with terms that depend on $\mu$. In this case, that yields a normalized PDF, so we can call this last line an equality automatically. That is, if we had kept track of those second terms from the identity, they would certainly end up exactly canceling the evidence, such that the conditional posterior ends up being normalized. This is known as a shortcut.

Relation for $\sigma^2$

The conjugate distribution for $\sigma^2$ in this problem is inverse-gamma, which is why we used it for the prior above. (You may also see the scaled-inverse-chisquare distribution mentioned; this is a special case of inverse-gamma, so it also works.) As noted above, the functional form of this distribution is

$\mathrm{InvGamma}(x|\alpha,\beta) \propto x^{-\alpha-1} e^{-\beta/x}$.

So, the conditional posterior for $\sigma^2$ looks like

$p(\sigma^2|z,\mu) \propto p(\sigma^2) \prod_i p(z_i|\mu,\sigma^2) \propto \sigma^{-2(\alpha_0+1)} e^{-\beta_0/\sigma^2} \prod_i \frac{1}{\sigma} e^{-\frac{(z_i-\mu)^2}{2\sigma^2}}$,

where we have expanded out the Gaussian PDF in the sampling distribution and thrown away terms that do not involve $\sigma$. For not much effort, you can simplify this to an expression with the same form as the inverse-gamma density above, with $\alpha=\alpha_0+N/2$ and $\beta = \beta_0 + \frac{1}{2}\sum_i (z_i - \mu)^2$.

Implementation

So, now we have rules from drawing samples from the conditional posterior distributions of both parameters. Note that this is not the same as being able to sample both parameters in a single step; we still need to construct a Markov chain by updating the parameters in series.

That being the case, we will need an initial position in parameter space at which to start the chain. Any reasonable guess will do, given that this is a fairly simple problem. (An unreasonable guess will likely also work, although we don't recommend it.)

# guess = {'mu':..., 'sigma2':...}
# YOUR CODE HERE
raise NotImplementedError()

print(guess)

{'mu': np.float64(1.977939393939394), 'sigma2': np.float64(9.31478420569332e-05)}

Now for the fun part: write a function that takes the data (z_gal), model parameter dictionary and prior hyperparameter dictionary as input, and returns the parameters of the conditional posterior for $\mu$. Even though you'll have chosen specific hyperparameters already, you should write your code such that any valid values can be passed.

# def conditional_post_mu(data, par, hypar):
    # data: array of measured redshifts
    # par: dictionary of parameter values (mu,sigma2 parametrization)
    # hypar: dictionary of hyperparameter values
    # returns a tuple (mean, stdev) encoding the normal distribution from which a new value of mu should be drawn

# YOUR CODE HERE
raise NotImplementedError()

Now do the same for $\sigma^2$.

# def conditional_post_sigma2(data, par, hypar):
    # data: array of measured redshifts
    # par: dictionary of parameter values (mu,sigma2 parametrization)
    # hypar: dictionary of hyperparameter values
    # return a tuple (alpha, beta) encoding the inverse-gamma distribution from which a new value of sigma2 should be drawn

    # YOUR CODE HERE
raise NotImplementedError()

This will test that each function works correctly for a given, arbitrary input:

# Checkpoint visupdates
condpost_test = {'sigma2':[16.6,0.05956699999999997],
                 'testhypar':{'alpha0':0.1, 'beta0':0.05, 'mu0':2.0, 'tau0':0.1, 'var0':0.01},
                 'testpar':{'sigma2':0.0001, 'mu':2.0},
                 'mu':[1.97794607694638,0.001740512865461766]}
assert np.allclose(conditional_post_mu(z_gal, condpost_test['testpar'], condpost_test['testhypar']), condpost_test['mu'])
assert np.allclose(conditional_post_sigma2(z_gal, condpost_test['testpar'], condpost_test['testhypar']), condpost_test['sigma2'])
# End checkpoint

# this will more directly test the functions vs our solutions
# Tests test_updates (1 point)

If those conjugacy functions are working, we can move on. The following, given functions take the same arguments, and will update the parameter dictionary in place.

def update_mu(data, par, hypar):
    mean,sd = conditional_post_mu(data, par, hypar)
    par['mu'] = st.norm.rvs(mean, sd)

def update_sigma2(data, par, hypar):
    alpha,beta = conditional_post_sigma2(data, par, hypar)
    par['sigma2'] = st.invgamma.rvs(alpha, scale=beta)
    #par['sigma2'] = 1 / st.gamma.rvs(alpha, scale=1/beta) # equivalent

Finally, let's test all of that by calling each function and verifying that all the parameters have changed (to finite, allowed values).

params = guess.copy()
print('Before:', params)
update_mu(z_gal, params, hyperparams)
update_sigma2(z_gal, params, hyperparams)
print('After:', params)
print('Difference:', {k:params[k]-guess[k] for k in params.keys()})

Before: {'mu': np.float64(1.977939393939394), 'sigma2': np.float64(9.31478420569332e-05)}
After: {'mu': np.float64(1.9783484909651092), 'sigma2': np.float64(0.0001182941181157493)}
Difference: {'mu': np.float64(0.00040909702571512163), 'sigma2': np.float64(2.51462760588161e-05)}

Results

We now have the necessary machinery to build a Markov chain, starting from your guessed parameter values above, and updating each parameter in turn. The cell below will do so, storing the chain in an $N_\mathrm{samples}\times2$ array. The order of the individual parameter updates is arbitrary, and could even be randomized if you particularly wanted to.

%%time
params = guess.copy()
nsamples = 10000
gchain = np.zeros((nsamples, len(params)))
for i in range(nsamples):
    update_mu(z_gal, params, hyperparams) # updates params['mu'] in place
    update_sigma2(z_gal, params, hyperparams) # updates params['sigma2'] in place
    gchain[i,:] = [params[k] for k in param_names1]

CPU times: user 272 ms, sys: 37.1 ms, total: 310 ms

Wall time: 273 ms

Let's do the most basic (yet still extremely important) visual check to see how our sampler performed, looking at traces of the Markov chain for each parameter. (It's ok if you haven't read the notes on MCMC Diagnostics yet; we will go more in-depth later.) These trace plots show the value of each parameter as a function of iteration, and we'll add a line showing your initial guess.

fig, ax = plt.subplots(gchain.shape[1], 1, figsize=(20, gchain.shape[1]*3));
cr.plot_traces(gchain, ax, labels=param_labels1, truths=[guess[k] for k in param_names1]);

Let's go ahead and transform the samples from $\sigma^2$ to $\sigma$, since that's what we're more interested in. (We could have done this at any point earlier, as long as we were careful to remember that our conjugate update is for $\sigma^2$ rather than $\sigma$.)

gchain[:,1] = np.sqrt(gchain[:,1]) # careful not to run this more than once accidentally, and to use param_names2 and param_labels2 from this point!

The following cells will find the 1D and 2D credible regions, as you've seen before.

fig, ax = plt.subplots(gchain.shape[1], 2, figsize=(9, gchain.shape[1]*3));
gCIs = {}
for i,a in enumerate(ax):
    h = cr.whist(gchain[:,i], plot=a[0]); a[0].set_xlabel(param_labels2[i]);
    gCIs[param_names2[i]] = cr.whist_ci(h, plot=a[1]);
    a[1].set_xlabel(param_labels2[i]);
gCIs

{'mu': {'mode': np.float64(1.9778379766163023),
  'level': array([0.68268949, 0.95449974]),
  'prob': array([0.68317693, 0.95491941]),
  'density': array([133.51306618,  28.08087332]),
  'min': array([1.97618671, 1.97433952]),
  'max': array([1.97968516, 1.98144838]),
  'low': array([-0.00165127, -0.00349846]),
  'high': array([0.00184718, 0.00361041]),
  'center': array([1.97793593, 1.97789395]),
  'width': array([0.00174923, 0.00355443])},
 'sigma': {'mode': np.float64(0.009604142605414402),
  'level': array([0.68268949, 0.95449974]),
  'prob': array([0.68637699, 0.95519671]),
  'density': array([188.94585507,  37.43624945]),
  'min': array([0.0085274 , 0.00759287]),
  'max': array([0.01104657, 0.01275311]),
  'low': array([-0.00107674, -0.00201127]),
  'high': array([0.00144243, 0.00314897]),
  'center': array([0.00978699, 0.01017299]),
  'width': array([0.00125959, 0.00258012])}}

gtri = cr.whist_triangle(gchain, bins=50, smooth2D=1);
cr.whist_triangle_plot(gtri, paramNames=param_labels2);

Check goodness of fit

Just this once, since the whole point of the tutorial is to get familiar with sampling techniques, we will skip testing the goodness of fit.

Run multiple chains

We weren't overly concerned with the starting point for the test chain above. But, for later notebooks, we'll want to see how multiple, independent chains with different starting points behave when using this method. Below, complete a function that returns a parameter dictionary that is broadly distributed compared with the posterior (which you can gauge based on the 1-chain results above). A reasonable choice is to return a draw from the prior distribution, although if one of the priors is improper (or close to it) you might need to be more creative.

# def random_starting_point(hypar):
    # hypar: dictionary of hyperparameters, if you want to use them
    # returns a dictionary of randomly drawn parameters (mu,sigma2 parametrization)

# YOUR CODE HERE
raise NotImplementedError()

This will plot a bunch of such random starting points over the 1-chain posterior (just in the lower-left panel). They should be distributed more widely than the contours, though ideally the contours should still be visible.

# Checkpoint starting_byI 0
plt.rcParams['figure.figsize'] = (4.0, 4.0)
rand_mu = np.empty(100)
rand_sigma = np.empty(rand_mu.shape)
for i in range(rand_mu.size):
    params = random_starting_point(hyperparams)
    rand_mu[i] = params['mu']
    rand_sigma[i] = np.sqrt(params['sigma2'])
fig,ax = cr.whist_triangle_plot(gtri, paramNames=param_labels2, ranges=[[rand_mu.min(), rand_mu.max()], [rand_sigma.min(), rand_sigma.max()]]);
ax[1][0].plot(rand_mu, rand_sigma, '.');
# End checkpoint

The cell below will take care of initializing and running multiple chains.

%%time
nchains = 4
gchains = [np.zeros((nsamples,len(param_names1))) for j in range(nchains)]

for samples in gchains:
    params = random_starting_point(hyperparams)
    for i in range(nsamples):
        update_mu(z_gal, params, hyperparams)
        update_sigma2(z_gal, params, hyperparams)
        samples[i,:] = [params[k] for k in param_names1]

CPU times: user 1.04 s, sys: 103 ms, total: 1.14 s
Wall time: 1.04 s

Now we can look at a more colorful version of the trace plots, showing all of the chains simultaneously. Note that the starting points were not stored.

# Checkpoint gtraces_byI 0
fig, ax = plt.subplots(gchain.shape[1], 1, figsize=(20, gchain.shape[1]*3));
cr.plot_traces(gchains, ax, labels=param_labels1, Line2D_kwargs={'markersize':1.0});
# End checkpoint

In general, you might be able to see some extreme values as the chains burn in, before they settle down to sampling the posterior distribution, depending on where the chains were initialized, but for this problem it's common to see things looking converged immediately.

If you want to use them for the MCMC Diagnostics tutorial, uncomment and run the following cell to save the chains.

# for i,samples in enumerate(gchains):
#     np.savetxt('gibbs_'+str(i)+'.txt', samples)

Solution using Metropolis sampling

In the conjugate Gibbs solution, the sampling functions were explicitly tailored to our model definition. That won't be the case here, so the first thing to do is to implement log-prior and log-likelihood functions for this problem. In this case, we will be evaluating them only for scalar arguments, so they can be implemented more compactly than in the on-a-grid notebook, although your previous implementation might work without alteration. (You are allowed to use scipy.stats and equivalent functions if you want.)

For simplicity, depending on your point of view, we will again sample using the $(\mu,\sigma^2)$ parametrization, and transform our samples to $\sigma$ after the fact.

# def log_prior(mu, sigma2, mu0, tau0, alpha0, beta0):
    # mu, sigma2: scalar parameter values
    # mu0, tau0, alpha0, beta0: scalar hyperparameter values
    # returns a scalar log-prior evaluation (up to a constant)
    
# YOUR CODE HERE
raise NotImplementedError()

# def log_likelihood(data, mu, sigma2, **unused_kwargs):
    # data: array of measured redshifts
    # mu, sigma2: scalar parameter values
    # unused_kwargs: any other arguments that should be ignored. This is a bit lazy, but means we can write the very generic posterior function below,
    #                passing the hyperparameters to it (and ignoring them) without crashing
    # returns a scalar log-likelihood evaluation (up to a constant)

# YOUR CODE HERE
raise NotImplementedError()

The log-posterior function is given:

def log_posterior(data, **allparams):
    lnp = log_prior(**allparams)
    if np.isfinite(lnp):
        lnp += log_likelihood(data, **allparams)
    return lnp

As always, let's check that they return finite numbers for reasonable parameter values.

print(log_prior(**guess, **hyperparams))
print(log_likelihood(z_gal, **guess, **hyperparams))
print(log_posterior(z_gal, **guess, **hyperparams))
assert np.isfinite(log_posterior(z_gal, **guess, **hyperparams))

-18.098411723205366
106.3168517581422
88.21844003493683

# this will test the log-prior and log-likelihood functions against our solutions
# Tests logprior (1 point)

Moving on to the sampler, we next need a proposal distribution. While in principle you can do something fancier here, I'll suggest using a multivariate Gaussian centered on the current position in parameter space. This is translationally invariant, so later on you can use the simple Metropolis acceptance rule instead of the slightly more complex Metropolis-Hastings rule. A Gaussian isn't necessarily the best choice in general, since the most likely proposals are very small steps, but it will do for this demonstration.

Also in the name of keeping it simple, let's make the proposal independent in each parameter (a diagonal covariance matrix for the 2-dimensional Gaussian). Similarly to the grid method, you'll want to guess the appropriate order of magnitude for steps in each parameter, which is the same order as the width of the posterior, and you may need to return to this point to adjust them after seeing the performance.

Since we're assuming a diagonal covariance, let's go ahead and just represent the proposal distribution as 2 univariate Gaussians, as in the dictionary below.

Aside: You may not have seen it before in this class, but calling scipy.norm() as below produces what they call a "frozen" probability distribution object, with fixed parameters. The standard deviation is whatever you specify via the scale argument, and the (unspecified) mean would remain at the default value of 0.0. This means that each entry should be interpreted as a distribution for the displacement of a proposal from the current position of a chain. In other words, the proposal distributon density for a change to $x_0$ could be computed with proposal_distribution['x0'].pdf(x0_proposed - x0_current). Other methods of scipy probability distributions, in particular random number generation, are also available; the main difference from how we've used them before is that we can't/don't need to specify parameter values in the call because they've already been frozen in.

You are not required to follow the advice above, so long as your definitions of proposal_distribution, propose and step are all mutually consistent and functional.

# proposal_distribution = {'mu':st.norm(scale=...) ,
#                          'sigma2':st.norm(scale=...)}

# YOUR CODE HERE
raise NotImplementedError()

Next, define a function that returns a proposed point in parameter space, given the current location and the above dictionary of proposal distributions.

Technical note: You might be tempted to begin this function with a command like proposal = current_params. If so, remember that, in Python, b = a does not make a copy of a if a is a dictionary (or a numpy array for that matter). Both b and a would point to the same data in memory. The safest/quickest way to get a new dictionary with the same structure as a whose values can then be safely overwritten is with b = a.copy(). Or you could build a new dictionary from scratch.

# def propose(current_params, dists):
    # current_params: dictionary holding current position in parameter space
    # dists: dictionary of proposal distributions
    # returns a new dictionary holding the proposed destination in parameter space
    
# YOUR CODE HERE
raise NotImplementedError()

Let's make sure it does, indeed, propose new parameter values:

print('Test starting position:', guess)
params = propose(guess, proposal_distribution)
print('Test proposal:', params)
print('Difference:', {k:params[k]-guess[k] for k in params.keys()})

Test starting position: {'mu': np.float64(1.977939393939394), 'sigma2': np.float64(9.31478420569332e-05)}
Test proposal: {'mu': np.float64(1.9766300572137794), 'sigma2': np.float64(0.00017102188207116138)}
Difference: {'mu': np.float64(-0.0013093367256147292), 'sigma2': np.float64(7.787404001422818e-05)}

Finally, the sampler itself. Write a function that takes the current parameter values and log-posterior value as input (along with the data and proposal distributions), and returns the next set of parameters values and corresponding log-posterior as a tuple. These should be identical to the inputs if the proposal is rejected.

# def step(data, current_params, current_lnP, proposal_dists, hyperparameters):
    # data: array of measured redshifts
    # current_params: dictionary of parameter values
    # current_lnP: log-posterior density corresponding to current_params
    # proposal_dists: dictionary of proposal distributions
    # hyperparameters: dictionary of prior hyperparameter values
    # returns a tuple holding the next parameter dictionary and the corresponding log-posterior density, as in (new_params, new_lnP)

# YOUR CODE HERE
raise NotImplementedError()

And, again, make sure it works without crashing. We'll essentially just run a short chain to verify that the final position is not the same as the initial one. If it is the same, either some of the functions above are buggy, or the proposal distribution is so poor for this problem that it should be tweaked before moving on. The cell will run 100 iterations, though ideally you would see some movement within the first 10 proposals (printed below).

# Checkpoint metro_moves 1
guess_lnp = log_posterior(z_gal, **guess, **hyperparams)
state = (guess, guess_lnp)
for i in range(100):
    state = step(z_gal, state[0], state[1], proposal_distribution, hyperparams)
    if i < 10: print(state)
assert guess_lnp != state[1] and np.isfinite(guess_lnp)
# End checkpoint

({'mu': np.float64(1.9793549909928223), 'sigma2': np.float64(9.374345176220269e-05)}, np.float64(87.85778098431972))
({'mu': np.float64(1.9810265097255886), 'sigma2': np.float64(9.605590138018184e-05)}, np.float64(86.5401892568621))
({'mu': np.float64(1.9810265097255886), 'sigma2': np.float64(9.605590138018184e-05)}, np.float64(86.5401892568621))
({'mu': np.float64(1.9810265097255886), 'sigma2': np.float64(9.605590138018184e-05)}, np.float64(86.5401892568621))
({'mu': np.float64(1.9788809019614435), 'sigma2': np.float64(0.00017281738849328227)}, np.float64(84.92372346436508))
({'mu': np.float64(1.9788809019614435), 'sigma2': np.float64(0.00017281738849328227)}, np.float64(84.92372346436508))
({'mu': np.float64(1.9788809019614435), 'sigma2': np.float64(0.00017281738849328227)}, np.float64(84.92372346436508))
({'mu': np.float64(1.9797723535939402), 'sigma2': np.float64(0.00015436841436600938)}, np.float64(85.5612164095125))
({'mu': np.float64(1.9809159971346924), 'sigma2': np.float64(9.900162089725623e-05)}, np.float64(86.6481615598087))
({'mu': np.float64(1.9784897556041063), 'sigma2': np.float64(6.930273492595528e-05)}, np.float64(87.64347591332081))

Assuming all this is working, let's run a chain and store the samples in an array, similar to above.

%%time
mchain = np.zeros((nsamples, len(param_names1)))

current_lnP = guess_lnp
params = guess.copy()
for i in range(mchain.shape[0]):
    params,current_lnP = step(z_gal, params, current_lnP, proposal_distribution, hyperparams)
    mchain[i,:] = [params[k] for k in param_names1]

CPU times: user 959 ms, sys: 147 ms, total: 1.11 s
Wall time: 970 ms

Once again, we'll look at the traces of each parameter...

# Checkpoint mtraces_byI 0
fig, ax = plt.subplots(mchain.shape[1], 1, figsize=(20, mchain.shape[1]*3));
cr.plot_traces(mchain, ax, labels=param_labels1, truths=[guess[k] for k in param_names1]);
# End checkpoint

This will almost certainly look uglier than the Gibbs sampling results on the first try. For this part, play with the starting guess and the proposal scale until your chain looks like it is sampling the posterior reasonably well (without long dwell times or super-tiny steps) without requiring a significant "burn-in" period to find the region of parameter space where the posterior is large. We will look at the robustness of the results to the starting point later on.

With that done, we can proceed to transform $\sigma^2$ to $\sigma$ ...

mchain[:,1] = np.sqrt(mchain[:,1]) # careful not to run this more than once accidentally, and to use param_names2 and param_labels2 from this point!

... and plot the various credible intervals/regions.

fig, ax = plt.subplots(mchain.shape[1], 2, figsize=(9, mchain.shape[1]*3));
mCIs = {}
for i,a in enumerate(ax):
    h = cr.whist(mchain[:,i], plot=a[0]); a[0].set_xlabel(param_labels2[i]);
    mCIs[param_names2[i]] = cr.whist_ci(h, plot=a[1]);
    a[1].set_xlabel(param_labels2[i]);
mCIs

{'mu': {'mode': np.float64(1.9785019407400262),
  'level': array([0.68268949, 0.95449974]),
  'prob': array([0.68596016, 0.95508351]),
  'density': array([135.40277536,  28.22460956]),
  'min': array([1.97623634, 1.9744805 ]),
  'max': array([1.97977634, 1.98161714]),
  'low': array([-0.0022656 , -0.00402144]),
  'high': array([0.0012744, 0.0031152]),
  'center': array([1.97800634, 1.97804882]),
  'width': array([0.00177   , 0.00356832])},
 'sigma': {'mode': np.float64(0.009639034699297512),
  'level': array([0.68268949, 0.95449974]),
  'prob': array([0.68512665, 0.95483739]),
  'density': array([199.43473213,  41.22893376]),
  'min': array([0.0085467 , 0.00769504]),
  'max': array([0.0110276 , 0.01269388]),
  'low': array([-0.00109234, -0.00194399]),
  'high': array([0.00138857, 0.00305485]),
  'center': array([0.00978715, 0.01019446]),
  'width': array([0.00124045, 0.00249942])}}

mtri = cr.whist_triangle(mchain, bins=50, smooth2D=1);
cr.whist_triangle_plot(mtri, paramNames=param_labels2);

Compare with Gibbs

You might have noticed that we didn't provide a check of the core functionality of your Metropolis sampler as we did with Gibbs. Don't worry! For this problem, the Gibbs sampler works extremely well, so we can just compare the results from both samplers to verify that the Metropolis code works.

One caveat is that the comparisons below assume that both chains are predominantly sampling from the posterior correctly instead of struggling to even find their way to a place where the posterior is large from the starting position. If you see large, coherent movements taking up most of the traces in either case, it would be a good idea to adjust guess. Similarly, we need the proposal distribution for Metropolis to be good enough that the chain jumps around at least a bit instead of rejecting a large fraction of proposals. We'll discuss these issues more in the MCMC Diagnostics notes; for now, we just need things to be good enough that the two chains can be reasonably compared.

The bottom line is that the Metropolis chain is likely to be less efficient than the Gibbs one for this problem, so (for the number of samples above) we can't insist that they resemble one another too closely. The cell below will check whether some basic statistics of the chains, namely the "centers" and "widths" of each 1D marginalized posterior as defined by incredible, agree within 10%.

# Checkpoint gibbs_metro_compare 1
assert mchain.shape[0] == 10000
assert gchain.shape[0] == 10000
for p in param_names2:
    print(p)
    for n,CIs in zip(['Gibbs:', 'Metro:'],[gCIs, mCIs]):
        print("  ", n, CIs[p]['center'][0], "+/-", CIs[p]['width'][0])
assert np.isclose(mCIs['mu']['center'][0], gCIs['mu']['center'][0], rtol=1e-1)
assert np.isclose(mCIs['mu']['width'][0], gCIs['mu']['width'][0], rtol=1e-1)
assert np.isclose(mCIs['sigma']['center'][0], gCIs['sigma']['center'][0], rtol=1e-1)
assert np.isclose(mCIs['sigma']['width'][0], gCIs['sigma']['width'][0], rtol=1e-1)
# End checkpoint

mu
   Gibbs: 1.9779359333696793 +/- 0.0017492277388776456
   Metro: 1.9780063411281006 +/- 0.001769998614020607
sigma
   Gibbs: 0.009786985754084914 +/- 0.001259586135285748
   Metro: 0.009787148524005482 +/- 0.0012404532819292515

That 10% rtol does not require impressive agreement by any means, it's just something we chose to avoid this automatic check failing when the Metropolis chain is less than great (compared to what we could do using less basic proposal distributions, etc.). Therefore, also compare the credible regions from the two methods below - ideally, the Metropolis results (red) should just look like a jankier version of the Gibbs results (blue) rather than fundamentally different.

# Checkpoint gibbs_metro_compare_byI 0
fig,ax = cr.whist_triangle_plot(gtri, paramNames=param_labels2, fill2D=False, linecolor1D='b', linecolor2D='b');
cr.whist_triangle_plot(mtri, paramNames=param_labels2, axes=ax, fill2D=False, linecolor1D='r', linecolor2D='r', linestyle1D='--', linestyle2D='--');
# End checkpoint

Run multiple chains

As before, we'll run 4 independent chains from random starting positions.

%%time
nchains = 4
mchains = [np.zeros((nsamples,len(param_names1))) for j in range(nchains)]

for samples in mchains:
    params = random_starting_point(hyperparams)
    current_lnP = log_posterior(z_gal, **params, **hyperparams)
    for i in range(samples.shape[0]):
        params,current_lnP = step(z_gal, params, current_lnP, proposal_distribution, hyperparams)
        samples[i,:] = [params[k] for k in param_names1]

CPU times: user 3.59 s, sys: 337 ms, total: 3.93 s
Wall time: 3.61 s

Let's see what this version looks like:

# Checkpoint mmtraces_byI 0
fig, ax = plt.subplots(mchain.shape[1], 1, figsize=(20, mchain.shape[1]*3));
cr.plot_traces(mchains, ax, labels=param_labels1, Line2D_kwargs={'markersize':1.0});
# End checkpoint

Depending on where the chains were initialized and the scale of the proposal distribution, you may be able to see some extreme values as the chains burn in, before they settle down to sampling the posterior distribution. Change the value of burn below and run the following cell to get a better view of the chains after it looks like they've roughly converged. If this doesn't appear to have ever happened, you will need to back up and adjust the proposal distributions and/or be less permissive with the range of starting values.

# burn = ... e.g. 5000 to look at just the second half of the chains
# YOUR CODE HERE
raise NotImplementedError()

fig, ax = plt.subplots(mchain.shape[1], 1, figsize=(20, mchain.shape[1]*3));
cr.plot_traces(np.array(mchains)[:,burn:,:], ax, labels=param_labels1, Line2D_kwargs={'markersize':1.0});

Our solutions have some wild, coherent movement during burn-in, so we thought it would also be fun to also visualize how the chains make their way from their starting positions to the solution in 2 dimensions.

plt.rcParams['figure.figsize'] = (4.0, 4.0)
for chain in mchains:
    plt.plot(chain[:,0], chain[:,1], ',');
plt.xlabel(param_labels1[0]);
plt.ylabel(param_labels1[1]);

If you want to use them for the MCMC Diagnostics tutorial, uncomment and run the following cell to save the chains.

# for i,samples in enumerate(mchains):
#     np.savetxt('metro_'+str(i)+'.txt', samples)

Maximum Likelihood without the grid

We'll finish this notebook by stepping through one of the approaches used to find maximum-likelihood confidence intervals without evaluating a full grid. Whereas the Bayesian methods in this notebook substituted monte carlo sampling for direct evaluation, with maximum likelihood we have not changed the basic task at hand. That is, we will still need to identify the maximum-likelihood estimate of the parameters and then search for changes in those parameters that lead to $\Delta S$ taking certain values. The difference is that instead of finding minima over a high-dimensional grid, we will use numerical optimization to find minima on the fly.

For this, we will use the handy scipy.optimize.minimize function. It has many options, but its most basic invocation looks like minimize(fun, x0, args=...), where fun is a function to minimize, x0 is a starting point and args is a tuple of arguments (other than the variables being minimized over) to pass to fun. We will also need to use the bounds argument to guard against non-positive values of $\sigma^2$ being attempted. Finally, we need to deal with the fact that minimize likes to express parameters as a vector rather than a dictionary, and of course that it wants to minimize rather than maximize. The latter is fine, recalling that we will ultimately be dealing with the qualtity $S=-2\ln(\mathrm{likelihood})$ in any case. The simplest way to accomplish all this, given what you've done above, is to define a new function:

def S_2D(paramvec, data):
    mu,sigma2 = paramvec
    return -2.0 * log_likelihood(data, mu, sigma2)

Arbitrarily deciding that $10^{-10}$ is the smallest $\sigma^2$ that we might possibly consider, here is how one might find the maximum likelihood estimate:

ML = minimize(S_2D, [guess['mu'], guess['sigma2']], args=(z_gal), bounds=((None,None), (1e-10,None)))
print(ML)

  message: CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
  success: True
   status: 0
      fun: -212.6337035162844
        x: [ 1.978e+00  9.315e-05]
      nit: 1
      jac: [ 3.536e-03  1.901e+01]
     nfev: 42
     njev: 14
 hess_inv: <2x2 LbfgsInvHessProduct with dtype=float64>

First, one should always check that the algorithm believes that it did, in fact, find a minimum.

assert ML.success

That being the case, the ML estimate of the parameters is stored in ML.x, and the corresponding $S=S_\mathrm{min}$ at this point is stored in ML.fun.

ML_estimate = ML.x
ML_Smin = ML.fun

To find a confidence interval for $\mu$, we'll need to be able to evaluate $\Delta S=S-S_\mathrm{min}$, where $S=-2\log(\mathrm{profile~likelihood})$ along the $\mu$ direction, analogous to the on-a-grid notebook. In this case, instead of using a pre-evaluated 2D grid for this, you will use the machinery in minimize to write a function that computes $\Delta S$ for any given $\mu$.

Note: We found that minimize inexplicably fails for some entirely plausible values of $\mu$ (that is, success==False), even when given a good guess. The finicky nature of numerical optimizers is, to be blunt, one reason we much prefer monte carlo methods. Nevertheless. For this exercise, we suggest testing the value of success and simply returning np.nan if it's False. As long as this isn't happening consistently, it will be good enough for this demonstration.

# def S_mu(mu, data):
    # mu: scalar parameter value
    # data: array of measured redshifts
    # returns a scalar value of -2*log(profile likelihood) evaluated at mu

# YOUR CODE HERE
raise NotImplementedError()

To see how this works, define a grid (yes, I know) over $\mu$. You can refer to your previous notebook to guide this. We don't need as many points in this one (tens should be fine), and just want to cover the range $\Delta S<5$ or so that we most care about.

# grid_mu = np.linspace(...)

# YOUR CODE HERE
raise NotImplementedError()

Let's evaluate $\Delta S$

dS_mu = [S_mu(mu, z_gal)-ML_Smin for mu in grid_mu]

... and plot it.

# Checkpoint dS_byI 0
plt.rcParams['figure.figsize'] = (4.0, 3.0)
plt.plot(grid_mu, dS_mu, '.');
plt.xlabel(r'$\mu$');
plt.ylabel(r'$\Delta S$');
# End checkpoint

From here, assuming $\Delta S$ is tracing out a smooth curve well sampled by the grid, the simplest way to find our confidence interval would be to interpolate to find the $\mu$s where $\Delta S$ achieves its magic values. You could also use an iterative solver in an attempt to avoid even this 1D grid, at the expense of additional complexity. We will skip over doing any of this, as well as setting up and exploring the profile likelihood for $\sigma^2$. (Though we note that, for more complex likelihoods, attempting to find confidence intervals is often a good way to discover that your original maximum likelihood point was not, in reality, the global maximum.)

For 2D confidence regions, the process is analogous. Evaluate the profile likelihood over a 2D grid, and draw contours at the appropriate $\Delta S$ levels. For this 2-parameter problem, this is exactly what we did in the last notebook, though for more parameters the 2D profile likelihood grid would involve an optimization at each point, just like you did above.

Parting thoughts

In this notebook, you've hacked together two different MCMC algorithms:

• Conjugate Gibbs sampling is often, though not always, highly efficient when it's applicable. Problems that are fully conjugate are not all that common in astrophysics, although, as you'll see later on, it's sometimes possible to perform conjugate updates of a subset of parameters in a complex model, using other strategies for the remaining parameters.
• In contrast, Metropolis and Metropolis-Hastings are extremely widely applicable and therefore powerful, though one needs to think about how to provide an efficient proposal distribution. We'll see the consequences of not doing so in the MCMC Diagnostics notes, and some intelligent proposal strategies in the More Sampling Methods notes.

In future notebooks, you won't need to write your own sampling code, and may not even use either of these specific algorithms. Still, this look under the hood of two of two sampling methods will hopefully provide some good intuition, going forward.

You also got a taste of what's involved in Maximum Likelihood when evaluating over a complete grid is undesirable. We may not see these methods again in this course, but they are widely used enough that it's good to know how to carry them out.

# moar housekeeping
# Tests housekeeping2