ProxNest: Proximal nested sampling for high-dimensional Bayesian model selection

ProxNest is an open source, well tested and documented Python implementation of the proximal nested sampling algorithm (Cai et al. 2022) which is uniquely suited for sampling from very high-dimensional posteriors that are log-concave and potentially not smooth (e.g. Laplace priors). This is achieved by exploiting tools from proximal calculus and Moreau-Yosida regularisation (Moreau 1962) to efficiently sample from the prior subject to the hard likelihood constraint. The resulting Markov chain iterations include a gradient step, approximating (with arbitrary precision) an overdamped Langevin SDE that can scale to very high-dimensional applications.

Basic Usage

The following is a straightforward example application to image denoising (Phi = I), regularised with Daubechies wavelets (DB6).

# Import relevant modules.
import numpy as np
import ProxNest

# Load your data and set parameters.
data = np.load(<path to your data.npy>)
params = params    # Parameters of the prior resampling optimisation problem.
options = options  # Options associated with the sampling strategy.

# Construct your forward model (phi) and wavelet operators (psi).
phi = ProxNest.operators.sensing_operators.Identity()
psi = ProxNest.operators.wavelet_operators.db_wavelets(["db6"], 2, (dim, dim))

# Define proximal operators for both your likelihood and prior.
proxH = lambda x, T : ProxNest.operators.proximal_operators.l1_projection(x, T, delta, Psi=psi)
proxB = lambda x, tau: ProxNest.optimisations.l2_ball_proj.sopt_fast_proj_B2(x, tau, params)

# Write a lambda function to evaluate your likelihood term (here a Gaussian)
LogLikeliL = lambda sol : - np.linalg.norm(y-phi.dir_op(sol), 'fro')**2/(2*sigma**2)

# Perform proximal nested sampling
BayEvi, XTrace = ProxNest.sampling.proximal_nested.ProxNestedSampling(
    np.abs(phi.adj_op(data)), LogLikeliL, proxH, proxB, params, options
    )

At this point you have recovered the tuple BayEvi and dict Xtrace which contain

Live = options["samplesL"] # Number of live samples
Disc = options["samplesD"] # Number of discarded samples

# BayEvi is a tuple containing two values:
BayEvi[0] = 'Estimate of Bayesian evidence (float).'
BayEvi[1] = 'Variance of Bayesian evidence estimate (float).'

# XTrace is a dictionary containing the np.ndarrays:
XTrace['Liveset'] = 'Set of live samples (shape: Live, dim, dim).'
XTrace['LivesetL'] = 'Likelihood of live samples (shape: Live).'

XTrace['Discard'] = 'Set of discarded samples (shape: Disc, dim, dim).'
XTrace['DiscardL'] = 'Likelihood of discarded samples (shape: Disc).'
XTrace['DiscardW'] = 'Weights of discarded samples (shape: Disc).'

XTrace['DiscardPostProb'] = 'Posterior probability of discarded samples (shape: Disc)'
XTrace['DiscardPostMean'] = 'Posterior mean solution (shape: dim, dim)'

from which one can perform e.g. Bayesian model comparison.

Referencing

A BibTeX entry for ProxNest is:

@article{Cai:ProxNest:2021,
   author = {Cai, Xiaohao and McEwen, Jason~D. and Pereyra, Marcelo},
    title = {"High-dimensional Bayesian model selection by proximal nested sampling"},
  journal = {ArXiv},
   eprint = {arXiv:2106.03646},
     year = {2021}
}