---
title: "Model summaries for a Bayesian linear regression"
author: "Sam Clifford"
date: "`r Sys.Date()`"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Model summaries for a Bayesian linear regression}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

The packages `mmcc` creates tidy summaries of Bayesian models, in the fashion of `broom`, with one important difference - `mmcc` uses a `data.table` instead of a `tibble`, due to the size of the output that is all too easily possible in Bayesian models.
  
The aim of this vignette is to demonstrate how to use the two key functions of `mmcc`, `mcmc_to_dt()` and `tidy` (which actually calls `mmcc:::tidy.mcmc.list` under the hood).

```{r setup, echo = FALSE}
knitr::opts_chunk$set(
collapse = TRUE,
comment = "#>"
)
```

First, we simulate some data to fit a Bayesian model to.

``` {r data-simulate, fig.align="center"}
set.seed(4000)
N <- 20
x <- sort(runif(n = N))
y <- rnorm(n = N, 
           mean = 2*x + 1, 
           sd = 0.25)
dat <- data.frame(x = x, y = y)

library(ggplot2)
ggplot(data = dat,
       aes(x = x, 
           y = y)) +
    geom_point() +
    theme_bw()
```

Then, we simulate some values for predicting

``` {r sim-predict}
M <- 10
x_pred <- seq(from = min(x), 
              to = max(x), 
              length.out = M)
```

Next, we fit the model, specified as

```{r specify-model}
jags_model <- 
"model{
    # model block
    for (i in 1:n){
        y[i] ~ dnorm(mu[i], tau_y)
        mu[i] <- beta_0 + beta_1*x[i]
    }

    # prediction block
    for (i in 1:m){
        y_pred[i] ~ dnorm(mu_pred[i], tau_y)
        mu_pred[i] <- beta_0 + beta_1*x_pred[i]
    }

    # priors
    beta_0 ~ dunif(-1e12, 1e12)
    beta_1 ~ dunif(-1e12, 1e12)
    tau_y <- exp(2*log_sigma)
    log_sigma ~ dunif(-1e12, 1e12)
}"
```

and then generate the `mcmc_object` with the `rjags` package.

``` {r rjags-run}
library(rjags)
model <- jags.model(file = textConnection(jags_model),
                    data = list(n = N,
                                x = x,
                                y = y,
                                m = M,
                                x_pred = x_pred),
                    n.chains = 3)
```

We draw burn-in samples and posterior inference samples for all terms in the model.

``` {r rjags-sample}
burn <- jags.samples(model = model,
                     variable.names = c("beta_0", 
                                        "beta_1", 
                                        "tau_y", 
                                        "mu"),
                     n.iter = 5000)

samples <- coda.samples(model = model,
                        variable.names = c("beta_0", 
                                           "beta_1", 
                                           "tau_y", 
                                           "mu_pred", 
                                           "y_pred"),
                        n.iter = 10000)

```

We can now convert the posterior samples to a `data.table` and summarise the regression parameters. A `data.table` object is very useful in this case when you have many samples for many parameters.

``` {r rjags-summarise}
library(mmcc)

# convert to a data.table
samples_dt <- mcmc_to_dt(samples)
samples_dt

pars_dt <- tidy(samples, 
                conf_level = 0.95, 
                colnames = c("beta_0", 
                             "beta_1", 
                             "tau_y"))

pars_dt
```

Summarise the line of best fit, `mu`, and the predictions, `y_pred`,

``` {r rjags-summarise-mu-y}
mu_dt <- tidy(samples, 
              conf_level = 0.95, 
              colnames = "mu_pred")

y_dt <- tidy(samples, 
             conf_level = 0.95, 
             colnames = "y_pred")
```

For plotting, we add the prediction $\boldsymbol{x}$ values to these data tables.

``` {r rjags-add-preds}
mu_dt[ , x:= x_pred]
y_dt[ , x:= x_pred]
y_dt
```

Now we'll generate a plot that shows the data, a 95% credible interval for the predictions, ${\hat{\bm{y}}}_{pred}$, and a 95% credible interval for their means, ${\hat{\bm{\mu}}}_{pred}$.

``` {r mmcc-plot-cr-i, fig.align="center", echo=F}
ggplot(data = dat,
       aes(x = x)) +
    geom_point(aes(y = y)) +
    geom_ribbon(data = mu_dt,
                aes(ymin = `2.5%`,
                    ymax = `97.5%`),
                fill = "salmon",
                alpha = 0.5) +
    geom_ribbon(data = y_dt,
                aes(ymin = `2.5%`,
                    ymax = `97.5%`),
                fill = "lightskyblue",
                alpha = 0.25) +
    geom_line(data = y_dt,
              aes(y = mean)) +
    theme_bw()
```

If we tidy the `samples` object, we can look at the distribution of values

``` {r mmcc-value-distr, fig.height=2, fig.width=7}
tidy_samples <- mcmc_to_dt(samples, 
                           colnames = c("beta_0", 
                                        "beta_1", 
                                        "tau_y"))

ggplot(data = tidy_samples, 
       aes(x = value)) +
    geom_density(color = "black", 
                 fill = "grey90") +
    facet_wrap(~parameter,
               nrow = 1,
               scales = "free") +
    theme_bw() +
    geom_segment(data = pars_dt,
                   aes(x = `2.5%`,
                       xend = `97.5%`),
                 y = 0, 
                 yend = 0,
                 size = 2) +
    geom_point(data = pars_dt,
               aes(x = mean),
               y = 0,
               color = "white")
```

We can also thin to create trace plots and plot per chain

``` {r}
tidy_samples_10 <- thin_dt(tidy_samples, thin = 10)

ggplot(data=tidy_samples_10, aes(x=iteration, y=value)) +
    geom_line(aes(group=chain, color=factor(chain))) +
    facet_wrap( ~ parameter, ncol=1, scales="free_y") +
    theme_bw() +
    theme(legend.position = "bottom") +
    scale_color_discrete(name="Chain")
```