---
title: "Approximate Bayesian inference in epidist"
description: |
  "A demonstration of how to fit models in epidist using approximate Bayesian
  inference methods such as the Laplace approximation, ADVI, and Pathfinder."
output:
  bookdown::html_vignette2:
    fig_caption: yes
    code_folding: show
    number_sections: true
pkgdown:
  as_is: true
# csl: https://raw.githubusercontent.com/citation-style-language/styles/master/apa-numeric-superscript-brackets.csl
link-citations: true
vignette: >
  %\VignetteIndexEntry{Approximate Bayesian inference in epidist}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
bibliography: references.bib
---



# Background

The `epidist` package uses Bayesian inference to estimate delay distributions and other quantities.
Doing Bayesian inference amounts to approximating the posterior distribution of each parameter in the statistical model^[We are currently developing a vignette explaining the statistical model in more detail!].
A range of methods exist to perform this approximation.

By default, `epidist` uses the No-U-Turn Sampler [NUTS; @hoffman2014no] Hamiltonian Monte Carlo (HMC) algorithm.
NUTS is an example of a broader class of Markov chain Monte Carlo (MCMC) methods.
These methods work by simulating from a Markov chain with the target posterior distribution as its stationary distribution.
When MCMC algorithms are run for sufficiently many iterations, and have reached convergence, the samples can be treated as being drawn from the posterior distribution.
Relevant posterior quantities such as expectations may then be computed using these samples.
A drawback of MCMC methods, like NUTS, is that simulations can be quite computational intensive, especially for complex models or large data.

The `epidist` package is built using `brms` [@brms], which stands for "Bayesian Regression Models using Stan", where Stan [@carpenter2017stan] is a probabilistic programming language.
Although NUTS is the primary inference algorithm used in Stan, additional options are available.
These additional inference algorithms are also available in `epidist` due to its dependence on `brms`.

In this vignette, we first briefly describe the alternative algorithms available (Section \@ref(other)) as well as directing you to more detailed resources.
Then (Section  \@ref(demo)) we demonstrate their application to fitting simulated data, before extracting and comparing posterior distributions.
By comparing the resulting inferences to those from NUTS, we hope to help you make informed decisions about which algorithm to use in your applied problem.

# Alternative approximate inference algorithms {#other}

Here we describe three alternative approximate Bayesian inference algorithms that are available to use in `brms`, and therefore also available in `epidist`.
It's worth noting that further inference algorithms may have become available since this vignette was last updated.
Please check `brms` package updates if interested!

## Laplace method

The Laplace method approximates a posterior distribution by a Gaussian distribution centred (by default) at the posterior mode.
In Stan, the Gaussian approximation is constructed on the unconstrained parameter space (as the domain of a Gaussian distribution is the real line).
Samples from the Gaussian approximation may then be transformed to the constrained parameter space.
To access the Laplace method, specify `algorithm = "laplace"` within `brms::brm()`.
See the section [Laplace Sampling](https://mc-stan.org/docs/cmdstan-guide/laplace_sample_config.html) of the `CmdStan` User's Guide for more information.

## Variational inference using ADVI

Automatic differentiation variational inference [ADVI; @kucukelbir2017advi] is a type of variational inference [VI; @blei2017variational] algorithm.
VI works by restricting to a family of distributions, and then selecting the member of that family which is the most similar to the posterior distribution.
Most commonly, and in Stan, (dis-)similarity is measured using the Kullback–Leibler (KL) divergence.
There are two options for the family of distributions, either a fully factorised Gaussian with `algorithm = "meanfield"` or a Gaussian with a full-rank covariance matrix with `algorithm = "fullrank"`.
See the section [Variational Inference using ADVI](https://mc-stan.org/docs/cmdstan-guide/variational_config.html) of the `CmdStan` User's Guide for more information.

## Pathfinder

Pathfinder is a method closely related to variational inference, which has been more recently developed by @zhang2022pathfinder.
It works by generating Gaussian approximations along each step of an iterative optimisation algorithm (such as L-BFGS).
The KL divergence from each approximation to the posterior is measured, with the best approximation chosen.
Pareto smoothed importance sampling [PSIS; @vehtari2015pareto] is optionally used to resample draws from the chosen Gaussian distribution.
When multiple paths are specified (using `num_paths`) then the Pathfinder algorithm is run multiple times, initialising the optimisation at different points.
The resulting approximation is a mixture of Gaussian distributions, rather than a single Gaussian distribution.
See the section [Pathfinder Method for Approximate Bayesian Inference](https://mc-stan.org/docs/cmdstan-guide/pathfinder_config.html) of the `CmdStan` User's Guide for more information.

# Demonstration {#demo}

In this demonstration, we use the following packages:


``` r
library(epidist)
library(ggplot2)
library(dplyr)
library(purrr)
library(tidyr)
library(tibble)
library(tidybayes)
library(cmdstanr) # nolint
```

To access the approximate inference methods used in this vignette we will need to use the `cmdstanr` backend for `brms` (we generally recommend using this backend for fitting models).
To do this, we first need to install CmdStan (see the README for more details).
We can check we have everything we need as follows:


``` r
cmdstanr::cmdstan_version()
#> [1] "2.40.0"
```

We can simulate data to use for fitting models.
The example data simulation process follows that used in the [Getting started with epidist](https://epidist.epinowcast.org/articles/epidist.html#data) vignette, so we will not detail exactly what is happening here, but please consult that vignette if interested:

<details><summary>Click to expand for data simulation code</summary>


``` r
meanlog <- 1.8
sdlog <- 0.5
obs_time <- 25
sample_size <- 200

obs_cens_trunc_samp <- simulate_gillespie(seed = 101) |>
  simulate_secondary(
    meanlog = meanlog,
    sdlog = sdlog
  ) |>
  mutate(
    ptime_lwr = floor(.data$ptime),
    ptime_upr = .data$ptime_lwr + 1,
    stime_lwr = floor(.data$stime),
    stime_upr = .data$stime_lwr + 1,
    obs_time = obs_time
  ) |>
  filter(.data$stime_upr <= .data$obs_time) |>
  slice_sample(n = sample_size, replace = FALSE)
```

</details>

We now prepare the data for fitting with the marginal model.
We first pre-compile the Stan model so that compilation time is excluded from the timing comparisons below.


``` r
linelist_data <- as_epidist_linelist_data(
  obs_cens_trunc_samp$ptime_lwr,
  obs_cens_trunc_samp$ptime_upr,
  obs_cens_trunc_samp$stime_lwr,
  obs_cens_trunc_samp$stime_upr,
  obs_time = obs_cens_trunc_samp$obs_time
)

data <- as_epidist_marginal_model(linelist_data)

# Pre-compile the model so compilation time is not included
# in the timing comparisons
fit_compile <- epidist(
  data = data, backend = "cmdstanr",
  chains = 1, iter = 5
)
```

We now perform inference with HMC:


``` r
t <- proc.time()
fit_hmc <- update(fit_compile, chains = 4, iter = 2000)
time_hmc <- proc.time() - t
```

Note that for clarity above we specify `algorithm = "sampling"`, but if you were to call `epidist(data = data)` the result would be the same since `"sampling"` (i.e. HMC) is the default value for the `algorithm` argument.

Now, we fit^[Note that in this section, and above for the MCMC, the output of the call is hidden, but if you were to call these functions yourself they would display information about the fitting procedure as it occurs] the same marginal model using each method in Section \@ref(other).
To match the four Markov chains of length 1000 in HMC above, we then draw 4000 samples from each approximate posterior.


``` r
t <- proc.time()
fit_laplace <- update(
  fit_compile,
  algorithm = "laplace", draws = 4000
)
time_laplace <- proc.time() - t

t <- proc.time()
fit_advi <- update(
  fit_compile,
  algorithm = "meanfield", draws = 4000
)
time_advi <- proc.time() - t
```

For the Pathfinder algorithm we will set `num_paths = 1`.


``` r
t <- proc.time()
fit_pathfinder <- update(
  fit_compile,
  algorithm = "pathfinder", draws = 4000,
  chains = 1
)
time_pathfinder <- proc.time() - t
```

We now extract the posterior distribution for the delay parameters from the fitted model for each inference method.
Thankfully, each algorithm is implemented to sample draws from the posterior distribution, making post-processing straightforward.

<details><summary>Click to expand for code to extract posterior draws</summary>


``` r
fits <- list(
  HMC = fit_hmc,
  Laplace = fit_laplace,
  ADVI = fit_advi,
  Pathfinder = fit_pathfinder
)

draws <- imap(fits, function(fit, name) {
  draws <- fit |>
    epidist_strata() |>
    add_delay_parameter_draws(fit) |>
    add_summaries() |>
    as.data.frame() |>
    pivot_longer(
      cols = c("mu", "sigma", "mean", "sd"),
      names_to = "parameter",
      values_to = "value"
    ) |>
    filter(parameter %in% c("mu", "sigma")) |>
    mutate(method = as.factor(name))
  return(draws)
})

draws <- bind_rows(draws)
```

</details>

## Comparison of parameter posterior distributions

The mean estimated value of each parameter, from each method, is as follows.


``` r
pars <- draws |>
  group_by(method, parameter) |>
  summarise(value = mean(value)) |>
  pivot_wider(names_from = parameter, values_from = value) |>
  ungroup()

pars
#> # A tibble: 4 × 3
#>   method        mu sigma
#>   <fct>      <dbl> <dbl>
#> 1 HMC         1.75 0.460
#> 2 Laplace     1.74 0.455
#> 3 ADVI        1.77 0.420
#> 4 Pathfinder  1.74 0.456
```

More comprehensively, the estimated posterior distributions are shown in Figure \@ref(fig:posterior).

<details><summary>Click to expand for code to create posterior distribution plot</summary>


``` r
p_posterior <- draws |>
  ggplot(aes(x = value, col = method)) +
  stat_slabinterval(density = "histogram", breaks = 30, alpha = 0.8) +
  scale_colour_manual(values = c("#56B4E9", "#009E73", "#E69F00", "#CC79A7")) +
  facet_grid(method ~ parameter, scales = "free_x") +
  theme_minimal() +
  guides(fill = "none") +
  labs(x = "", y = "", col = "Method") +
  theme(legend.position = "bottom")
```

</details>

(ref:posterior) Estimated posterior distributions for the `mu` and `sigma` parameters using each inference method, shown using `tidybayes::stat_slabinterval()`.


``` r
p_posterior
```

<div class="figure">
<img src="figures/approx-inference-posterior-1.png" alt="(ref:posterior)"  />
<p class="caption">(\#fig:posterior)(ref:posterior)</p>
</div>

## Comparison of resulting delay distributions

Figure \@ref(fig:delay-pdf) shows how the different `mu` and `sigma` posterior mean estimates from each inference method alter an estimated delay distribution.

<details><summary>Click to expand for code to create delay PDF plot</summary>


``` r
p_delay_pdf <- pmap_df(
  filter(pars), ~ tibble(
    x = seq(0, 25, by = 0.1),
    method = ..1, density = dlnorm(x, ..2, ..3)
  )
) |>
  ggplot(aes(x = x, y = density, col = method)) +
  geom_line(linewidth = 1) +
  scale_color_manual(values = c("#56B4E9", "#009E73", "#E69F00", "#CC79A7")) +
  theme_minimal() +
  labs(x = "", y = "", col = "Method") +
  theme(legend.position = "bottom")
```

</details>

(ref:delay-pdf) Delay probability density functions obtained based on the posterior mean estimated `mu` and `sigma` parameters.


``` r
p_delay_pdf
```

<div class="figure">
<img src="figures/approx-inference-delay-pdf-1.png" alt="(ref:delay-pdf)"  />
<p class="caption">(\#fig:delay-pdf)(ref:delay-pdf)</p>
</div>

## Comparison of time taken

In this example, HMC took a longer time to run than the other methods and Pathfinder was the fastest running method.
That said, even for HMC the computation time in this case is unlikely to be prohibitive.


``` r
times <- list(
  HMC = time_hmc,
  Laplace = time_laplace,
  ADVI = time_advi,
  Pathfinder = time_pathfinder
)

times
#> $HMC
#>    user  system elapsed
#>  12.551   0.206  12.771
#>
#> $Laplace
#>    user  system elapsed
#>   1.296   0.110   1.502
#>
#> $ADVI
#>    user  system elapsed
#>   1.029   0.082   1.122
#>
#> $Pathfinder
#>    user  system elapsed
#>   0.396   0.083   0.528
```

# Conclusion

The range of alternative approximation algorithms available, and their ease of use, is an attractive feature of `brms`.
We found that these algorithms do produce reasonable approximations in far less time than HMC.
Of course, this vignette only includes one example, so results may differ for other models and datasets.

## References {-}