1. Sampling with CUQIpy: five little stories

1.1. The two targets, again¶

Here we will add the code for building the two targets that we discussed in the previous notebook. The details of the two targets can be found in that notebook. After running the following (hidden) code cell, we will have the two targets ready for sampling:

• The donut target target_donut (distribution)

• The Poisson 1D target target_poisson (posterior)

In the following (hidden) code cell we define the method plot_pdf_2D for plotting bi-variate PDFs.

1.2. Sampling the posterior¶

After defining the posterior distribution (or a target distribution in general) for the parameters of interest $\mathbf{x}$, we can characterize the parameter and its uncertainty by samples from the posterior (target) distribution. However, in general the posterior (target) is not a simple distribution that we can easily sample from. Instead, we need to rely on Markov Chain Monte Carlo (MCMC) methods to sample from the posterior.

In CUQIpy, a number of MCMC samplers are provided in the sampler module that can be used to sample probability distributions. All samplers have the same signature, namely Sampler(target, ...), where target is the target CUQIpy distribution and ... indicates any (optional) arguments.

In CUQIpy we have a number of samplers available, that we can list using:

1.2.1. Story 1 - It pays off to adapt: Comparing Metropolis Hastings (MH) with and without adaptation ¶

• In the class, you learned about MH method.

• In MH, for each iteration, a new sample is proposed based on the current sample (essintially adding a random perturbation to the corrent sample).

• Then the porposed sample is accepted or rejected based on the acceptance probability: $\min\left(1, \frac{p(\mathbf{x}_{k+1})q(\mathbf{x}_k|\mathbf{x}_{k+1})}{p(\mathbf{x}_k)q(\mathbf{x}_{k+1}|\mathbf{x}_k)}\right)$, where $p(\mathbf{x})$ is the target distribution and $q(\mathbf{x}|\mathbf{y})$ is the proposal distribution.

• Here we test the MH method with and without adaptation.

• For adaptation, MH uses vanishing adaptation to set the step size (scale) to achieve a desired acceptance rate (0.234 in this case).

Let us sample the target_donut with a fixed scale. We first set up the number of samples, number of burn-in samples, and the fixed scale.

We create the MH sampler with the target distribution target_donut, the scale scale and an initial sample x0.

We sample using sample method:

The samples are stored in a Samples object that is returned by the sample method. Let us visualize the samples using the plot_pair method, along with the density of the target distribution using the plot_2D_density method:

Observations

• The samples captures only part of the target distribution.

• Once the chain reaches the high density region, it stays there.

• The samples starts at the center and takes some time to move to the donut shape. This part of the chain is the burn-in phase.

Now we try the solving the same problem with adaptation. We make sure to set the scale to the previous value 0.05.

Then we sample this time using sample_adapt method which adapts the step size. Note that the adaptation happen during the burin-in phase so we need to set Nb to some large value to see the effect of adaptation.

Here we visualize the results (both the ones with and without adaptation). We set the color of the samples and the transparency alpha using the arviz argument scatter_kwargs:

Let us plot the chains for the two cases, removing the burin-in from the fixed scale case:

Notes

• The chains in the second case (with adaptation) has much better mixing compared to the first case (closer to a fuzzy warm shape).

• The estimate density of the target distribution is much better in the second case, captures the 2 modes.

We can also use plot_pair with additional arviz arguments kind, kde_kwargs, and marginals to visualize the marginals and the kernel density estimate of the target distribution. Let us do this for the two cases:

It is clear from these figures how using adaptation of the step size is very beneficial to capture the target distribution in this scenario.

1.2.2. Story 2 - Not all unknowns are created equal: Component-wise Metropolis Hastings CWMH¶

To get to the next sample, CWMH updates the components of the current sample one at a time and applies the accept/reject step for each component. Updating the scale happens for each component separately depending on the acceptance rate of the component.

Let us test sampling the Poisson problem target with CWMH. We fix the seed for reproducibility and set the number of samples, burn-in samples, and the scale.

For comparison, we sample the Poisson problem target with the MH sampler first:

Then we create a CWMH sampler with the target distribution target_poisson, the scale scale and an initial sample x0 and sample from it.

Note that we use adaptation in both cases. Let us visualize the credible intervals in both cases: for MH

And for CWMH

1.2.3. Story 3 - May the force guide you: Unadjusted Langevin algorithm (ULA)¶

• In the Unadjusted Langevin algorithm (ULA), new states are proposed using (overdamped) Langevin dynamics.

• Let $p$ denote a probability density on $\mathbb{R}^d$, from which it is desired to draw an ensemble of i.i.d. samples. We consider the overdamped Langevin Itô diffusion

driven by the time derivative of a standard Brownian motion $W$ and force given as gradient of a potential $\log p(\boldsymbol{x}_t)$.

• In the limit, as $t\to\infty$, the probability distribution of $X(t)$ approaches a stationary distribution, which is also invariant under the diffusion. It turns out that this distribution is $p$.

• Approximate sample paths of the Langevin diffusion can be generated by the Euler--Maruyama method with a fixed time step $\varepsilon>0$. We set $\boldsymbol{x}_{0}$ and then recursively define an approximation to the true solution by

where each $\boldsymbol{\xi}_k\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I}_d)$.

Let us create the ULA sampler in CUQIpy with the target distribution target_donuts, a fixed scale scale, and an initial sample x0:

We set the number of samples, and sample the target distribution using the sample method:

And visualize the samples using the plot_pair method:

We note that the samples are all over the domain and do not capture the target distribution. This is because the ULA sampler will result in approximate (biased) distribution that as $\varepsilon \rightarrow 0$ will converge to $p(\mathbf{x})$.

1.2.4. Story 4: This is a bias we can fix - Metropolis-adjusted Langevin algorithm (MALA)¶

• In contrast to the Euler--Maruyama method for simulating the Langevin diffusion. MALA incorporates an additional step. We consider the above update rule as defining a proposal $\tilde{\boldsymbol{x}}$ for a new state:

this proposal is accepted or rejected according to the MH algorithm.

• That is, the acceptance probability is:

where the proposal has the form

• The combined dynamics of the Langevin diffusion and the MH algorithm satisfy the detailed balance conditions necessary for the existence of a unique, invariant, stationary distribution.

• For limited classes of target distributions, the optimal acceptance rate for this algorithm can be shown to be 0.574; this can be used to tune $\varepsilon$.

Let us get back to the uni-variate Gaussian target distribution x_uni that was suggested in the exercise of the previous section. We saw that when the scale was 1, the resulting samples did not capture the target distribution, but an appreciably different one, which shows the bias in ULA.

We will now use MALA to sample from this target distribution. Let us create x_uni first:

Then we create a MALA sampler object MALA_uni for the target distribution x_uni with scale 1 and initial sample 0 and sample from it:

Lastly, we visualize the KDE obtained from samples along with the target distribution PDF:

We notice in this case that the KDE of the samples captures the target distribution well. This is because MALA corrects the bias in ULA by using the MH acceptance step.

1.2.5. Story 5 - Sampling goes nuts with NUTS: Hamiltonian Monte Carlo (HMC) and No-U-Turn Sampler (NUTS)¶

• Hamiltonian Monte Carlo (HMC): use Hamiltonian dynamics to simulate particle trajectories.

• Define a Hamiltonian function in terms of the target distribution.

• Introduce an auxiliary momentum variables, which typically have independent Gaussian distributions.

• Hamiltonian dynamics operate on a $d$-dimensional position vector $\boldsymbol{x}$, and a $d$-dimensional momentum vector $\boldsymbol{r}$, so that the full state space has $2d$ dimensions. The system is described by a function of $\boldsymbol{x}$ and $\boldsymbol{r}$ known as the Hamiltonian $H(\boldsymbol{x}, \boldsymbol{r})$.

• In HMC, one uses Hamiltonian functions that can be written as (closed-system dynamics):

• The potential energy is completely determined by the target distribution, indeed $U(\boldsymbol{x})$ is equal to the logarithm of the target distribution $p$.

• The kinetic energy is unconstrained and must be specified by the implementation.

• The Hamiltonian is an energy function for the joint state of position-momentum, and so defines a joint distribution for them as follows:

• There are several ways to set the kinetic energy (density of the auxiliary momentum):

  • Euclidean--Gaussian kinetic energy: using a fixed covariance $\boldsymbol{M}$ estimated from the position parameters, $K(\boldsymbol{r},\boldsymbol{x}) = \dfrac{1}{2} \boldsymbol{r}^T\boldsymbol{M}^{-1}\boldsymbol{r} + \log( \left\vert\boldsymbol{M}\right\vert) + \text{const}$.

  • Riemann--Gaussian kinetic energy: unlike the Euclidean metric, varies as one moves through parameter space, $K(\boldsymbol{r}, \boldsymbol{x}) = \dfrac{1}{2} \boldsymbol{r}^T\boldsymbol{\Sigma}(\boldsymbol{x})^{-1}\boldsymbol{r} + \dfrac{1}{2}\log( \left\vert\boldsymbol{\Sigma}(\boldsymbol{x})\right\vert) + \text{const}$.

  • Non-Gaussian kinetic energies.

• Hamilton’s equations read as follows:

where $\dfrac{\partial U}{\partial \boldsymbol{x}} = \dfrac{\partial \text{log}p(\boldsymbol{x})}{\partial \boldsymbol{x}}$ is the gradient of the logarithm of the target density.

• Discretizing Hamilton’s equations:

  • Symplectic integrators: the leapfrog method (the standard choice).

• Hamiltonian dynamics are time-reversible and volume-preserving.

• The dynamics keep the Hamiltonian invariant. A Hamiltonian trajectory will (if simulated exactly) move within a hypersurface of constant probability density.

• Each iteration of the HMC algorithm has two steps. Both steps leave the joint distribution of $p(\boldsymbol{x}, \boldsymbol{r})$ invariant (detailed balance).

  • In the first step, new values of $\boldsymbol{r}$ are randomly drawn from their Gaussian distribution, independently of the current $\boldsymbol{x}$.

  • In the second step, a Metropolis update is performed, using Hamiltonian dynamics to propose a new state.

  • Optimal acceptance rate is 0.65. The step size $\varepsilon$ and trajectory length $L$ need to be tuned.

No-U-Turn sampler¶

• HMC is an algorithm that avoids the random walk behavior and sensitivity to correlated parameters that plague many MCMC methods by taking a series of steps informed by first-order gradient information.

• However, HMC’s performance is highly sensitive to two user-specified parameters: a step size $\varepsilon$ and a desired number of steps $L$.

• The No-U-Turn Sampler (NUTS), an extension to HMC that eliminates the need to set a number of steps $L$, as well as the step-size.

• We simulate in discrete time steps, and to make sure you explore the parameter space properly you simulate steps in one direction and the twice as many in the other direction, turn around again, etc. At some point you want to stop this and a good way of doing that is when you have done a U-turn (i.e., appear to have gone all over the place).

• NUTS begins by introducing an auxiliary variable with conditional distribution. After re-sampling from this distribution, NUTS uses the leapfrog integrator to trace out a path forwards and backwards in fictitious time. First running forwards or backwards 1 step, then forwards or backwards 2 steps, then forwards or backwards 4 steps, etc.

• This doubling process implicitly builds a balanced binary tree whose leaf nodes correspond to position-momentum state. The doubling is stopped when the subtrajectory from the leftmost to the rightmost nodes of any balanced subtree of the overall binary tree starts to double back on itself (i.e., the fictional particle starts to make a U-turn).

• At this point NUTS stops the simulation and samples from among the set of points computed during the simulation, taking care to preserve detailed balance.

• To adapt the step-size, NUTS uses a modified dual averaging algorithm during the burn-in phase.

• The good thing about NUTS is that proposals are made based on the shape of the posterior and they can happend at the other end of the distribution. In contrast, MH makes proposals within a ball, and Gibbs sampling only moves along one (or at least very few) dimensions at a time.

Let us try the NUTS sampler in CUQIpy for the donut target distribution target_donut. We create the NUTS sampler object NUTS_donut for the target distribution target_donut with initial sample [0, 0]:

We set the number of samples and burin-in and sample the target distribution: Note: for NUTS in CUQIpy, using sample or sample_adapt is the same. NUTS will always adapt if the number of burn-in samples is larger than 0 and the parameter adapt_step_size is passed as True, the later is a default behavior.

Let us visualize the samples using the plot_pair method over the target distribution PDF. We will also add connetions between the samples to visualize the path of the NUTS sampler:

We can see that NUTS can take rapid steps in exploring the target distribution. The samples are spread all over the domain and capture the target distribution efficiently.

We see that although the number of NUTS samples is much smaller than in the CWMH samples, the ESS is much larger. This however comes at a cost of computational time, which is due to building the tree structure in each iteration as well as the finite difference approximation of the gradient. Note that other approaches that are more computationally efficient can be used to compute the gradient of the target distribution, e.g. computing the adjoint-based gradient.