---
engine: knitr
---

```{python}
#| label: setup
#| include: false

import liesel.goose as gs
import pandas as pd

gs.Summary.__repr__ = gs.Summary._repr_html_
gs.Summary._repr_markdown_ = gs.Summary._repr_html_
pd.options.display.float_format = "{:.3f}".format
pd.options.display.html.border = 0
```

# Comparing samplers

In this tutorial, we are comparing two different sampling schemes on the `mcycle` dataset with a Gaussian location-scale regression model and two splines for the mean and the standard deviation. The `mcycle` dataset is a "data frame giving a series of measurements of head acceleration in a simulated motorcycle accident, used to test crash helmets" (from the help page). It contains the following two variables:

- `times`: in milliseconds after impact
- `accel`: in g

We start off in R by loading the dataset and setting up the model with the `rliesel::liesel()` function.

```{r}
#| label: model

library(MASS)
library(rliesel)

data(mcycle)
with(mcycle, plot(times, accel))

model <- liesel(
  response = mcycle$accel,
  distribution = "Normal",
  predictors = list(
    loc = predictor(~ s(times)),
    scale = predictor(~ s(times), inverse_link = "Exp")
  ),
  data = mcycle
)
```

## Metropolis-in-Gibbs

First, we try a Metropolis-in-Gibbs sampling scheme with IWLS kernels for the regression coefficients ($\boldsymbol{\beta}$) and Gibbs kernels for the smoothing parameters ($\tau^2$) of the splines.

```{python}
#| label: iwls-sampling

import liesel.model as lsl

model = r.model

builder = lsl.dist_reg_mcmc(model, seed=42, num_chains=4)
builder.set_duration(warmup_duration=5000, posterior_duration=1000)
builder.show_progress = False

engine = builder.build()
engine.sample_all_epochs()
```

Clearly, the performance of the sampler could be better, especially for the intercept of the mean. The corresponding chain exhibits a very strong autocorrelation.

```{python}
#| label: iwls-summary

import liesel.goose as gs

results = engine.get_results()
gs.Summary(results)
```

```{python}
#| label: iwls-traces

fig = gs.plot_trace(results, "loc_p0_beta")
fig = gs.plot_trace(results, "loc_np0_tau2")
fig = gs.plot_trace(results, "loc_np0_beta")

fig = gs.plot_trace(results, "scale_p0_beta")
fig = gs.plot_trace(results, "scale_np0_tau2")
fig = gs.plot_trace(results, "scale_np0_beta")
```

To confirm that the chains have converged to reasonable values, here is a plot of the estimated mean function:

```{python}
#| label: iwls-summary-df

summary = gs.Summary(results).to_dataframe().reset_index()
```

```{r}
#| label: iwls-spline

library(dplyr)
library(ggplot2)
library(reticulate)

summary <- py$summary

beta <- summary %>%
  filter(variable == "loc_np0_beta") %>%
  group_by(var_index) %>%
  summarize(mean = mean(mean)) %>%
  ungroup()

beta <- beta$mean
X <- py_to_r(model$vars["loc_np0_X"]$value)
f <- X %*% beta

beta0 <- summary %>%
  filter(variable == "loc_p0_beta") %>%
  group_by(var_index) %>%
  summarize(mean = mean(mean)) %>%
  ungroup()

beta0 <- beta0$mean

ggplot(data.frame(times = mcycle$times, mean = beta0 + f)) +
  geom_line(aes(times, mean), color = palette()[2], size = 1) +
  geom_point(aes(times, accel), data = mcycle) +
  ggtitle("Estimated mean function") +
  theme_minimal()
```

## NUTS sampler

As an alternative, we try using NUTS kernels for all parameters.
To do so, we first need to log-transform the smoothing parameters.
This is the model graph before the transformation:

```{python}
#| label: untransformed-graph

lsl.plot_vars(model)
```

To transform the smoothing parameters with the method {meth}`.Var.transform`,
we need to retrieve the nodes and vars form the model. This is necessary, because
while they are part of a model, the inputs and outputs of nodes and vars cannot
be changed.
We retrieve the nodes and vars using {meth}`.Model.pop_nodes_and_vars`,
which renders the model empty.

Note, RLiesel automatically populates the parameter variables with intormation
how to conduct the MCMC inference. To trasform the variables, we can drop it.
We setup the MCMC sampler in the next chunk.

After transformation, there are two additional nodes in the new model graph.


```{python}
#| label: transformed-graph

import tensorflow_probability.substrates.jax.bijectors as tfb

nodes, _vars = model.pop_nodes_and_vars()

_vars["loc_np0_tau2"].transform(tfb.Exp(), inference='drop')
_vars["scale_np0_tau2"].transform(tfb.Exp(), inference='drop')

gb = lsl.GraphBuilder().add(_vars["response"])

model = gb.build_model()
lsl.plot_vars(model)
```

Now we can set up the NUTS sampler. In complex models like this one it can be
very beneficial to use individual NUTS samplers for blocks of parameters. This
is pretty much the same strategy that we apply to the IWLS sampler, too.



```{python}
#| label: nuts-sampling

builder = gs.EngineBuilder(seed=42, num_chains=4)

builder.set_model(gs.LieselInterface(model))

# add NUTS kernels
parameters = [name for name, var in model.vars.items() if var.parameter]

for parameter in parameters:
  builder.add_kernel(gs.NUTSKernel([parameter]))


builder.set_initial_values(model.state)

builder.set_epochs(
  gs.stan_epochs(warmup_duration=5000, posterior_duration=1000, init_duration=750, term_duration=500)
)

builder.show_progress = False

engine = builder.build()
engine.sample_all_epochs()
```

The NUTS sampler overall seems to do a good job - and even yields higher
effective sample sizes than the IWLS sampler, especially for the spline
coefficients of the scale model.

```{python}
#| label: nuts-summary
#| output: asis

results = engine.get_results()
gs.Summary(results)
```

```{python}
#| label: nuts-traces

fig = gs.plot_trace(results, "loc_p0_beta")
fig = gs.plot_trace(results, "loc_np0_tau2_transformed")
fig = gs.plot_trace(results, "loc_np0_beta")

fig = gs.plot_trace(results, "scale_p0_beta")
fig = gs.plot_trace(results, "scale_np0_tau2_transformed")
fig = gs.plot_trace(results, "scale_np0_beta")
```

Again, here is a plot of the estimated mean function:

```{python}
#| label: nuts-summary-df

summary = gs.Summary(results).to_dataframe().reset_index()
```

```{r}
#| label: nuts-spline

library(dplyr)
library(ggplot2)
library(reticulate)

summary <- py$summary
model <- py$model

beta <- summary %>%
  filter(variable == "loc_np0_beta") %>%
  group_by(var_index) %>%
  summarize(mean = mean(mean)) %>%
  ungroup()

beta <- beta$mean
X <- model$vars["loc_np0_X"]$value
f <- X %*% beta

beta0 <- summary %>%
  filter(variable == "loc_p0_beta") %>%
  group_by(var_index) %>%
  summarize(mean = mean(mean)) %>%
  ungroup()

beta0 <- beta0$mean

ggplot(data.frame(times = mcycle$times, mean = beta0 + f)) +
  geom_line(aes(times, mean), color = palette()[2], size = 1) +
  geom_point(aes(times, accel), data = mcycle) +
  ggtitle("Estimated mean function") +
  theme_minimal()
```
