Homework 4.1: Writing your own MCMC sampler (70 pts)

---

import numpy as np
import numba

a) Write your own MCMC sampler that employs a Metropolis-Hastings algorithm to sample continuous parameters (it does not need to handle discrete parameters) that uses a Normal proposal distribution. Since you are sampling multiple parameters, your proposal distribution will be multi-dimensional. You can use a Normal proposal distribution with a diagonal covariance. In other words, you generate a proposal for each variable in the target distribution independently.

You can organize your code how you like, but here is a suggestion.

• Write a function that takes a Metropolis-Hastings step. It should look something like the below (obviously where it does something instead of passing).

def mh_step(x, logtarget, logtarget_current, sigma, args=()):
    """
    Parameters
    ----------
    x : ndarray, shape (n_variables,)
        The present location of the walker in parameter space.
    logtarget : function
        The function to compute the log target. It has call
        signature `logtarget(x, *args)`.
    logtarget_current : float
        The current value of the log target.
    sigma : ndarray, shape (n_variables, )
        The standard deviations for the proposal distribution.
    args : tuple
        Additional arguments passed to `logtarget()` function.

    Returns
    -------
    new_x : ndarray, shape (n_variables,)
        The position of the walker after the Metropolis-Hastings
        step. If no step is taken, returns the inputted `x`.
    new_logtarget : float
        The updated log probability density of the target.
    accepted : bool
        True is the proposed step was taken and False otherwise.
    """

    pass

• Write another function that calls that function over and over again to do the sampling. It should look something like the below. (Note that I am using n_burn as opposed to n_warmup, because you are just going to step as you normally would, and then “burn” the samples. This is in contrast to Stan, which tunes the stepping procedure during the warm-up phase.)

def mh_sample(
    logtarget, x0, sigma, args=(), n_burn=1000, n_steps=1000, variable_names=None
):
    """
    Parameters
    ----------
    logtarget : function
        The function to compute the log target. It has call
        signature `logtarget(x, *args)`.
    x0 : ndarray, shape (n_variables,)
        The starting location of a walker in parameter space.
    sigma : ndarray, shape (n_variables, )
        The standard deviations for the proposal distribution.
    args : tuple
        Additional arguments passed to `logtarget()` function.
    n_burn : int, default 1000
        Number of burn-in steps.
    n_steps : int, default 1000
        Number of steps to take after burn-in.
    variable_names : list, length n_variables
        List of names of variables. If None, then variable names
        are sequential integers.

    Returns
    -------
    output : DataFrame
        The first `n_variables` columns contain the samples.
        Additionally, column 'lnprob' has the log target value
        at each sample.
    """

    pass

b) To test your code, we will get samples out of a known distribution. We will use a bivariate Normal with a mean of \(\boldsymbol{\mu} = (10, 20)\) and covariance matrix of

\begin{align} \mathsf{\Sigma} = \begin{pmatrix} 4 & -2 \\ -2 & 6 \end{pmatrix} \end{align}

I have written the log-PDF function to be unnormalized and JITted with numba for optimal speed.

Do not be confused: In this test function we are sampling \(\mathbf{x}\) out of \(P(\mathbf{x}\mid \boldsymbol{\mu},\mathsf{\Sigma})\). This is not sampling a posterior; it is just a test for your code. (Remember, MCMC samplers can sample out of arbitrary distributions, not just posteriors. In general, the distribution we sample out of is called a target distribution.) Furthermore, though both Normal, this test function is not the proposal distribution. You will pass log_test_distribution as the logtarget argument in the above functions.

mu = np.array([10.0, 20])
cov = np.array([[4, -2],[-2, 6]])
inv_cov = np.linalg.inv(cov)

@numba.njit
def log_test_distribution(x, mu, inv_cov):
    """
    Unnormalized log target of a multivariate Gaussian.
    """
    return -np.dot((x-mu), np.dot(inv_cov, (x-mu))) / 2

If you compute the means and covariance (using np.cov()) of your samples, you should come close to the inputted means and covariance. You might also want to plot your samples using bebi103.viz.corner() to make sure everything makes sense.

You may want to add in some logic to your Metropolis-Hastings sampler to enable tuning. This is the process of automatically adjusting the \(\sigma\) in the proposal distribution such that the acceptance rate is desirable. A good target acceptance rate is about 0.4. The developers of PyMC use the scheme below, which is reasonable.

Acceptance rate	Standard deviation adaptation
< 0.001	\(\times\) 0.1
< 0.05	\(\times\) 0.5
< 0.2	\(\times\) 0.9
> 0.5	\(\times\) 1.1
> 0.75	\(\times\) 2
> 0.95	\(\times\) 10

Be sure to test your code to demonstrate that it works.