PyVBMC Example 3: Output diagnostics and saving results#
In this notebook, we demonstrate extended usage of PyVBMC. We will take a brief look at PyVBMC’s diagnostic output, and show you how to save the results of optimization to disk.
This notebook is Part 3 of a series of notebooks in which we present various example usages for VBMC with the PyVBMC package. The code used in this example is available as a script here.
import sys
if "google.colab" in sys.modules: # Colab lacks PyVBMC: install it
%pip install --upgrade pyvbmc
import numpy as np
import scipy.stats as scs
from scipy.optimize import minimize
from pyvbmc import VBMC
from pyvbmc.formatting import format_dict
1. Model definition and setup#
For demonstration purposes, we will run PyVBMC with a restricted budget of function evaluations, insufficient to achieve convergence. Then we will inspect the output diagnostics, and resume optimization.
We use a higher-dimensional analogue of the same toy target function in Example 1, a broad Rosenbrock’s banana function in \(D = 4\).
D = 4 # A four-dimensional problem
prior_mu = np.zeros(D)
prior_var = 3 * np.ones(D)
def log_prior(theta):
"""Multivariate normal prior on theta."""
cov = np.diag(prior_var)
return scs.multivariate_normal(prior_mu, cov).logpdf(theta)
The likelihood function of your model will in general depend on the observed data. This data can be fixed as a global variable, as we did directly above for prior_mu and prior_var. It can also be defined by a default second argument: to PyVBMC there is no difference so long as the function can be called with only a single argument (the parameters theta):
def log_likelihood(theta, data=np.ones(D)):
"""D-dimensional Rosenbrock's banana function."""
# In this simple demo the data just translates the parameters:
theta = np.atleast_2d(theta)
theta = theta + data
x, y = theta[:, :-1], theta[:, 1:]
return -np.sum((x**2 - y) ** 2 + (x - 1) ** 2 / 100, axis=1)
def log_joint(theta, data=np.ones(D)):
"""log-density of the joint distribution."""
return log_likelihood(theta, data) + log_prior(theta)
LB = np.full(D, -np.inf) # Lower bounds
UB = np.full(D, np.inf) # Upper bounds
PLB = np.full(D, prior_mu - np.sqrt(prior_var)) # Plausible lower bounds
PUB = np.full(D, prior_mu + np.sqrt(prior_var)) # Plausible upper bounds
In a typical inference scenario, we recommend starting from a “good” point (i.e. one near the mode). We can run a quick preliminary optimization, though a more extensive optimization would not harm.
np.random.seed(41)
x0 = np.random.uniform(PLB, PUB) # Random point inside plausible box
x0 = minimize(
lambda t: -log_joint(t),
x0,
bounds=[
(-np.inf, np.inf),
(-np.inf, np.inf),
(-np.inf, np.inf),
(-np.inf, np.inf),
],
).x
np.random.seed(42)
# Limit number of function evaluations
options = {
"max_fun_evals": 10 * D,
}
# We can specify either the log-joint, or the log-likelihood and log-prior.
# In other words, the following lines are equivalent:
vbmc = VBMC(
log_likelihood,
x0,
LB,
UB,
PLB,
PUB,
options=options,
log_prior=log_prior,
)
# vbmc = VBMC(
# log_joint,
# x0, LB, UB, PLB, PUB, options=options,
# )
Reshaping x0 to row vector.
Reshaping lower bounds to (1, 4).
Reshaping upper bounds to (1, 4).
Reshaping plausible lower bounds to (1, 4).
Reshaping plausible upper bounds to (1, 4).
(PyVBMC expects the bounds to be (1, D) row vectors, and the initial point(s) to be of shape (n, D), but it will accept and re-shape vectors of shape (D,) as well.)
2. Running the model and checking convergence diagnostics#
Now we run PyVBMC with a very small budget of 40 function evaluations:
vp, results = vbmc.optimize()
Beginning variational optimization assuming EXACT observations of the log-joint.
Performance settings: saved calibration.
Tip: Run VBMC 3–4 times with different starting points. Compare the posterior plots across runs; large differences can indicate that a run missed part of the posterior.
https://acerbilab.org/pyvbmc/_examples/pyvbmc_example_4_validation.html
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
0 10 -12.91 2.29 272806.70 2 inf start warm-up
1 15 -4.96 1.34 3.36 2 inf
2 20 -5.12 0.83 1.52 2 28.7
3 25 -5.49 0.19 0.31 2 6.94
4 30 -5.41 0.26 0.21 2 4.67
5 35 -5.34 0.06 0.04 2 1.1
6 40 -5.15 0.12 0.11 2 2.87 end warm-up
inf 40 -5.23 0.07 0.12 50 1.1 finalize
Inference terminated: reached maximum number of function evaluations options.max_fun_evals.
Estimated ELBO: -5.225 +/-0.066.
Caution: Returned variational solution may have not converged.
PyVBMC is warning us that convergence is doubtful. We can look at the output for more information and diagnostics.
print(results["success_flag"])
False
False means that PyVBMC has not converged to a stable solution within the given number of function evaluations.
print(
format_dict(
{
"convergence_status": results["convergence_status"],
"r_index": results["r_index"],
}
)
)
{
'convergence_status': 'no',
'r_index': 1.099153773374601,
}
In the results dictionary:
the
convergence_statusfield says ‘no’ (probable lack of convergence);the reliability index
r_indexis above 1 (it should be less than 1). Our diagnostics tell us that this run has not converged, suggesting to increase the budget.
Note that convergence to a solution does not mean that it is a good solution. You should always check the returned variational posteriors, and ideally should compare across multiple runs of PyVBMC.
3. Saving results#
We can also save the VBMC instance to disk and reload it later, in order to check the results and convergence diagnostics, sample from the posterior, or resume the optimization from checkpoint etc. If you are only interested in the final (best) variational solution, as opposed to the full iteration history of the optimization, then you may wish to save only the final VariationalPosterior instead.
VBMC instance contains Python bytecode — of the target function, among others — which belongs to the minor version of Python that saved it (3.12, say). Under another minor version the instance can be loaded and inspected, but the optimization should not be resumed or the instance saved again there, which can end the interpreter. To move results between Python versions, save the variational posterior alone (vbmc.vp.save(...)), whose file holds no bytecode.# Here we specify `overwrite=True`, since we don't care about overwriting our
# test file. By default, `overwrite=False` and PyVBMC will raise an error if
# the file already exists.
vbmc.save("vbmc_test_save.pkl", overwrite=True)
# We could also save just the final variational posterior:
# vbmc.vp.save("vp_test_save.pkl")
4. Loading results and resuming the optimization process#
The VBMC.load(file) class method will load a previously-saved VBMC instance from the specified file. We can load the instance saved above and resume the optimization process, increasing maximum number of function evaluations. The default budget is \(50(d+2)\) evaluations, where \(d\) is the dimension of the parameter space (though this example will not require the full budget). We can change this or other options by passing the new_options keyword to VBMC.load(...). Here we increase max_fun_evals and resume the optimization process:
new_options = {
"max_fun_evals": 50 * (D + 2),
}
vbmc = VBMC.load(
"vbmc_test_save.pkl",
new_options=new_options,
iteration=None, # the default: start from the last stored iteration.
set_random_state=False, # the default: don't modify the random state
# (can be set to True for reproducibility).
)
vp, results = vbmc.optimize()
Beginning variational optimization assuming EXACT observations of the log-joint.
Continuing optimization from previous state.
Performance settings: saved calibration.
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
7 45 -4.82 0.15 0.51 2 10.2
8 50 -4.73 0.08 0.13 2 2.68
9 55 -4.60 0.07 0.05 3 1.56
10 60 -4.53 0.03 0.06 4 1.36
11 65 -4.47 0.02 0.04 5 0.955
12 70 -4.30 0.01 0.07 8 1.73 rotoscale, undo rotoscale
13 75 -4.27 0.01 0.01 9 0.226
14 80 -4.22 0.00 0.01 12 0.37
15 85 -4.21 0.01 0.00 15 0.103
16 90 -4.20 0.00 0.01 18 0.269
17 95 -4.18 0.00 0.00 21 0.0767
18 100 -4.17 0.00 0.01 22 0.184
19 105 -4.17 0.00 0.00 23 0.101
20 110 -4.17 0.00 0.00 23 0.0495
21 115 -4.16 0.00 0.00 23 0.0575 stable
inf 115 -4.13 0.00 0.00 50 0.0575 finalize
Inference terminated: variational solution stable for options.tol_stable_count fcn evaluations.
Estimated ELBO: -4.134 +/-0.002.
print(
format_dict(
{
"convergence_status": results["convergence_status"],
"r_index": results["r_index"],
}
)
)
{
'convergence_status': 'probable',
'r_index': 0.057535275703608446,
}
With the default budget of function evaluations, we can see that the convergence_status is ‘probable’ and the r_index is much less than 1, suggesting convergence has been achieved. We can save the result to file and load it later, e.g. to perform further validation or to sample from the variational posterior.
vbmc.save("vbmc_test_save.pkl", overwrite=True)
vbmc = VBMC.load("vbmc_test_save.pkl")
samples, components = vbmc.vp.sample(5)
# `samples` are samples drawn from the variational posterior.
# `components` are the index of the mixture components each
# sample was drawn from.
print(samples)
print(components)
[[ 0.02386004 -0.01340836 -0.05273027 -0.29004221]
[ 0.1591856 -0.32992674 -0.81694259 -1.29970093]
[-0.44658878 -1.51796274 -0.54322082 -0.82897637]
[-0.61505072 -0.33724184 -0.73800971 -0.18671574]
[-0.90011042 -1.13544916 -2.0923922 -0.02907238]]
[12 16 19 13 14]
5. Conclusions#
In this notebook, we have given a brief overview of PyVBMC’s output diagnostics, and shown how to save and load results and resume optimization from the last stored iteration.
In the next notebook, we will illustrate running PyVBMC multiple times in order to validate the results.
Acknowledgments#
Work on the PyVBMC package is supported by the Research Council of Finland (grants 356498 and 358980 to Luigi Acerbi) and its Flagship programme: Finnish Center for Artificial Intelligence FCAI.