Skip to content

Commit 4e705eb

Browse files
mikkelbuegmingas
andauthored
Advanced features for the MLDA step class (#4125)
* Add a subsample() function in sampling.py to improve the speed of the MLDA sampler. It is a stripped down version of sample(). Also modify RecursiveDAProposal to use subsample(). * Add a second gravity surveying example, supporting code, and two fenics notebooks * Add support for DEMetropolisZ as base sampler. Changes: - Adds new class DEMetropolisZMLDA that inherits from DEMetropolisZ - Modifies MLDA class to accepte DEMetropolisZ as base sampler (and makes it default) - Adds extra stats - Adds more tests for DEMetropolisZ base sampler in test_step.py and test_examples.py * Adjust notebooks to new MLDA class notation, re-run notebooks, add/remove notebooks * Add variance reduction feature to MLDA Changes: - Adds various parts of mlda.py in order to allow keeping quantities of interest Q in memory from all levels and outputing them to stats - Adds store_Q_fine argument and new stats - Modify MetropolisMLDA and DEMetropolisZMLDA to allow them to be used for VR - Add notebooks with VR demonstration - Add test for variance reduction feature * Add adaptive error model to MLDA Changes: - Add code to MLDA class to support AEM - Add supporting class for error estimation to mlda.py - Add test for AEM - Add two example notebooks which use AEM * Repair some minor discrepancies after rebase * Adjust notebooks to new MLDA class notation, re-run notebooks, add/remove notebooks * Extend the docstring of MLDA in the VR and AEM sections * Fix pre-commit errors * Removed unnecessary test from subsampling function * Removed notebooks with external dependencies and their utility code * Added new notebook on variance reduction to examples list * Moved subsample function to mlda.py * Add random selection of subsample for proposal * Remove utility vr function to fix test and add some comments * Add extract_Q_estimate utility function and add description of it in VR notebook * Rename Q utility function * Remove redundant import from test file * Make MLDA create subchains of random length and select the last sample when VR is not used. Rerun notebooks * Revert to previous version of sampling for non-VR settings. Now, for non-VR we have fixed subchain length and we select the last subchain sample as a proposal. For VR we have fixed subchain length and we select a random subchain sample as a proposal. Also, rerun notebooks. * Tidy up code (remove outdated comments, fix VR notebook) * Add MLDA introduction notebook with links to all other notebooks * Tidyed up the tuning_end_trigger functionality of DEMetropolisZ for MLDA * Refactored the astep function for the MLDA base steppers. * Fix issues in intro notebook * Changed the MLDA naming convention from next_model to model_below and next_step_method to step_method_below * Reorganised the variance reduction trigger in the astep method to be more uniform with the error model * Tidyed up the code and removed duplicate assignments of self.model_below and self.model * Added typing to new variables in MLDA step class * Blacked mlda.py and test_step.py * Updated the MLDA_introduction notebook to reflect the current status of notebooks * Rewrote the base tuning stats routines to produce results more consistent with upstream PyMC3 * Reran core MLDA notebooks * Blacked MLDA source and notebooks * Rebased on upstream PyMC3 * Fixed pre-commit errors and unused dependecies * Added extract_Q_estimate() to tests to improve codecov Co-authored-by: Greg Mingas <[email protected]>
1 parent 1d89609 commit 4e705eb

9 files changed

+3359
-1029
lines changed

docs/source/notebooks/MLDA_gravity_surveying.ipynb

Lines changed: 534 additions & 803 deletions
Large diffs are not rendered by default.
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# MLDA sampler: Introduction and resources\n",
8+
"\n",
9+
"This notebook contains an introduction to the Multi-Level Delayed Acceptance MCMC algorithm (MLDA) proposed in [1]. It explains the main idea behind the method, gives an overview of the problems it is good for and points to specific notebooks with examples of how to use it within PyMC3. \n",
10+
"\n",
11+
"[1] Dodwell, Tim & Ketelsen, Chris & Scheichl, Robert & Teckentrup, Aretha. (2019). Multilevel Markov Chain Monte Carlo. SIAM Review. 61. 509-545. https://doi.org/10.1137/19M126966X"
12+
]
13+
},
14+
{
15+
"cell_type": "markdown",
16+
"metadata": {},
17+
"source": [
18+
"### The MLDA sampler\n",
19+
"\n",
20+
"The MLDA sampler is designed to deal with computationally intensive problems where we have access not only to the desired (fine) posterior distribution but also to a set of approximate (coarse) posteriors of decreasing accuracy and decreasing computational cost (we need at least one of those). \n",
21+
"\n",
22+
"Its main idea is to run multiple chains, where each chain samples from a different version of the posterior (the lower the level, the coarser the posterior). Each chain generates a number of samples and the last of them is passed as a proposal to the chain one level up. The latter accepts or rejects the sample and then gives back control to the first chain. This strategy is applied recursively so that each chain uses the chain below as a source of proposed samples. \n",
23+
"\n",
24+
"MLDA improves the effective sample size of the finest chain compared to standard samplers (e.g. Metropolis) and this allows us to reduce the number of expensive fine-chain likelihood evaluations while still exploring the posterior adequately. Note that the bottom level sampler is a standard MCMC sampler like Metropolis or DEMetropolisZ."
25+
]
26+
},
27+
{
28+
"cell_type": "markdown",
29+
"metadata": {},
30+
"source": [
31+
"### Problems it is good for\n",
32+
"\n",
33+
"In many real-world problems, we use models to represent spatially or temporally varying quantities. We often have the ability to modify the spatial and/or temporal granularity of the models. For example, when representing a high-resolution image, we can use a coarse 64x64 grid but we can also use a much finer 512x512 grid, which is more accurate and more computationally demanding when working with it. In those cases it is often possible to apply multilevel modeling to infer unknown quantities in the model. In multilevel modeling, a hierarchy of models of increasing accuracy/resolution and increasing computational cost are used together to perform inference more efficiently than doing inference using the finest model only. \n",
34+
"\n",
35+
"Example applications include inverse problems for physical, natural or other systems, e.g. subsurface fluid transportation, predator-prey models in ecology, impedance imaging, ultrasound imaging, emission tomography, flow field of ocean circulation. \n",
36+
"\n",
37+
"In many of those inverse problems, evaluating the Bayesian likelihood requires solving a partial differential equation (PDE) numerically. Doing this on a fine-resolution model can be orders of magnitude slower than doing it on a coarse-resolution model. This is the ideal scenario for multilevel modeling and MLDA as MLDA allows us to get away with only calculating a fraction of the expensive fine-resolution likelihoods in exchange for many cheap coarse-resolution likelihoods. "
38+
]
39+
},
40+
{
41+
"cell_type": "markdown",
42+
"metadata": {},
43+
"source": [
44+
"### PyMC3 implementation\n",
45+
"\n",
46+
"MLDA is one of the MCMC inference methods available in PyMC3. You can instantiate an MLDA sampler using the `pm.MLDA(coarse_models=...)`, where you need to pass at least one coarse model within a list.\n",
47+
"\n",
48+
"The PyMC3 implementation of MLDA supports any number of levels, tuning parameterization for the bottom-level sampler, separate subsampling rates for each level, choice between blocked and compound sampling for the bottom-level sampler, two types of bottom-level samplers (Metropolis, DEMetropolisZ), adaptive error correction and variance reduction.\n",
49+
"\n",
50+
"For more details about the MLDA sampler and the way it should be used and parameterised, the user can refer to the docstrings in the code and to the other example notebooks (links below) which deal with more complex problem settings and more advanced MLDA features.\n",
51+
"\n",
52+
"Please note that the MLDA sampler is new in PyMC3. The user should be extra critical about the results and report any problems as issues in the PyMC3's github repository."
53+
]
54+
},
55+
{
56+
"cell_type": "markdown",
57+
"metadata": {},
58+
"source": [
59+
"### Notebooks with example code\n",
60+
"\n",
61+
"\n",
62+
"[Simple linear regression](./MLDA_simple_linear_regression.ipynb): This notebook demonstrates the workflow for using MLDA within PyMC3. It employes a very simple toy model.\n",
63+
"\n",
64+
"[Gravity surveying](./MLDA_gravity_surveying.ipynb): In this notebook, we use MLDA to solve a 2-dimensional gravity surveying inverse problem. Evaluating the likelihood requires solving a PDE, which we do using [scipy](https://www.scipy.org/). We also compare the performance of MLDA with other PyMC3 samplers (Metropolis, DEMetropolisZ).\n",
65+
"\n",
66+
"[Variance reduction 1](./MLDA_variance_reduction_linear_regression.ipynb) and [Variance reduction 2](https://github.com/alan-turing-institute/pymc3/blob/mlda_all_notebooks/docs/source/notebooks/MLDA_variance_reduction_groundwater.ipynb) (external link): Those two notebooks demonstrate the variance reduction feature in a linear regression model and a groundwater flow model. This feature allows the user to define a quantity of interest that they need to estimate using the MCMC samples. It then collects those quantities of interest, as well as differences of these quantities between levels, during MLDA sampling. The collected quentities can then be used to produce an estimate which has lower variance than a standard estimate that uses samples from the fine chain only. The first notebook does not have external dependencies, while the second one requires FEniCS. Note that the second notebook is outside the core PyMC3 repository because FEniCS is not a PyMC3 dependency.\n",
67+
"\n",
68+
"[Adaptive error model](https://github.com/alan-turing-institute/pymc3/blob/mlda_all_notebooks/docs/source/notebooks/MLDA_adaptive_error_model.ipynb) (external link): In this notebook we use MLDA to tackle another inverse problem; groundwarer flow modeling. The aim is to infer the posterior distribution of model parameters (hydraulic conductivity) given data (measurements of hydraulic head). In this example we make use of Theano Ops in order to define a \"black box\" likelihood, i.e. a likelihood that uses external code. Specifically, our likelihood uses the [FEniCS](https://fenicsproject.org/) library to solve a PDE. This is a common scenario, as PDEs of this type are slow to solve with scipy or other standard libraries. Note that this notebook is outside the core PyMC3 repository because FEniCS is not a PyMC3 dependency. We employ the adaptive error model (AEM) feature and compare the performance of basic MLDA with AEM-enhanced MLDA. The idea of Adaptive Error Model (AEM) is to estimate the mean and variance of the forward-model error between adjacent levels, i.e. estimate the bias of the coarse forward model compared to the fine forward model, and use those estimates to correct the coarse model. Using the technique should improve ESS/sec on the fine level.\n",
69+
"\n",
70+
"[Benchmarks and tuning](https://github.com/alan-turing-institute/pymc3/blob/mlda_all_notebooks/docs/source/notebooks/MLDA_benchmarks_tuning.ipynb) (external link): In this notebook we benchmark MLDA against other samplers using different parameterizations of the groundwater flow model. We also give some advice on tuning MLDA. Note that this notebook is outside the core PyMC3 repository because FEniCS is not a PyMC3 dependency."
71+
]
72+
},
73+
{
74+
"cell_type": "code",
75+
"execution_count": 1,
76+
"metadata": {},
77+
"outputs": [
78+
{
79+
"name": "stdout",
80+
"output_type": "stream",
81+
"text": [
82+
"last updated: Sat Oct 10 2020 \n",
83+
"\n",
84+
"CPython 3.6.9\n",
85+
"IPython 7.16.1\n",
86+
"watermark 2.0.2\n"
87+
]
88+
}
89+
],
90+
"source": [
91+
"%load_ext watermark\n",
92+
"%watermark -n -u -v -iv -w"
93+
]
94+
}
95+
],
96+
"metadata": {
97+
"kernelspec": {
98+
"display_name": "Python 3",
99+
"language": "python",
100+
"name": "python3"
101+
},
102+
"language_info": {
103+
"codemirror_mode": {
104+
"name": "ipython",
105+
"version": 3
106+
},
107+
"file_extension": ".py",
108+
"mimetype": "text/x-python",
109+
"name": "python",
110+
"nbconvert_exporter": "python",
111+
"pygments_lexer": "ipython3",
112+
"version": "3.6.9"
113+
}
114+
},
115+
"nbformat": 4,
116+
"nbformat_minor": 4
117+
}

docs/source/notebooks/MLDA_simple_linear_regression.ipynb

Lines changed: 261 additions & 63 deletions
Large diffs are not rendered by default.

docs/source/notebooks/MLDA_variance_reduction_linear_regression.ipynb

Lines changed: 1145 additions & 0 deletions
Large diffs are not rendered by default.

docs/source/notebooks/table_of_contents_examples.js

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,5 +60,6 @@ Gallery.contents = {
6060
"ODE_API_introduction": "Inference in ODE models",
6161
"probabilistic_matrix_factorization": "Case Studies",
6262
"MLDA_simple_linear_regression": "MCMC",
63-
"MLDA_gravity_surveying": "MCMC"
63+
"MLDA_gravity_surveying": "MCMC",
64+
"MLDA_variance_reduction_linear_regression": "MCMC"
6465
}

pymc3/step_methods/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@
2121
from .metropolis import BinaryMetropolis
2222
from .metropolis import BinaryGibbsMetropolis
2323
from .metropolis import CategoricalGibbsMetropolis
24-
from .mlda import MLDA, MetropolisMLDA, RecursiveDAProposal
24+
from .mlda import MLDA, MetropolisMLDA, DEMetropolisZMLDA, RecursiveDAProposal
2525
from .metropolis import NormalProposal
26+
from .metropolis import UniformProposal
2627
from .metropolis import CauchyProposal
2728
from .metropolis import LaplaceProposal
2829
from .metropolis import PoissonProposal

0 commit comments

Comments
 (0)