---
title: "More Examples with ernest"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{More Examples with ernest}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---
```{r, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
```

```{r, include = FALSE, eval = FALSE}
#' @srrstats {G1.5} These examples are based on classical test problems for
#' nested sampling used within the nestle package. These problems are
#' also used within ernest's test suite to validate its behaviour.
#' @srrstats {G5.1} Tests and test data are provided here for users to
#' interact with.
#' @srrstats {BS1.1} Demonstrates how to enter data through a likelihood
#' function.
```

```{r setup, message = FALSE}
library(ernest)
library(posterior)
library(ggplot2)
```

This vignette describes the test problems used by ernest to validate its sampling behaviour. By reading this, you will learn how different models can be expressed within ernest and how to examine the `ernest_run` produced by a sampling run.

## Multimodal Bivariate Gaussian "Blobs"

*From the Python package [nestle](https://github.com/kbarbary/nestle/blob/master/).*

The goal is to estimate the evidence of a multimodal distribution, made up of a mixture of two well-separated bivariate Gaussian distributions. As nested sampling is designed to handle multimodal distributions, this provides a good test of ernest's ability to explore a complicated likelihood surface.

```{r}
# Log-likelihood for two Gaussian blobs
sigma <- 0.1
mu1 <- c(1, 1)
mu2 <- -c(1, 1)
sigma_inv <- diag(2) / 0.1**2

gaussian_blobs_loglik <- function(x) {
  dx1 <- -0.5 * mahalanobis(x, c(1, 1), sigma_inv, inverted = TRUE)
  dx2 <- -0.5 * mahalanobis(x, c(-1, -1), sigma_inv, inverted = TRUE)
  matrixStats::colLogSumExps(rbind(dx1, dx2))
}

# Uniform prior over [-5, 5] in each dimension
prior <- create_uniform_prior(lower = -5, upper = 5, names = c("A", "B"))

# Run the model
blob_sampler <- ernest_sampler(
  create_likelihood(vectorized_fn = gaussian_blobs_loglik),
  prior,
  nlive = 100,
  seed = 42L
)
blob_result <- generate(blob_sampler, show_progress = FALSE)
```

The expected log-evidence, found analytically, is $-6.679$. We may use the summary and calculate methods to examine how accurate our results from sampling are:

```{r}
summary(blob_result)
calculate(blob_result, ndraws = 500)$log_evidence |>
  tail(1)
```

```{r include = FALSE}
rm(blob_sampler, blob_result)
```

## Eggbox Distribution

*From the Python package [nestle](https://github.com/kbarbary/nestle/blob/master/).*

The eggbox likelihood is another useful test case that demonstrates ernest's ability to properly integrate across multimodal likelihood surfaces.

```{r}
# Log-likelihood for the eggbox function
eggbox_loglik <- function(x) {
  tmax <- 5.0 * pi
  if (!is.matrix(x)) dim(x) <- c(1, length(x))
  t <- sweep(2.0 * tmax * x, 2, tmax, "-")
  (2.0 + cos(t[, 1] / 2.0) * cos(t[, 2] / 2.0))^5.0
}

# Uniform prior over [0, 1] in each dimension
eggbox_prior <- create_uniform_prior(names = c("A", "B"))
```

```{r echo = FALSE}
eggbox_sample <- expand.grid(
  "A" = seq(0, 1, by = 0.01),
  "B" = seq(0, 1, by = 0.01)
)
eggbox_sample$logl <- mapply(
  function(a, b) eggbox_loglik(c(a, b)),
  eggbox_sample$A,
  eggbox_sample$B
)

library(ggplot2)
ggplot(eggbox_sample, aes(A, B, fill = logl)) +
  geom_tile() +
  scale_fill_viridis_c("Log-Lik.")
rm(eggbox_sample)
```

Analytic results indicate that we should expect a log-evidence of $235.895$.

```{r}
egg_sampler <- ernest_sampler(
  eggbox_loglik,
  eggbox_prior,
  sampler = multi_ellipsoid(),
  seed = 42L
)
egg_result <- generate(egg_sampler, show_progress = FALSE)
summary(egg_result)
```

We can further visualize the multimodal structure of each variable:
```{r}
visualize(egg_result, .which = "trace")
```

```{r include = FALSE}
rm(egg_sampler, egg_result)
```

## Incorporating Data within Likelihood Functions

Nested sampling runs are capable of generating estimates of the posterior distribution. To examine this, we provide a test problem from the National Institute of Standards and Technology ([NIST](https://www.itl.nist.gov/div898/strd/mcmc/mcmc01.html)).
```{r}
y <- c(0.2, 0.1, 0.3, 0.1, 0.3, 0.1, 0.3, 0.1, 0.3, 0.1) + 1e+08
```

The certified posterior quantiles for the mean and standard deviation of `y` are:

| Parameter | 2.5%                  | 50.0%                 | 97.5%                |
|-----------|-----------------------|-----------------------|----------------------|
| mu        | 100000000.13281908588 | 100000000.20000000000 | 100000000.26718091412|
| sigma     | 0.069871704416342     | 0.103462818336964     | 0.175493354741336    |

```{r include = FALSE}
ref <- vctrs::data_frame(
  variable = c("mu", "sigma"),
  `2.5%` = c(100000000.13281908588, 0.069871704416342),
  `50%` = c(100000000.20000000000, 0.103462818336964),
  `97.5%` = c(100000000.26718091412, 0.175493354741336)
)
```

To run this model, we must incorporate the data `y` into the likelihood function. We do this through a function factory that captures a vector of observations from its environment. The produced function then computes the log-likelihood of a Gaussian model with parameters mean `mu` and standard deviation `sigma`.

```{r}
gaussian_log_lik <- function(data) {
  force(data)

  function(theta) {
    if (theta[2] <= 0) return(-Inf)
    sum(stats::dnorm(data, mean = theta[1], sd = theta[2], log = TRUE))
  }
}

log_lik <- gaussian_log_lik(y)
prior <- create_uniform_prior(
  names = c("mu", "sigma"),
  lower = c(1e+08 - 1, 0.01),
  upper = c(1e+08 + 1, 1)
)
nist_sampler <- ernest_sampler(log_lik, prior, seed = 42L)
nist_result <- generate(nist_sampler, show_progress = FALSE)
```

We rely on the posterior R package to examine the posterior distribution by summarizing each variable by its quantiles. To better visualize the correspondence between the expected and actual values, we can plot their respective IQRs.

```{r}
post <- as_draws(nist_result) |>
  resample_draws() |>
  summarise_draws(\(x) quantile(x, probs = c(0.025, 0.5, 0.975)))
post
```

```{r echo = FALSE}
all <- rbind(post, ref)
all$src <- rep(c("est", "act"), each = 2)
all

ggplot(all, aes(x = src)) +
  geom_crossbar(aes(ymin = `2.5%`, y = `50%`, ymax = `97.5%`)) +
  facet_grid(rows = vars(variable), scales = "free_y") +
  scale_y_continuous("Value") +
  scale_x_discrete(
    "Data Source",
    breaks = c("act", "est"),
    labels = c("Actual", "Expected")
  )
```