Beyond Metropolis: state-of-the-art MCMC

Back in notebook 4 we looked at the Metropolis algorithm, which is the basic building block of MCMC. It is guaranteed to converge on the target distribution, but we know that convergence will be more efficient if the jump (or proposal) distribution is similar to the target. This motivates the development of adaptive algorithms, where the proposal distribution is "tuned" as the Markov chain progresses.

There are lots of ways of implemeting adaptation (as well as various other innovations for improving the sampling efficiency), but a detailed description of modern MCMC algorithms is beyond the scope of these notes. What's more, if you're a hydrologist or an ecologist (rather than a statistician), you're probably more interested in applying MCMC than learning about all the mathematical details. Fortunately, if you understand the principles of the Metropolis algorithm (see notebook 4), you should be in a good position to use some more sophisticated MCMC packages without having to worry too much about the algorithms themselves.

In this notebook, we'll start off by introducing two Python packages that implement a variety of state-of-the-art MCMC algorithms. We'll then use one of these to calibrate the hydrological model developed in notebook 5.

Note for anyone using this notebook "interactively": the MCMC code in this notebook may take a long time to run (the whole notebook takes more than an hour on my desktop machine).

1. Useful Python packages

1.1. PyMC3

PyMC3 is a Python package offering a variety of sophisticated MCMC samplers, including Hamiltonian Monte Carlo (HMC) and the No U-Turn Sampler (NUTS). It also provides a very clean syntax for model specification (defining priors, likelihoods etc.), as well as convenience functions for a wide range of distributions. This makes setting up a model and running MCMC very simple: as an illustration, check out the linear regression example in the PyMC3 tutorial. This is very similar to the example covered in notebook 4, except there we coded everything from scratch and used only a basic Metropolis algorithm. With PyMC3, the number of lines of code is dramatically reduced and we have the added advantage of being able to switch easily between several modern samplers.

However, a key property of the HMC and NUTS algorithms is they use the gradient of the target function to help tune the proposal distribution. This gives excellent performance for high-dimensional problems, but requires the target function to be differentiable. Unfortunately, as we have already seen in notebook 3, for complex environmental models the likelihood function is often not well-behaved, so the posterior distribution is not differentiable and many of the advantages of the HMC and NUTS algorithms are lost.

1.2. emcee

emcee is Python package which implements an affine invariant MCMC ensemble sampler (a more detailed description of the algorithm from the package authors is here). This sampler uses multiple chains and does not require gradient information, which makes it ideal for the kinds of complex likelihoods commonly encountered in environmental modelling. Furthermore, although in high-dimensional parameter spaces the HMC or NUTS algorithms may be more efficient, the multiple chains used by emcee are easy to parallelise, which offers a significant performance advantage.

The downside of emcee is that, unlike PyMC3, there are no easy ways to specify your model: you have to write Python functions from scratch to represent your priors, likelihood and posterior. For simple models, this is not a major limitation, but for more complex likelihoods (e.g. more sophisticated error structures) the process can become time consuming.

It is possible that the emcee algorithm will eventually be integrated into PyMC3, but in the meantime I recommend investigating both packages as they're useful in different situations.

In this notebook we're going to use emcee.

2. Choosing a likelihood function

Before we run any MCMC, we need to decide on a likelihood function. In the previous examples in notebooks 2 and 4 we assumed a simple likelihood consisting of our deterministic model, $f(x, \theta)$ plus independent and identically distributed Gaussian errors with mean zero, $\epsilon \sim \mathcal{N}(0, \sigma_\epsilon)$. However, as we will see, there are problems with using the same likelihood function for our hydrological model. First, we'll set up our model using the code from notebook 5.

We also need to read in the Tarland meteorological data that we will use to drive the model.

The likelihood function we have used previously adds one additional parameter, $\sigma_\epsilon$, to our list of calibrating parameters. If we select values for $\alpha$, $\tau_s$, $\tau_g$ and $\sigma_\epsilon$, we can therefore run the model and generate some example output.

One problem with this error structure is immediately obvious from the plot above: by adding random noise with a fixed standard deviation, our model frequently predicts negative runoff when flows are low. This clearly doesn't make sense, and we can't solve the problem by simply setting values less than zero back to zero, because if we did our error distribution would no longer be Gaussian with mean zero.

Another problem with this likelihood can be seen from looking at the very last plot in notebook 5, which shows simulated versus observed values. The model parameters used for this plot were chosen at random, so we don't expect them to give a good fit to the observed data. However, it is clear even from this rough example that the errors are likely to be larger when flows are high and smaller when flows are low. In other words, the errors are heteroscedastic and it is therefore not appropriate to use the same value of $\sigma_\epsilon$ for all flows.

One possible way to address both these issues is to make $\sigma_\epsilon$ proportional to the flow, $Q(t)$

$$\sigma_\epsilon = mQ(t)$$

If we do this, when $Q(t)$ is large we expect the error distribution to be wider, but as $Q(t)$ tends to zero, $\sigma_\epsilon$ also tends to zero, so there is a much lower chance of the model predicting negative values. Note that this formulation still only includes a single "error parameter", $m$, but the role of the error parameter in the likelihood function has changed slightly. Once we have estimated $m$, we can caluate $\sigma_\epsilon$ based on the flow time series and proceed as before.

This looks much better: the error variance is large during storm events and decreases as the hydrograph recesses, which matches our expections based the real data and the practical difficulties of estimating flow. Also note that, although it is still possible for the model to predict negative values, this is much less likely to happen than in the first example. Because negaive values happen very rarely, we can safely ignore them without dramatically truncating our error distribution.

When we come to calibrate the hydrological model, we will therefore use a likelihood function that assumes the error variance increases with the flow.

3. Parallel processing

To make the most of emcee, we will use IPython.parallel to write code that can utilise a user-specified number of cores/processors. The first step is to launch multiple Python processes from the Jupyter Clusters tab on the Home screen, as shown in the image below

Once you've found the Clusters tab, set the # of engines equal to the number of cores or processors you have available and then click Start. This will create several new Python processes, which we can use to distribute the MCMC work.

Next, we set up our IPython notebook to make use of these extra processes. The view object in the code below will allow us to allocate our jobs.

Each process has its own namespace, so in order to run our MCMC, each one will need access to all the data necessary to evaluate our log posterior function. There are several ways to ensure this, but one of the simplest is to use IPython's %%px --local magic operator. As described in more detail here, placing this at the start of a notebook code cell will execute the same code on all the available processes (i.e. all the processes created above under the Clusters tab, plus the main Python process hosting the current notebook). We can use this to make sure that each process has access to all the necessary modules, functions and variables.

In the sections below, I've reproduced the hydrological model code from above. We need to run it again, because this time we're also distributing it to our "worker" processes.

Note also that the code here could (and should) be streamlined a lot to reduce the amount of processing required for each model evaluation. In any real-world application, performance tweaks at this stage are usually worth doing, because ultimately the model may be run tens of thousands of times by the MCMC algorithm. It is therefore worth spending a bit of time tidying up the code (and perhaps even optimising bits of it into e.g. Fortran or C), because small gains here can make a big difference later. In the code that appears below I haven't done this, because I want to maintain continuity with the ideas and code presented in the previous notebook.

3.1. Import modules

3.2. Define the ODE system

The function below is used by scipy.integrate.odeint to solve the model equations at each time step.

3.3. Define the hydrological model

This function wraps the model equations in a loop to evaluate the drainage volume at each step in the input time series.

3.4. Define log likelihood

As discussed above, we will use a log likelihood that assumes independent Gaussian errors, but with the standard deviations, $\sigma_\epsilon$, proportional to the flow, $Q(t)$. This is a pretty sensible starting point for hydrological modelling, but there are some other issues we are not addressing here. For example, because we are working with time series data we know the errors are unlikely to be independent - most environmental time series exhibit autocorrelation to some degree. It is possible to modify the likelihood function to include autocorrelation too, but we won't go into this yet, because there are some subtleties to consider.

Note that in the function below we've introduced the parameter $m$, which is used to estimate the value of $\sigma_\epsilon$ for each point in the time series.

3.5. Define log prior

We must define prior distributions for the model parameters we intend to calibrate (plus the standard deviation of the error term in our likelihood function, which is now controlled by $m$). The model parameters of interest are $\alpha$, $\tau_s$ and $\tau_g$ and, for the sake of simplicity, we will assume uniform distributions for all of them.

Note that the constants alpha_max etc. in the code below will be specified later.

3.6. Define log posterior

The log posterior distribution is (proportional to) the sum of the log prior and the log likelihood. The code below is almost identical to that used in Section 7.5 of notebook 4.

3.7. Read the input data

Ultimately, we want to use the real data from the Tarland catchmemnt to calibrate our model. Before we do that, however, it is worth testing our entire workflow using some synthetic data to assess our approach.

As discussed previously, one of the trickiest parts of any MCMC analysis is formulating an appropriate likelihood function. If we get the likelihood function wrong, our parameter estimates and their associated uncertainty will be biased, and it is difficult to assess how biased without knowing the true likelihood function. Because we never have this information in reality (assuming it even exists), it can be difficult to know whether our method is produing sensible results.

One way to gain confidence in a statistical methodology is to generate "fake" data with a known error structure, which we can use to test our method. If our MCMC algorithm can recover the "fake true" parameters with reasonable accuracy under ideal circumstances (i.e. when we know the exact form of the likelihood function), then perhaps it will also have a chance of identifying meaningful parameter sets in the real world, when we don't know the error structure in detail. On the other hand, if the algorithm struggles even under ideal test conditions (e.g. due the model being over-parameterised compared to the amount of test data available), it is very unlikely to yield meaningful results in any real application.

This may sound like a lot of extra work, but it's actually very easy. Once you have your model and have decided on an error structure (i.e. once you've written down your likelihood function, which you need to do anyway), it is very easy to simulate fake observations that can be used for testing. In fact, we've already done this above in section 2, where we assumed some "true" values for $\alpha$, $\tau_s$, $\tau_g$ and $m$ and then ran our model to check the error structure seemed sensible.

In the code that follows, we will use the dataset created above as a synthetic observed time series. If the MCMC algorithm can recover the "fake true" parameter values using the synthetic data, it should be easy to re-run the analysis using the real observed dataset instead; if it can't, we'll need to work out why and refine the approach.

Note that, once again, some of the code from above is repeated here as we need to run it on the "worker" processes we intend to use for parallelising emcee.

3.8. Define prior ranges

We have already decided to use uniform distributions with lower limits of 0 for all our priors, but we still need to specify the upper limit for each distribution.

The parameter $\alpha$ is a correction factor applied to the PET data. We know from notebook 5 that the PET estimates are likely to be too high, so we expect $\alpha$ to be no bigger than about 1. Searching within a prior range from 0 to 2 therefore seems reasonable.

$$\alpha \sim \mathcal{U}(0, 2)$$

$\tau_s$ and $\tau_g$ are both very poorly constrained, so we will set wide priors on those.

$$\tau_s \sim \mathcal{U}(0, 100) \qquad and \qquad \tau_g \sim \mathcal{U}(0, 500)$$

Finally, $m$ is a multiplicative factor applied to $Q(t)$ to estimate $\sigma_\epsilon$. An upper limit of 1 seems like a reasonable starting point, but we can always return later and broaden this range if it appears to be too narrow.

$$m \sim \mathcal{U}(0, 1)$$

3.10. Find the MAP

As noted above, emcee makes use of multiple chains (also called "walkers"), which are the equivalents of the "drunken mountaineers with walkie-talkies" in my dodgy analogy from notebook 4. We therefore need to decide on sensible starting locations for each of our walkers.

One way to do this is to use an optimiser to estimate the Maximum a Posteriori (MAP) density, just as we did in section 7.7 of notebook 4. This is the maximum of the posterior distribution, so starting our walkers in the vicinity of this point is equivalent to helicoptering all the "mountaineers" to top of the highest summit in the region. This is as good a place as any for them to begin their random search, and starting in an area of high posterior density improves the chances of the chains converging efficiently. Of course, we don't want to start all the walkers from exactly the same place, so having found the MAP we'll add a small amount of random noise, so that each starts from a slightly different location that is close to (but not actually on) the posterior "summit". From here, each walker can head off in his/her own direction to explore the surrounding region.

Although the optimiser takes a while to run, it does an excellent job of identifying the "fake true" parameter values we used to generate the synthetic observations. Ultimately, this step should actually save time, because the MCMC algorithm will hopefully converge more quickly due to using the MAP as an approximate starting point.

3.10. Sample from the posterior using emcee

We're now ready to use emcee to draw samples from the posterior. Although the optimiser has already done a good job of identifying the "best" parameter values, we still need to know how much confidence we should have in this result. This is where the MCMC comes in.

Note that the code in the cell below does not need the %%px --local magic operator, because emcee will handle distributing the MCMC chains to our "worker" processes (each of which now has all the information necessary to run the model).

3. Evaluating the results

The plots above are reassuring: the chains have all converged nicely and the posterior distributions for each parameter (the histograms along the diagonal of the bottom plot) nicely identify the "true" parameter values. On these histograms, blue vertical lines mark the "true" values and the black dotted lines mark the 2.5%, 50% and 97.5% quantiles. It is nice to see that all the true values lie within the 95% credible intervals, and in most cases the median MCMC estimate is very close correct value.

The acceptance fraction is the average acceptance rate across all the chains and the value here seems pretty reasonable - see section 4 of notebook 4 for some guidance regarding acceptance rates. The autocorrelation time gives an indication of the number of evaluations of the posterior needed to produce an independent sample of the target density. Again, there are no hard-and-fast rules regarding how this metric should be interpreted, but a useful guide is that sampling for approximately 10 autocorrelation times should ensure you have a representative posterior sample. In this case, the autocorrelation times for each parameter suggest we should have allowed the algorithm to run for a bit longer - perhaps 300 steps in total and discarding the first ~200 as burn-in. For a real application it might be worth re-running the MCMC, but for this illustrative example these chains seem sufficiently well-converged.

Having sampled from the posterior, we are now in a position to estimate the parameter-related and predictive uncertainties in our model output. The parameter-related uncertainty is the uncertainty due to variability in our three calibrating model parameters $\alpha$, $\tau_s$ and $\tau_g$. This represents our confidence in the mean of the model output, ignoring the stochastic error term. The overall predictive uncertainty is the parameter-related uncertainty plus the additional stochastic uncertainty introduced by the error parameter $m$. The overall predictive uncertainty is what matters most in terms of assessing the fit of our model, but a useful feature of the Bayesian approach is that it allows partitioning of uncertainty into parameter-related and random components.

To visualise the results in the domain of our original data (i.e. as a time series) we can draw $N$ parameter sets at random from the MCMC chains (after discarding the burn-in, so we're only selecting from the "converged" portion). For each parameter set we run the model, and the median and range of the model output before adding on the stochastic error component gives an indication of the parameter-related uncertainty. We can then add-on the additional uncertainty from $m$ to give an overall 95% credible interval around the median simulation.

In the code below I've used the excellent mpld3 library to create the plot, so you should be able to click the magnifying glass near the bottom left corner and zoom in to see the uncertainty bounds on the simulated curve.

The darker red shading on the plot above represents the 95% credible interval for the parameter-related uncertainty, which in this case is very narrow because our marginal posteriors are so sharply defined (you'll need to zoom in a long way to even see the parameter-related uncertainty interval on the plot). The lighter red shading defines the 95% credible interval for the overall predictive uncertainty, which in this case is dominated by the stochastic error term.

The next step is to check whether our assumptions regarding the likelihood function have been met. There are many ways to do this, some being more sophisticated than others. One simple but useful diagnostic plot is to look at the standardised residuals, $\epsilon_{std}$ of the model compared to the observations

$$\epsilon_{std} = \frac{\epsilon_{raw}}{\sigma_\epsilon}$$

where $\epsilon_{raw}$ is the difference between modelled and (fake) observed values and $\sigma_\epsilon$ is the standard deviation used at that particular time point (calculated as $mQ(t)$). It is usual to plot the residuals for just a single parameter set taken from somewhere near the MAP. The plot below uses residuals for the median set of parameters from the converged MCMC chains.

It's good to see that there is no obvious structure to the residuals, which also seem to be normally distributed with mean zero (as we originally assumed). We should also check for autocorrelation here, but because this is just an example where we know the residuals are independent (because we generated them that way in section 2, above), we'll postpone looking at autocorrelation until we're working with some real data.

Another easy-to-calculate metric that can be informative is the coverage. This is the proportion of the observations lying within the overall predictive 95% credible interval. If all is well with our approach, we can expect roughly 95% of the observations to be located in this region.

Again, it is reassuring to see that the actual coverage is in close agreement with the theoreticaly expected value of 95%.

4. Real data

So far, we have demonstrated that our approach using emcee is capable of identifying correct parameter sets and uncertainty ranges in our hydrological model as long as the likelihood function is known exactly. It's now time to try the analysis using some real observations, where the assumptions of our likelihood function are less likely to be satisfied.

4.1. Read the input data

We'll begin by reading-in a new copy of the meteorological data to all of our "worker" processes, replacing the modified "fake" dataset we created earlier.

4.2. Adjust the priors

We will also broaden the prior on $\tau_g$, by increasing the upper limit from $500$ to $2000 \; days$.

4.3. Find the MAP

Next, we'll run the optimiser again to find a good satrting point for our "walkers".

4.4. Run the MCMC

And now we can run the MCMC. Note that in the code below truths=None because, unlike our previous test example, we no longer know what the answer should be.

Given that we're using real data, the MCMC chains above have converged very nicely. However, it is worth noting that the posterior histogram for $\tau_g$ is probably unrealistically "spiked", because it is clear from the trace plot for this parameter that the chains didn't properly diverge from the MAP until after about step 150. Because we've only discarded the first 100 steps as burn-in, our sample for $\tau_g$ will not be truly independent of the original starting location. Although this would probably not affect the median estimate for $\tau_g$, for any real application it would be worth re-running the MCMC with a larger number of steps to ensure the posterior samples are genuinely representative.

The median value of $\alpha = 0.78$ seems physically reasonable based on the discussion in notebook 5, and the value of $\tau_s = 12.3 \; days$ is also perfectly credible. The value of $\tau_g = 1830 \; days \approx 5 \; years$ seems rather large, but is not out-of-the-question.

All the posteriors are well-defined and unimodal, which is nice to see.

Next, we'll plot the median simulated time series and the 95% credible intervals, which we can compare with the observed runoff series.

The 95% credible interval for parameter-relatated uncertainty on the plot above is absolutely tiny, emphasising just how narrow the marginal posteriors are. The median simulation also does a reasonable job of matching the recession rate in the real dataset, and the results are clearly better than the example we started with at the end of notebook 5.

On the other hand, it is also true that the model is failing to properly simulate the timing and magnitude of the main flood peaks, and it's also significantly overestimating baseflow during dry periods (e.g. summer 2003). Both these features suggest the estimated groundwater residence time may be too high - see below for further discussion.

5. Evaluate the results

A plot of standardised residuals is shown below. Although clearly not as neat as the example above using synthetic data, the overall appearance is still broadly of Gaussian errors with mean zero, as we assumed by our likelihood function.

The coverage is also about right, which is reassuring.

The final assumption made by our likelihood function is that of serially independent errors i.e. there should be no autocorrelation in the rediduals. In the plot below, the solid and dashed grey lines respectively represent the 95% and 99% confidence bands for autocorrelation. It is clear that significant autocorrelation is present in the residuals, which will have led to bias in our MCMC results.

Dealing with autocorrelation is possible, but not straightforward, and the hydrological literature includes a number of recent papers describing the advantages and disadvantages of different ways of modifying the likelihood.

In these notes I won't go any further with autocorrelation, but it may be worth investigating the literature (e.g. here) if you're interested.

6. A final assessment

The model does a pretty good job of simulating the essential features of the observed data, including getting the recession time about right and correctly identifying most major flow events to within the range of predictive uncertainty. In addition, although there are certainly opportunities for improvements, the overall approach has been successful and the results seem statistically coherrent.

The key question is, of course, is the model actually useful? The answer depends on what we hope to use it for, but it's fair to say that this model wouldn't be much use in most situations of genuine interest, such as investigating high flows (e.g. to assess flooding) or low flows (e.g. to assess drought indices or ecological impacts). It's clear from the time series plot above that, although the model does OK on average, it doesn't accurately represent the magnitude of major flood peaks, nor does it do a very good job of capturing baseflow during long dry spells (especially the significant drought during 2003).

Part of the reason for this is that the model is too simple - there are literally dozens of processes we know are important in catchment hydrology that we haven't included in our conceptualisation. For example, during major storm events much of the incoming water may run over the surface directly to the stream, without entering the soil "reservoir" at all. This water is often referred to as infiltration excess and it can be a significant component of the fast, "flashy" response of some streams during periods of high rainfall. Without at least incorporating some representation of this overland flow into our model, it's probably unreasonable to expect it to accurately represent the size of flow peaks during large storms.

Similarly, the calibrated groundwater residence time, $\tau_g$, of $1800 \; days$ is very large for a catchment such as Tarland. The result of this in our model is to produce a steady, largely unresponsive level of baseflow, which is maintained even during long drought periods (e.g. 2003) - unlike in the observed dataset. One explanation for this issue relates to the way we parameterised our likelihood function at the start of this notebook. Because the error standard deviation, $\sigma_\epsilon$, is assumed to be proportional to flow, if the flow goes to zero then so does $\sigma_\epsilon$ and our code assigns $-\infty$ to the posterior probability. This implies that any parameter set resulting in zero drainage to the stream at any time step is automatically deemed to be impossible, so the algorithm is likely to favour longer groundwater residence times that are capable of maintaining flows through dry periods (like 2003). One common alternative is to introduce an additional parameter to the error structure by estimating $\sigma_\epsilon$ as

$$\sigma_\epsilon = c + mQ(t)$$

A second reason for the algorithm to favour large values of $\tau_g$ is that we have only used the model to simulate 5 years' worth of data, which is insufficent for the groundwater to properly equilibrate if $\tau_g = 1830 \; days$. Once near equilibrium, a groundwater store with such a large residence time would have a high outflow rate, which would likely be incompatible with the observed data. It is possible that either running the model for longer, or setting initial conditions closer to a plausible equilibrium state, would yield smaller estimates for $\tau_g$.

7. Summary

1. Getting the likelihood function (approximately) right is the only way to get your inference (approximately) right.
   
   

2. Before using a complicated model with a complicated MCMC procedure and your real dataset, it's worth generating some test data to see if you can identify the correct parameters under ideal conditions. If not, you'll need to modify your approach.
   
   

3. The Bayesian method makes it possible to apportion your uncertainty to different sources. Here we've only looked at parameter-related and overall predictive uncertainty, but more sophisticated approaches (e.g. BATEA can do more).
   
   

4. It's important to check your assumptions and modify them if necessary. We've only considered the most basic diagnostic plots - there are lots of other options to investigate.
   
   

5. Just because your marginal posteriors look nice and your credible intervals look credible, doesn't mean your model is actually useful!