Nested Sampling with ernest

Initialization

As with most Bayesian computation, nested sampling requires a likelihood function and prior for the model parameters. In ernest, this is done through the create_likelihood and create_prior functions. We also need to select an LRPS method to be used during the run.

data("epilepsy")
frame <- model.frame(count ~ zAge + zBase * Trt, epilepsy)
X <- model.matrix(count ~ zAge + zBase * Trt, epilepsy)
Y <- model.response(frame)

poisson_log_lik <- function(predictors, response, link = "log") {
  force(predictors)
  force(response)
  link <- make.link(link)
  
  function(theta) {
    eta <- predictors %*% theta
    mu <- link$linkinv(eta)
    sum(dpois(response, lambda = mu, log = TRUE))
  }
}

epilepsy_log_lik <- create_likelihood(poisson_log_lik(X, Y))
epilepsy_log_lik
#> Scalar Log-likelihood Function
#> function (theta) 
#> {
#>     eta <- predictors %*% theta
#>     mu <- link$linkinv(eta)
#>     sum(dpois(response, lambda = mu, log = TRUE))
#> }

epilepsy_log_lik(c(1.94, 0.15, 0.57, -0.20, 0.05))
#> [1] -859.9659

epilepsy_log_lik(c(1.94, 0.15, 0.57, -0.20, Inf))
#> Warning: Replacing `NaN` with `-Inf`.
#> [1] -Inf

poisson_vec_lik <- function(predictors, response, link = "log") {
  force(predictors)
  force(response)
  link <- make.link(link)

  function(theta_mat) {
    eta_mat <- predictors %*% t(theta_mat)
    mu_mat <- link$linkinv(eta_mat)
    colSums(apply(mu_mat, 2, \(col) dpois(response, lambda = col, log = TRUE)))
  }
}

epilepsy_vec_lik <- create_likelihood(vectorized_fn = poisson_vec_lik(X, Y))
epilepsy_vec_lik
#> Vectorized Log-likelihood Function
#> function (theta_mat) 
#> {
#>     eta_mat <- predictors %*% t(theta_mat)
#>     mu_mat <- link$linkinv(eta_mat)
#>     colSums(apply(mu_mat, 2, function(col) dpois(response, lambda = col, 
#>         log = TRUE)))
#> }

theta_mat <- matrix(
  c(1.94, 0.15, 0.57, -0.20, 0.05, 1.94, 0.0, 0.0, 0.0, 0.0),
  byrow = TRUE,
  nrow = 2
)
epilepsy_vec_lik(theta_mat)
#> [1]  -859.9659 -1668.7106
epilepsy_log_lik(theta_mat)
#> [1]  -859.9659 -1668.7106

coef_names <- c("Intercept", "zAge", "zBase", "Trt1", "zBase:Trt1")

norm_transform <- function(unit) {
  qnorm(unit, sd = 2.5)
}
custom_prior <- create_prior(norm_transform, names = coef_names)
custom_prior
#> custom prior distribution with 5 dimensions (Intercept, zAge, zBase, Trt1, and zBase:Trt1)

model_prior <- create_normal_prior(names = coef_names, sd = 2.5)
model_prior
#> normal prior distribution with 5 dimensions (Intercept, zAge, zBase, Trt1, and zBase:Trt1)

multi_ellipsoid()
#> Uniform sampling within bounding ellipsoids (enlarged by 1.25):
#> # Dimensions: Uninitialized
#> # Calls since last update: 0
#> 
multi_ellipsoid(enlarge = 1.5)
#> Uniform sampling within bounding ellipsoids (enlarged by 1.5):
#> # Dimensions: Uninitialized
#> # Calls since last update: 0
#> 

rwmh_cube()
#> 25-step random walk sampling (acceptance target = 50%):
#> # Dimensions: Uninitialized
#> # Calls since last update: 0
#> 
rwmh_cube(steps = 30, target_acceptance = 0.4)
#> 30-step random walk sampling (acceptance target = 40%):
#> # Dimensions: Uninitialized
#> # Calls since last update: 0
#> 

Generating Samples

To initialize a sampler, call ernest_sampler with the likelihood, prior, and LRPS objects. This creates an ernest_sampler object, which contains (among other metadata) an initial live set of points ordered by their likelihood values. Calling this function also tests your likelihood and prior transformation functions and reports unexpected behaviour.

sampler <- ernest_sampler(
  epilepsy_log_lik,
  model_prior,
  sampler = rwmh_cube(),
  nlive = 300,
  seed = 42
)
sampler
#> Nested sampling run specification:
#> * No. points: 300
#> * Sampling method: 25-step random walk sampling (acceptance target = 50%)
#> * Prior: normal prior distribution with 5 dimensions (Intercept, zAge, zBase,
#> Trt1, and zBase:Trt1)

The generate function runs the nested sampling loop and controls when a run will stop. For example, you can perform 1000 sampling iterations by setting max_iterations:

run_1k <- generate(sampler, max_iterations = 1000, show_progress = FALSE)
run_1k
#> Nested sampling run:
#> * No. points: 300
#> * Sampling method: 25-step random walk sampling (acceptance target = 50%)
#> * Prior: normal prior distribution with 5 dimensions (Intercept, zAge, zBase,
#> Trt1, and zBase:Trt1)
#> ── Results ─────────────────────────────────────────────────────────────────────
#> * Iterations: 1000
#> * Likelihood evals.: 16401
#> * Log-evidence: -1614.9184 (± 27.7466)
#> * Information: 802.4

generate produces an ernest_run object, which inherits from ernest_sampler and contains data describing the run. It is not recommended to overwrite values within this object, but you can explore its internals. For example, ernest calculates the contribution of each shell to the final log-evidence estimate as

\[ w_i = f(L_i) * \Delta V_i \]

where \(\Delta V_i = V_i - V_{i-1}\) and \(f(L_i) = (L_{i-1} + L_i)/2\). Summing these unnormalized log-weights gives the model evidence estimate, which we can view by accessing them within the run_1k object:

run_1k$samples$log_weight |> summary()
#> Length  Class   Mode 
#>      0   NULL   NULL

Robustness

Traversing and integrating over an arbitrary parameter space \(\Theta\) is challenging, especially due to high dimensionality—the main motivation for nested sampling (Skilling 2004). Even in low dimensions, \(\Theta\) may have features such as non-convexity or multimodality that hinder exploration (Buchner 2023). Such pathologies often confound MCMC and related Monte Carlo methods (Freeman and Dale 2013).

Nested sampling avoids many of these problems: by generating samples from the entire prior region using LRPS, it is less sensitive to local behaviour and can globally explore the parameter space, rather than getting stuck on features in the likelihood, such as local maxima.

Complexity

Nested sampling is computationally intensive, mainly due to its need to explore the entire parameter space. The exact complexity of a run is problem-specific, depending on the LRPS method and the shape of the likelihood-restricted prior at each step; however, Skilling (2009) notes that the number of iterations needed to successfully integrate the posterior is proportional to \(HN\), where \(H\) is the KL divergence between the posterior and prior:

\[ H = D_\text{KL}(P(\theta) | \pi(\theta)) = \sum_{ \Theta } P(\theta) \, \log \frac{ P(\theta) }{ \pi(\theta) } \]

This gives a rough estimate of complexity; other work refines this for specific likelihoods and priors Skilling (2009). In practice, uninformative or weakly informative priors greatly increase run time, as more iterations are spent traversing uninformative regions.

Stopping Criteria

We can use naive methods to stop a nested sampling run, such as by setting max_iterations or max_evaluations when calling generate. Fortunately, nested sampling provides a more nuanced method for terminating a run. At each iteration, we can estimate the amount of evidence that remains un-integrated as (Speagle 2020):

\[\begin{equation} \Delta \hat{\mathbf{Z}_{i}} = L_{max} V_i \end{equation}\]

This estimate assumes the remaining space can be represented as a single shell with likelihood \(L_{max}\) and volume \(V_i\). As the run continues, this contribution shrinks. Once it is small relative to the current evidence estimate, the run can be considered to have successfully traversed the informative region of the parameter space. Numerically, this is represented as a log-ratio (Skilling 2006):

\[\begin{equation} \Delta \ln{\epsilon_i} = \ln{(\hat{\mathbf{Z}_i} - \Delta \hat{\mathbf{Z}_i})} - \ln{\hat{\mathbf{Z}_i}} \end{equation}\]

Since \(\Delta \hat{\mathbf{Z}_i}\) tends to overestimate the remaining contribution, this criterion results in more iterations than strictly necessary, allowing the sampler to better detect irregularities in the likelihood surface.

In ernest, we set this criterion with the min_logz argument in generate. By default, generate will produce samples until min_logz falls below \(0.05\). We can call generate on the earlier-produced run_1k object to continue sampling until the run satisfies \(\Delta \ln{\epsilon_i} < 0.05\).

run <- generate(run_1k, show_progress = FALSE)
run
#> Nested sampling run:
#> * No. points: 300
#> * Sampling method: 25-step random walk sampling (acceptance target = 50%)
#> * Prior: normal prior distribution with 5 dimensions (Intercept, zAge, zBase,
#> Trt1, and zBase:Trt1)
#> ── Results ─────────────────────────────────────────────────────────────────────
#> * Iterations: 7872
#> * Likelihood evals.: 188201
#> * Log-evidence: -883.1998 (± 0.3181)
#> * Information: 20.86

Uncertainty

Evidence estimates produced by nested sampling have two main sources of error: (a) error in \(\Delta V_i\), due to uncertainty in assigning volumes to shells over \(P(\theta)\); and (b) error in \(f(L_i)\), from using a point estimate for the likelihood across a shell. A key advantage of nested sampling is that both sources of error can be estimated without repeating the sampling procedure.

ernest provides methods for estimating uncertainty due to \(\Delta V_i\) with both experimental and analytic means. For the former, ernest reports a “cheap’n’cheerful” (Skilling 2009) approximation for the variance in a log-evidence estimate:

\[ \mathbf{V}(\hat{\ln \mathbf{Z}}) = \sqrt{\frac{\hat{H}}{N}} \]

where \(\hat{H}\) is the estimated information, calculated from a nested sampling run:

\[ \hat{H} = -\ln{\mathbf{Z}} + \frac{1}{\mathbf{Z}} \int_0^1 \mathbf{L}(\theta) \ln{\mathbf{L}(\theta)}dx \]

You can view this uncertainty by calling summary on an ernest_run object. We can additionally visualize how the log weights at each iteration changed across the run by using the plot method.

summary(run)
#> Summary of nested sampling run:
#> ── Run Information ─────────────────────────────────────────────────────────────
#> * No. points: 300
#> * Iterations: 7872
#> * Likelihood evals.: 188201
#> * Log-evidence: -883.1998 (± 0.3181)
#> * Information: 20.86
#> * RNG seed: 42
#> ── Posterior Summary ───────────────────────────────────────────────────────────
#> # A tibble: 5 × 6
#>   variable      mean    sd  median    q15   q85
#>   <chr>        <dbl> <dbl>   <dbl>  <dbl> <dbl>
#> 1 Intercept   1.41   1.22   1.89    0.623 2.01 
#> 2 zAge        0.112  0.841  0.151  -0.155 0.416
#> 3 zBase       0.445  0.832  0.567   0.195 0.768
#> 4 Trt1       -0.184  1.05  -0.185  -0.641 0.363
#> 5 zBase:Trt1  0.0254 0.856  0.0480 -0.263 0.369
#> ── Maximum Likelihood Estimate (MLE) ───────────────────────────────────────────
#> * Log-likelihood: -859.9751
#> * Original parameters: 1.9353, 0.1466, 0.5699, -0.196, and 0.05
plot(run, which = c("weight", "likelihood"))

Additionally, since the shrinkage in \(V_i\) between iterations follows \(\text{Beta}(N, 1)\), it is relatively inexpensive to simulate different shrinkage sequences to create a bootstrapped statistic. In ernest, we can do this using calculate, which has a similar plot method.

sim_run <- calculate(run, ndraws = 1000)
sim_run
#> # A tibble: 8,172 × 4
#>     log_lik        log_volume       log_weight     log_evidence
#>       <dbl>        <rvar[1d]>       <rvar[1d]>       <rvar[1d]>
#>  1 -1.36e21  -0.0034 ± 0.0034  -1.4e+21 ± 0.00  -1.4e+21 ± 0.00
#>  2 -3.67e17  -0.0069 ± 0.0048  -3.7e+17 ± 0.00  -3.7e+17 ± 0.00
#>  3 -2.49e16  -0.0101 ± 0.0060  -2.5e+16 ± 2.73  -2.5e+16 ± 2.73
#>  4 -2.62e15  -0.0134 ± 0.0066  -2.6e+15 ± 0.84  -2.6e+15 ± 0.84
#>  5 -7.60e13  -0.0167 ± 0.0077  -7.6e+13 ± 0.81  -7.6e+13 ± 0.81
#>  6 -5.04e13  -0.0200 ± 0.0085  -5.0e+13 ± 0.81  -5.0e+13 ± 0.81
#>  7 -4.45e13  -0.0232 ± 0.0092  -4.5e+13 ± 0.79  -4.5e+13 ± 0.79
#>  8 -1.73e13  -0.0268 ± 0.0098  -1.7e+13 ± 0.77  -1.7e+13 ± 0.77
#>  9 -1.49e13  -0.0302 ± 0.0104  -1.5e+13 ± 0.83  -1.5e+13 ± 0.83
#> 10 -2.48e12  -0.0337 ± 0.0110  -2.5e+12 ± 0.82  -2.5e+12 ± 0.82
#> # ℹ 8,162 more rows
plot(sim_run, which = c("weight", "likelihood"))

Uncertainty within \(L_i\) is harder to quantify, as it may arise from systematic errors in LRPS. In the future, ernest will allow for creating simulated runs through bootstrapping: a run is split into \(N\) threads (each with one live point), then regrouped into synthetic runs of size \(N\) with replacement (Higson et al. 2018). Currently, users interested in bootstrapping ernest_run objects are referred to the nestcheck python package (Higson et al. 2019).

Posterior Distributions

Nested sampling, while targeting model evidence, can also produce an estimated posterior distribution as a side-effect. As each point and its log-weight are discarded from the live set, we can estimate the posterior using this dead set of points, weighted by their importance weights:

\[\begin{align} p_i &\approx \frac{w_i}{\sum_{\forall i}{w_i}} \\ &\approx \frac{w_i}{\hat{Z}} \end{align}\]

In ernest, the posterior can be extracted using the as_draws object. After reweighting this sample by its importance weights, we can summarize the posterior using the posterior package, or produce visualizations.

as_draws(run) |>
  resample_draws() |>
  summarise_draws()
#> # A tibble: 5 × 10
#>   variable      mean  median     sd    mad        q5     q95  rhat ess_bulk
#>   <chr>        <dbl>   <dbl>  <dbl>  <dbl>     <dbl>   <dbl> <dbl>    <dbl>
#> 1 Intercept   1.93    1.93   0.0368 0.0362  1.87      1.99    1.09    1629.
#> 2 zAge        0.151   0.151  0.0249 0.0251  0.110     0.192   1.08    1776.
#> 3 zBase       0.572   0.572  0.0242 0.0244  0.532     0.613   1.12    1670.
#> 4 Trt1       -0.193  -0.195  0.0546 0.0573 -0.282    -0.105   1.11    1644.
#> 5 zBase:Trt1  0.0483  0.0486 0.0290 0.0278 -0.000692  0.0971  1.10    1727.
#> # ℹ 1 more variable: ess_tail <dbl>
visualize(run, .which = "trace")

visualize(run, -Intercept, .which = "density")