---
title: "Introduction to RTMB"
author: "Kasper Kristensen"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{Introduction to RTMB}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include=FALSE, message=FALSE}
##library(knitr)
library(RTMB)
set.seed(1)
formals(MakeADFun)$silent <- TRUE
```

# Introduction

This vignette demonstrates basic features of `RTMB` by implementing a *random regression* model from scratch and fit it to a built-in R dataset.

# Random regression

## Loading the package

```{r}
library(RTMB)
```

## Getting the data

We'll take as starting point the built-in data set 'ChickWeight':

```{r}
data(ChickWeight)
```

The data contains `weight` of 50 chicks by `Time` and individual identifier `Chick`. A useful plot of growth by individual is available from the help page by running `example(ChickWeight)`.

## Defining a model

Although not the most natural based on the plot, we'll setup a *random regression model*.
It is defined by first drawing slopes and intercepts by individual from an underlying distribution

\[ a_1,...,a_{50} \sim N(\mu_a , \sigma_a^2) \]

\[ b_1,...,b_{50} \sim N(\mu_b , \sigma_b^2) \]

and then state a normal regression given these random coefficients:

\[ \text{weight}_{i} \sim N(a_{\text{Chick}_i} * \text{Time}_i + b_{\text{Chick}_i}, \sigma^2) \]

## Implementing the model in RTMB

To implement the model in RTMB we have to set up the objective function and define the parameters and random effects.

Parameter objects are gathered in a list that also serves as initial guess when fitting the model:

```{r}
parameters <- list(
    mua=0,          ## Mean slope
    sda=1,          ## Std of slopes
    mub=0,          ## Mean intercept
    sdb=1,          ## Std of intercepts
    sdeps=1,        ## Residual Std
    a=rep(0, 50),   ## Random slope by chick
    b=rep(0, 50)    ## Random intercept by chick
)
```

The objective function takes as input such a parameter list, and is defined by

```{r}
f <- function(parms) {
    getAll(ChickWeight, parms, warn=FALSE)
    ## Optional (enables extra RTMB features)
    weight <- OBS(weight)
    ## Initialize joint negative log likelihood
    nll <- 0
    ## Random slopes
    nll <- nll - sum(dnorm(a, mean=mua, sd=sda, log=TRUE))
    ## Random intercepts
    nll <- nll - sum(dnorm(b, mean=mub, sd=sdb, log=TRUE))
    ## Data
    predWeight <- a[Chick] * Time + b[Chick]
    nll <- nll - sum(dnorm(weight, predWeight, sd=sdeps, log=TRUE))
    ## Get predicted weight uncertainties
    ADREPORT(predWeight)
    ## Return
    nll
}
```

This function calculates the negative log likelihood `nll` using straight forward R operations corresponding exactly to the model definition.
In addition, some RTMB specific statements are used:

- The `getAll` function makes all the list elements of data and parameters visible inside the function, so that one can write e.g. `weight` rather than `ChickWeight$weight`.
- The `weight <- OBS(weight)` statement tells RTMB that that `weight` is the response. This is needed to enable automatic simulation and residual calculations from the model object.
- The `ADREPORT(predWeight)` statement tells RTMB that we want uncertainties for this intermediate calculation.

The objective function `f` is processed by RTMB using the call

```{r}
obj <- MakeADFun(f, parameters, random=c("a", "b"))
```

where we also specify that `a` and `b` are random effects.

## Fitting the model

We optimize the model using `nlminb` (or any other suitable gradient based optimizer in R)

```{r}
opt <- nlminb(obj$par, obj$fn, obj$gr)
```

## Calculating model output

Uncertainties are now calculated using

```{r}
sdr <- sdreport(obj)
sdr
```

By default, the shown output is very brief, containing only the model parameters. The `sdr` object contains much more.
It is often convenient to inspect parameter estimates and other output as a list similar to the one containing the parameters.
For instance, to get parameters estimates and standard errors as separate lists use:

```r
as.list(sdr, "Est") ## parameter estimates
as.list(sdr, "Std") ## parameter uncertainties
```

Pass `report=TRUE` to get `ADREPORT`ed quantities:

```r
as.list(sdr, "Est", report=TRUE) ## ADREPORT estimates
as.list(sdr, "Std", report=TRUE) ## ADREPORT uncertainties
```

## Simulating from the model object

New datasets can be generated from the estimated model using  `obj$simulate()` assuming that the model is implemented in accordance with the principles defined in the help page `?Simulation`.

## Checking correctness of the implementation using simulation

When building random effect models from scratch, many mistakes can be made leading to wrong results. We always recommend to run the model through the completely automatic consistency check (this requires that `obj$simulate()` works for the implementation). By default, the check simulates 100 datasets and calculates gradients for each replicate. A standard output tells you whether the implementation is consistent with simulation (message 'simulation appears to be correct'). It also gives an idea of the parameter bias usually caused by the Laplace approximation (not an issue for the random regression model for which the Laplace approximation is exact). We run the standard check by:

```{r}
set.seed(1)
chk <- checkConsistency(obj)
chk
```

As expected for this case, everything looks fine.
A complete simulation study, that re-estimates parameters for each replicate, can also be performed although that takes longer to run:

```r
chk <- checkConsistency(obj, estimate=TRUE)
summary(chk)
```

For more details we refer to the help page `?TMB::checkConsistency`.

## Goodness-of-fit using quantile residuals

Quantile residuals can be generated automatically using the `oneStepPredict` function. These residuals are conceptually superior to other methods (e.g. Pearson residuals), but much trickier to calculate. It is very important to specify an appropriate `method` (see `?TMB::oneStepPredict`) because using an inappropriate method can give wrong residuals. For the random regression model the 'fullGaussian' method is computationally exact.

```{r}
osa <- oneStepPredict(obj, method="fullGaussian", discrete=FALSE)
qqnorm(osa$res); abline(0,1)
```

The argument `discrete=FALSE` is necessary in this case because the data has duplicates, which is a zero-probability event for continuous distributions.
If the model is correct, the residuals are standard, independent normally distributed, which is obviously not the case here.

# Debugging

For the random regression model, we could run `MakeADFun` without problems.
However, model implementation might not always work as smoothly. In case of errors it is useful to test the implementation step-wise.

1. Going back to our objective function `f`, first step is to check that you can evaluate the function as a normal R function:

    ```r
    f(parameters)
    ```

    An error at this point is obviously not due to `RTMB`.

2. Next, it is useful to check that `MakeADFun` can be run *without* random effects:

    ```r
    obj <- MakeADFun(f, parameters)
    ```

    Should an error occur at this point, you can enable standard R debugging `debug(f)` and run `MakeADFun` again to figure out which code line caused the error. A common cause is to that an unsupported (e.g. non-differentiable) operation has been used.

3. Once `obj` has been constructed successfully, you should evaluate it

    ```r
    obj$fn()
    ```

    and verify that it gives the same result as `f(parameters)`.

# Probabilistic syntax

The random regression model could alternatively have been written using the RTMB 'tilde operator' (`%~%`):

```{r}
f2 <- function(parms) {
    getAll(ChickWeight, parms, warn=FALSE)
    ## Optional (enables extra RTMB features)
    weight <- OBS(weight)
    ## Random slopes
    a %~% dnorm(mean=mua, sd=sda)
    ## Random intercepts
    b %~% dnorm(mean=mub, sd=sdb)
    ## Data
    predWeight <- a[Chick] * Time + b[Chick]
    weight %~% dnorm(predWeight, sd=sdeps)
    ## Get predicted weight uncertainties
    ADREPORT(predWeight)
}
```

This syntax is closer to other probabilistic languages (e.g. BUGS, JAGS and Stan).
But more importantly, it prevents some very common TMB mistakes, by passing `log=TRUE` automatically and making sure the sign of the objective function is correct.

Otherwise, `f2` is identical to `f`, and the model object can be constructed and fitted in the same way:

```{r}
obj <- MakeADFun(f2, parameters, random=c("a", "b"))
```

# Using an explicit data argument

A common cause of confusion is that the RTMB version of `MakeADFun` does not have a `data` argument. What if we want to change the data?
This can be handled in R using *closures*. Start by adding an explicit data argument to the objective function:

```{r}
f3 <- function(parms, data) {
    getAll(data, parms, warn=FALSE)
    ## Optional (enables extra RTMB features)
    weight <- OBS(weight)
    ## Random slopes
    a %~% dnorm(mean=mua, sd=sda)
    ## Random intercepts
    b %~% dnorm(mean=mub, sd=sdb)
    ## Data
    predWeight <- a[Chick] * Time + b[Chick]
    weight %~% dnorm(predWeight, sd=sdeps)
    ## Get predicted weight uncertainties
    ADREPORT(predWeight)
}
```

A general function to combine this objective with a *specific* data set is:

```{r}
cmb <- function(f, d) function(p) f(p, d)
```

We can now easily create model objects using different datasets:

```{r}
## Using the original ChickWeight
obj <- MakeADFun(cmb(f3, ChickWeight), parameters, random=c("a", "b"))
## Using a new dataset
ChickWeightNew <- transform(ChickWeight, weight=log(weight))
obj <- MakeADFun(cmb(f3, ChickWeightNew), parameters, random=c("a", "b"))
```

# Model predictions

We've already seen how `ADREPORT()` can be used to make predictions of derived quantities e.g. `predWeight` corresponding to each of our measurements.
But what if we need predictions outside of where we have data? For example, some natural prediction tasks could be:

1. Predict, with standard deviations, the weight of chick number `10` for times `20:24` using diet `3`.
2. Same as (1) for an *unseen chick* number `51`.

A general approach to make such predictions is to modify the objective function to handle rows with missing (`NA`) response, and then augment the original data with missing responses where predictions are needed.
We'll take as starting point the `f3()` implementation and modify (changes are marked with a comment):

```{r}
f_pred <- function(parms, data) {
    getAll(data, parms, warn=FALSE)
    ## Optional (enables extra RTMB features)
    weight <- OBS(weight)
    ## Random slopes
    a %~% dnorm(mean=mua, sd=sda)
    ## Random intercepts
    b %~% dnorm(mean=mub, sd=sdb)
    ## Data
    predWeight <- a[Chick] * Time + b[Chick]
    mis <- is.na(weight)                                 ## <-- Get missing responses
    weight[!mis] %~% dnorm(predWeight[!mis], sd=sdeps)   ## <-- Likelihood of non-missing
    ## Get predicted weight uncertainties
    ADREPORT(predWeight[mis])                            ## <-- Only report missing
}
```

## Prediction task 1

For the first prediction task we create a `newdata` object with missing response:

```{r}
newdata <- data.frame(Time=20:24, Diet="3", Chick="10", weight=NA)
```

We then refit the model using the original data *augmented with* `newdata` (Note: refitting can be avoided - see next section).

```{r}
augdata <- rbind(ChickWeight, newdata)
obj <- MakeADFun(cmb(f_pred, augdata), parameters, random=c("a", "b"))
opt <- nlminb(obj$par, obj$fn, obj$gr)
sdr <- sdreport(obj)
cbind(newdata, summary(sdr, select="report"))
```

## Prediction task 2

The second prediction task is slightly more tricky as it involves an unseen chick number 51. That is, we need to add additional random effects `a` and `b` corresponding to this individual.

```{r}
newdata <- transform(newdata, Chick="51")
augdata <- rbind(ChickWeight, newdata)
parameters$a <- rep(0, nlevels(augdata$Chick))
parameters$b <- rep(0, nlevels(augdata$Chick))
```

This is all we need and we could in principle proceed as before by refitting the model and running `sdreport()`. However, a shortcut is worth mentioning: The previous fit of *parameters* and *hessian* can be re-used as these can't possibly change by adding missing responses to the data:

```{r}
par.old <- opt$par
hessian.old <- optimHess(par.old, obj$fn, obj$gr)
```

We can then pass this information to `sdreport` and skip the optimization step:

```{r}
obj <- MakeADFun(cmb(f_pred, augdata), parameters, random=c("a", "b"))
sdr <- sdreport(obj, par.old, hessian.old)
cbind(newdata, summary(sdr, select="report"))
```
