Sampling with CUQIpy: five little stories

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.

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:

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.

Exercise:¶

• In CUQIpy we can remove the burin in in two ways. The first approach is to set Nb to the number of samples you want to remove. The second approach is after you already have the samples, you can use the Samples object burnthin method which removes the burn-in and applies thinning if the user asked for it. Try the second approach to remove a burn-in of size 1500 samples and plot the samples again (same as above, visualize both the density using plot_2D_density and the samples using the plot_pair method).

• Try sampling with another scale (much smaller 0.005 and much larger 1) and visualize the samples, how does that affect the samples and the acceptance rate, and why.

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:

Exercise:¶

• What do you notice about the acceptance rate and the scale in the case with adaptation compared to the previous case.

• Time the two sampling methods, which one is faster and by what factor?

• Compute the effective sample size for both cases. Note that you can use compute_ess method of the Samples object to compute the effective sample size. Which one is larger.

• Scale the ESS by the time so that you have the ESS per second. Which one is larger? Does it pay off to use adaptation?

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.

Story 2 - Not all unkowns 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

Exercise:¶

• We mentioned in the previous notebook that the credible interval (CI) visualized here is computed on the KL coefficients and then mapped to the function values (through computing the linear combination of the KL basis weighted by the KL coefficients and applying the map $x \rightarrow e^x$). We are interested in mapping all the samples to the function values and then computing the credible interval. To achieve this in CUQIpy, one can use the property funvals of the Samples object. This property returns the function values of the samples wrapped in another Samples object. Use this property to create the Samples objects of function values MH_samples_funvals and CWMH_samples_funvals and then compute the credible interval on these new Samples object. Do this for both cases, the MH and the CWMH cases. Which sampler resulted in samples with a better credible interval?

• Which approach of visualizing the credible interval do you prefer (mapping then computing the CI or the other way around)? Why?

• Visualize the credible interval for both cases in the coefficient space. To achieve this in CUQIpy, pass the argument plot_par=True to the plot_ci method. What do you notice? what do you observe about the credible intervals in the coefficient space as the coefficient index increases? Hint: use the original MH_samples and CWMH_samples objects and not the ones converted to function values. Which sampler performs better in the coefficient space?

• Plot the effictive sample size for both cases in the same plot (where the $x-\text{axis}$ is the sample index and the $y-\text{axis}$ is the effective sample size). What do you notice? Hint: use the compute_ess method of the Samples object to compute the effective sample size.

• Plot the sampler scale for both cases in the same plot (the $x-\text{axis}$ is the sample index and the $y-\text{axis}$ is the scale). What do you notice? Hint: you can use CWMH_sampler.scale and MH_sampler.scale to get the scale for each sampler. Keep in mind that the scale for the MH sampler is the same for all components.

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})$.

Exercise:¶

• Try a different scales 0.01 and visualize the samples in the same way as above. What do you notice about the samples? does it capture the target distribution?

• Using the new scale 0.01, sample only 10 samples. Using this scale as well, generate 10 MH samples using a MH sampler MH_10 = MH(target_donut, scale=0.01, x0=np.array([0, 0])). Visualize the samples of both samplers using the plot_pair method (along with the pdf plot plot_2D_density(target_donut, -4, 4, -4, 4)). What do you notice about the samples of the two samplers in terms of finding the high density region of the target distribution?

• Try ULA for a univariate target distribution x_uni = Gaussian(0, 1):

  • Create an ULA sampler object ULA_uni for the target distribution x_uni with scale 1 and initial sample 0.

  • Generate 40000 samples and store it in ULA_samples_uni.

  • Compare the kernel density estimation (KDE) of the samples with the target distribution x_uni PDF. What do you notice? does the KDE capture the target distribution? Hint: to plot the PDF use the plot_1D_density method and to plot the KDE generate a KDE using scipy.stats gaussian_kde (already imported here) kde_est = gaussian_kde(ULA_samples_uni.samples[0,:]), then create a grid grid = np.linspace(-4, 4, 1000) and plot the KDE using plt.plot(grid, kde_est(grid)).

  • Try a different scale 0.1 and repeat the above steps. What do you notice about the samples KDE? does it capture the target distribution?

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.

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.

Exercise:¶

• Compute the ESS of the NUTS samples and visualize the chain. What do you observe?

• Generate a plot similar to the one above (showing samples and connections) but with samples generated from MALA using scale = 0.01. What do you notice about the samples of the two samplers in terms of exploring the target distribution?

• Inspect how the step size in NUTS is adapted during the burn-in phase as well as the size of the tree that is built each step. Create a plot that visualizes both quantities over the sample index. What do you notice? Hint: you can use the properties epsilon_list and num_tree_node_list of the NUTS_donut object to get, respectively, the step size and the tree size lists for the samples.

Exercise:¶

• Try the NUTS sampler for the Poisson problem target target_poisson:

  • Enable finite difference approximation of the posterior gradient by calling target_poisson.enable_FD()

  • Create a NUTS sampler object NUTS_poisson for the target distribution target_poisson and limit the depth of the tree to 7 (use help(NUTS) to learn about the smapler paramters).

  • Sample the target distribution using the sample method setting Ns to 100 and Nb to 10 (sampling might take about 12 min).

  • Visualize the credible interval of the samples along with the exact solution using the plot_ci method. How does the results compare to the MH and CWMH samplers we used earlier?

  • Compute the ESS of the samples and comment on how they compare to the MH and CWMH cases.

  • Can you guess roughly how many times the forward operator was called?

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.