How to sample

We consider the random regression example from the RTMB introduction vignette:

## Data
data(ChickWeight)

## Parameters and random effects
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
)

## Negative log likelihood
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
}

We build the TMB model object (following the original example, we specify the random effects although that is redundant here):

obj <- MakeADFun(f, parameters, random=c("a", "b"))

A full Bayesian analysis is run by:

out <- tmbstan(obj, chains=1, seed=1)

## 
## SAMPLING FOR MODEL 'tmb_generic' NOW (CHAIN 1).
## Chain 1: 
## Chain 1: Gradient evaluation took 8.3e-05 seconds
## Chain 1: 1000 transitions using 10 leapfrog steps per transition would take 0.83 seconds.
## Chain 1: Adjust your expectations accordingly!
## Chain 1: 
## Chain 1: 
## Chain 1: Iteration:    1 / 2000 [  0%]  (Warmup)
## Chain 1: Iteration:  200 / 2000 [ 10%]  (Warmup)
## Chain 1: Iteration:  400 / 2000 [ 20%]  (Warmup)
## Chain 1: Iteration:  600 / 2000 [ 30%]  (Warmup)
## Chain 1: Iteration:  800 / 2000 [ 40%]  (Warmup)
## Chain 1: Iteration: 1000 / 2000 [ 50%]  (Warmup)
## Chain 1: Iteration: 1001 / 2000 [ 50%]  (Sampling)
## Chain 1: Iteration: 1200 / 2000 [ 60%]  (Sampling)
## Chain 1: Iteration: 1400 / 2000 [ 70%]  (Sampling)
## Chain 1: Iteration: 1600 / 2000 [ 80%]  (Sampling)
## Chain 1: Iteration: 1800 / 2000 [ 90%]  (Sampling)
## Chain 1: Iteration: 2000 / 2000 [100%]  (Sampling)
## Chain 1: 
## Chain 1:  Elapsed Time: 3.089 seconds (Warm-up)
## Chain 1:                1.618 seconds (Sampling)
## Chain 1:                4.707 seconds (Total)
## Chain 1:

Improving sampling speed

MCMC runs for large models can take a very long time. A simple strategy for improving sampling speed is to speed up the objective function and gradient by enabling vectorization.

TapeConfig(vectorize="enable")  ## Enable vectorization
obj_vect <- MakeADFun(f, parameters, random=c("a", "b"))
TapeConfig(vectorize="disable") ## Restore previous settings

We can check the effect of this optimization before starting the potentially long MCMC run. For accurate timing, excluding R overhead, we use the RTMB tape timer() method, and report an average over 1e4 evaluations of the total tape pass:

## Timing without vectorization
t <- GetTape(obj)$timer(rep=1e4,total=TRUE)
## Timing with vectorization
t_vect <- GetTape(obj_vect)$timer(rep=1e4,total=TRUE)
## speedup ?
c(speedup = unname(t / t_vect))

##  speedup 
## 2.745817

In this case vectorization appears to be beneficial. Let’s compare the MCMC runs:

system.time(out <- tmbstan(obj, chains=1, seed=1, refresh = 0))

##    user  system elapsed 
##   5.029   0.029   5.059

system.time(out_vect <- tmbstan(obj_vect, chains=1, seed=1, refresh = 0))

##    user  system elapsed 
##   2.220   0.002   2.222