Restarting NIMBLE MCMC

Back to NIMBLE Vignettes

This document gives a sketch of how to restart a NIMBLE MCMC algorithm from where it left off, assuming you have to end your R session between runs. If R didn’t have to be restarted, you could simply use the reset = FALSE argument to the Cmcmc$run() method. But if the R session must be stopped and restarted, it’s a little more complicated.

This approach is not covered in the NIMBLE User Manual, and requires some knowledge of the internals and implementation of NIMBLE and NIMBLE’s MCMC.

Creating Model and MCMC

###########################
### Start new R session ###
###########################

library(nimble)

Suppose you have the code, constants, data, and inits necessary for creating a NIMBLE model, for MCMC sampling. Here’s a simple example below of Poisson regression.

Note that we’ll use exactly these same code, constants, data, and inits after restarting R, to create new model and mcmc objects.

code <- nimbleCode({
    mu ~ dnorm(0, 0.0001)
    sigma ~ dunif(0, 1000)
    for(i in 1:2) {
        beta[i] ~ dnorm(0, 0.0001)
    }
    for(i in 1:N) {
        epsilon[i] ~ dnorm(mu, sd = sigma)
        log(lambda[i]) <- beta[1] + beta[2] * x[i] + epsilon[i]
        y[i] ~ dpois(lambda[i])
    }
})

set.seed(0)
N <- 5
x <- rnorm(N, 0, 2)
lambda <- exp(4 + .3*x)
y <- rpois(N, lambda)

constants <- list(N = N, x = x)
data <- list(y = y)
inits <- list(mu = 0, sigma = 1, beta = rep(0,2), epsilon = rep(0, N))

Rmodel <- nimbleModel(code, constants, data, inits)

For good measure, run Rmodel$calculate() to make sure the model is fully initialized.

Rmodel$calculate()

## [1] -1724.11

Now we’ll create and configure the MCMC algorithm. To make for a more interesting example, demonstrating how to handle different samplers, we’ll assign a slice sampler on sigma.

We’ll leave the conjugate sampler that is assigned to mu, and theRW samplers assigned to the beta and epsilon terms.

You can decide whether or not to include WAIC by setting useWAIC.

useWAIC <- TRUE
conf <- configureMCMC(Rmodel, enableWAIC = useWAIC)

## ===== Monitors =====
## thin = 1: beta, mu, sigma
## ===== Samplers =====
## RW sampler (8)
##   - sigma
##   - beta[]  (2 elements)
##   - epsilon[]  (5 elements)
## conjugate sampler (1)
##   - mu

conf$printSamplers()

## [1] conjugate_dnorm_dnorm_identity sampler: mu
## [2] RW sampler: sigma
## [3] RW sampler: beta[1]
## [4] RW sampler: beta[2]
## [5] RW sampler: epsilon[1]
## [6] RW sampler: epsilon[2]
## [7] RW sampler: epsilon[3]
## [8] RW sampler: epsilon[4]
## [9] RW sampler: epsilon[5]

conf$removeSamplers("sigma")
conf$addSampler("sigma", "slice")
conf$printSamplers()

## [1] conjugate_dnorm_dnorm_identity sampler: mu
## [2] RW sampler: beta[1]
## [3] RW sampler: beta[2]
## [4] RW sampler: epsilon[1]
## [5] RW sampler: epsilon[2]
## [6] RW sampler: epsilon[3]
## [7] RW sampler: epsilon[4]
## [8] RW sampler: epsilon[5]
## [9] slice sampler: sigma

That looks good. We’ll use this MCMC.

Next, we’ll build the MCMC, compile the model and the MCMC, and run the MCMC for 10,000 iterations.

Rmcmc <- buildMCMC(conf)

Cmodel <- compileNimble(Rmodel)
Cmcmc <- compileNimble(Rmcmc, project = Rmodel)

set.seed(0)
Cmcmc$run(10000)

## |-------------|-------------|-------------|-------------|
## |-------------------------------------------------------|

## NULL

samples <- as.matrix(Cmcmc$mvSamples)

What to Save

Now, before ending the R session, we’ll need to save several things:

The MCMC samples we’ve already run.
The values of all non-data stochastic nodes in the model. That is, anything in the model code defined using ~, which is not specified as data. In this case, that’s mu, sigma, beta, and epsilon.
The internal state variables from all the MCMC samplers. This will be the slightly tricky part, but not too difficult.
R’s random number seed.
If using WAIC, the internal state variables for the WAIC calculation.

Samples

Straight forward, make sure to save the MCMC samples that you’ve already run.

Model Variables

For any parameters or latent variables that are being sampled in the model, you’ll need to save the “current” values of these which are stored in the Cmodel object. These can be accessed from the Cmodel object using the syntax Cmodel$variableName.

I’ll assume you know for your model what these are, but you can also look at everything that’s being sampled using. conf$printSamplers(). Another way to tell would be:

Cmodel$getNodeNames(stochOnly = TRUE, includeData = FALSE)

## [1] "mu"         "sigma"      "beta[1]"    "beta[2]"    "epsilon[1]"
## [6] "epsilon[2]" "epsilon[3]" "epsilon[4]" "epsilon[5]"

Or even nicer:

nodes <- Cmodel$getNodeNames(stochOnly = TRUE, includeData = FALSE)
Cmodel$getVarNames(nodes = nodes)

## [1] "mu"      "sigma"   "beta"    "epsilon"

So, we’ll extract these and store them. Let’s use a list to hold the model variables.

modelVariables <- list()

modelVariables$mu <- Cmodel$mu
modelVariables$sigma <- Cmodel$sigma
modelVariables$beta <- Cmodel$beta
modelVariables$epsilon <- Cmodel$epsilon

This list should now contain the values of all non-data stochastic variables in the model:

modelVariables

## $mu
## [1] 2.716613
## 
## $sigma
## [1] 0.05292085
## 
## $beta
## [1] 1.3819946 0.2780219
## 
## $epsilon
## [1] 2.680143 2.764016 2.756193 2.777305 2.630557

You’ll want to save this list modelVariables object.

MCMC Sampler State Variables

NIMBLE’s MCMC samplers usually contain one or more “state variables”, which define the state of the sampler algorithm, and are updated as the MCMC progresses. You’ll need to extract and save the values of these state variables for each sampling algorithm.

We’ll use an R list to hold the state variables. The first element will itself be a list of the state variables for the first sampler. The second element will record the state variables from the second sampler, etc.

samplerStates <- list()

Again, you can see how many samplers there are, and what types they are, using:

conf$printSamplers()

## [1] conjugate_dnorm_dnorm_identity sampler: mu
## [2] RW sampler: beta[1]
## [3] RW sampler: beta[2]
## [4] RW sampler: epsilon[1]
## [5] RW sampler: epsilon[2]
## [6] RW sampler: epsilon[3]
## [7] RW sampler: epsilon[4]
## [8] RW sampler: epsilon[5]
## [9] slice sampler: sigma

The first sampler is a conjugate sampler (assigned to mu). Conjugate samplers don’t require any internal state. So the first element in our list (the state variables for the first sampler) won’t require anything.

To tell what the “state variables” for any other sampler are, your best bet is looking at the underlying code for each sampler, available in MCMC_samplers.R.

For each sampling algorithm, for example the RW sampler, find corresponding nimbleFunction definition, in this case sampler_RW. Near the end of the definition there will be a reset method. In this case, for the RW sampler, this looks like:

reset = function() {
    scale <<- scaleOriginal 
    timesRan      <<- 0
    timesAccepted <<- 0
    timesAdapted  <<- 0
    ##scaleHistory  <<- scaleHistory * 0    ## scaleHistory
    gamma1 <<- 0
}

The “state variables” for the sampler are all the variables that are reset to some initial value in this function. So in this case: scale, timesRan, timesAccepted, timesAdapted, and gamma1. For each of these, you’ll need to extract the value of this variable from the compiled sampler algorithm, and save it.

There happen to be seven RW samplers (two for the beta terms, and one for each epsilon) so we’ll use a loop to extract these state variables from each compiled sampler algorithm, and store them into the appropriate element of our list. Loop over sampler numbers 2 through 8.

for(i in 2:8) {
    samplerStates[[i]] <- list()
    samplerStates[[i]]$scale <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "scale")
    samplerStates[[i]]$timesRan <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesRan")
    samplerStates[[i]]$timesAccepted <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesAccepted")
    samplerStates[[i]]$timesAdapted <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesAdapted")
    samplerStates[[i]]$gamma1 <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "gamma1")
}

For the 9th sampler, a slicesampler on sigma, the code for the sampler_slice reset() method looks like:

reset = function() {
    width        <<- widthOriginal
    timesRan     <<- 0
    timesAdapted <<- 0
    sumJumps     <<- 0
}

So, we’ll store these four state variables:

samplerStates[[9]] <- list()
samplerStates[[9]]$width <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "width")
samplerStates[[9]]$timesRan <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "timesRan")
samplerStates[[9]]$timesAdapted <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "timesAdapted")
samplerStates[[9]]$sumJumps <- valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "sumJumps")

Now, the samplerStates list will also need to be saved.

R’s random number seed

We’ll also save the current value of R’s random number generation seed, so we can pick up the MCMC sampling precisely as it would have continued, had we actually kept running the MCMC chain now.

seed <- .Random.seed

WAIC State Variables

If you’re using WAIC, you’ll also need to save the current values of some WAIC-related quantities that are used in NIMBLE’s online WAIC algorithm.

waicState <- list()
if(useWAIC) {
    waicState$mcmcIter <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "mcmcIter")
    waicState$lppdSumMaxMat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "lppdSumMaxMat")
    waicState$lppdCurSumMat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "lppdCurSumMat")
    waicState$sspWAICmat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "sspWAICmat")
    waicState$meanpWAICmat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "meanpWAICmat")
    waicState$delta1pWAICmat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "delta1pWAICmat")
    waicState$delta2pWAICmat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "delta2pWAICmat")
    waicState$logProbMat <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "logProbMat")
    waicState$finalized <- valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "finalized")
}

Now, save all these to a file.

fileName <- tempfile()
save(samples, modelVariables, samplerStates, waicState, seed, file = fileName)

New R Session: Restoring your MCMC

Easy from here. Say you’ve now restarted R, and restored the samples, modelVariables, and samplerStates objects.

First, recreate your model and MCMC. This code is exactly the same as above:

###########################
### Start new R session ###
###########################

library(nimble)

code <- nimbleCode({
    mu ~ dnorm(0, 0.0001)
    sigma ~ dunif(0, 1000)
    for(i in 1:2) {
        beta[i] ~ dnorm(0, 0.0001)
    }
    for(i in 1:N) {
        epsilon[i] ~ dnorm(mu, sd = sigma)
        log(lambda[i]) <- beta[1] + beta[2] * x[i] + epsilon[i]
        y[i] ~ dpois(lambda[i])
    }
})

set.seed(0)
N <- 5
x <- rnorm(N, 0, 2)
lambda <- exp(4 + .3*x)
y <- rpois(N, lambda)

constants <- list(N = N, x = x)
data <- list(y = y)
inits <- list(mu = 0, sigma = 1, beta = rep(0,2), epsilon = rep(0, N))

Rmodel <- nimbleModel(code, constants, data, inits)

Rmodel$calculate()

## [1] -1724.11

conf <- configureMCMC(Rmodel, enableWAIC = useWAIC)

## ===== Monitors =====
## thin = 1: beta, mu, sigma
## ===== Samplers =====
## RW sampler (8)
##   - sigma
##   - beta[]  (2 elements)
##   - epsilon[]  (5 elements)
## conjugate sampler (1)
##   - mu

conf$removeSamplers("sigma")
conf$addSampler("sigma", "slice")
conf$printSamplers()

## [1] conjugate_dnorm_dnorm_identity sampler: mu
## [2] RW sampler: beta[1]
## [3] RW sampler: beta[2]
## [4] RW sampler: epsilon[1]
## [5] RW sampler: epsilon[2]
## [6] RW sampler: epsilon[3]
## [7] RW sampler: epsilon[4]
## [8] RW sampler: epsilon[5]
## [9] slice sampler: sigma

Rmcmc <- buildMCMC(conf)

Cmodel <- compileNimble(Rmodel)
Cmcmc <- compileNimble(Rmcmc, project = Rmodel)

Then load our saved file:

##fileName <- '~/temp/restartMCMC.RData'
load(fileName)

Restore the model variables

Restore the variable values into the Cmodel object, from our saved modelVariables list.

The syntax is just the reverse of extracting them earlier.

Cmodel$mu <- modelVariables$mu
Cmodel$sigma <- modelVariables$sigma
Cmodel$beta <- modelVariables$beta
Cmodel$epsilon <- modelVariables$epsilon

The following code block is new. A few things were overlooked before, since the reset = FALSE option was designed to continue running an MCMC that had already been run. The specific issues are pretty technical, but have to do with the mvSaved and mvSamples data structures internal to the Cmcmc algorithm, and initializing those datastructures correctly.

These lines take care of it:

Cmodel$calculate()

## [1] -34.02553

for(name in names(Cmcmc$mvSaved$sizes)) {
    Cmcmc$mvSaved[[name]] <- Cmodel[[name]]
}

resize(Cmcmc$mvSamples, 0)

Restore the sampler state variables

Now extract the values from the samplerStates list, and put them into the Cmcmc sampler functions.

Here, the syntax is slighly different:

for(i in 2:8) {
    valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "scale", samplerStates[[i]]$scale)
    valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesRan", samplerStates[[i]]$timesRan)
    valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesAccepted", samplerStates[[i]]$timesAccepted)
    valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "timesAdapted", samplerStates[[i]]$timesAdapted)
    valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[i]], "gamma1", samplerStates[[i]]$gamma1)
}

valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "width", samplerStates[[9]]$width)
valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "timesRan", samplerStates[[9]]$timesRan)
valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "timesAdapted", samplerStates[[9]]$timesAdapted)
valueInCompiledNimbleFunction(Cmcmc$samplerFunctions[[9]], "sumJumps", samplerStates[[9]]$sumJumps)

This has now restored the “state” into the MCMC sampler algorithms.

Restore the sampler state variables

And, we’ll restore R’s random number generation seed to where it left off.

.Random.seed <- seed

Restore WAIC state variables

If using WAIC, restore the WAIC state variables.

if(useWAIC) {
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "mcmcIter", waicState[["mcmcIter"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "lppdSumMaxMat", waicState[["lppdSumMaxMat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "lppdCurSumMat", waicState[["lppdCurSumMat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "sspWAICmat", waicState[["sspWAICmat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "meanpWAICmat", waicState[["meanpWAICmat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "delta1pWAICmat", waicState[["delta1pWAICmat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "delta2pWAICmat", waicState[["delta2pWAICmat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "logProbMat", waicState[["logProbMat"]])
    valueInCompiledNimbleFunction(Cmcmc$waicFun[[1]], "finalized", waicState[["finalized"]])
}

## [1] FALSE

Continue Running MCMC

You can now continue running the Cmcmc algorithm, and it will effectively “pick up where it left off”.

Just make sure to use the reset = FALSE argument, so it doesn’t reset all the samplers back to their initial states.

And if using WAIC, similarly set resetWAIC = FALSE.

Cmcmc$run(10000, reset = FALSE, resetWAIC = FALSE)

## |-------------|-------------|-------------|-------------|
## |-------------------------------------------------------|

## NULL

samples_continued <- as.matrix(Cmcmc$mvSamples)

Let’s look at the samples from the first run, and from the second.

samples_all <- rbind(samples, samples_continued)
dim(samples_all)

## [1] 20000     4

samples_all[9995:10005,]

##        beta[1]   beta[2]       mu      sigma
##  [1,] 1.394315 0.2578996 2.709307 0.03237251
##  [2,] 1.399039 0.2502140 2.724426 0.05516978
##  [3,] 1.400995 0.2494896 2.728632 0.03437194
##  [4,] 1.400995 0.2494896 2.719170 0.03925602
##  [5,] 1.381995 0.2288096 2.741263 0.05120657
##  [6,] 1.381995 0.2780219 2.716613 0.05292085
##  [7,] 1.381995 0.2220797 2.700665 0.04997873
##  [8,] 1.381995 0.2433144 2.690273 0.08012158
##  [9,] 1.381995 0.2410102 2.669353 0.07388723
## [10,] 1.381995 0.2388559 2.696224 0.04560240
## [11,] 1.331974 0.2388559 2.694611 0.03243129

If using WAIC, we can look at the WAIC value from the combined runs.

Cmcmc$getWAIC()

##   [Warning] There are 5 individual pWAIC values that are greater than 0.4. This may indicate that the WAIC estimate is unstable (Vehtari et al., 2017), at least in cases without grouping of data nodes or multivariate data nodes.

## nimbleList object of type waicNimbleList
## Field "WAIC":
## [1] 190.877
## Field "lppd":
## [1] -17.49842
## Field "pWAIC":
## [1] 77.94008

Verify this is the same as one long run

We’ll double-check these samples are precisely the same, as though we’d done one single long (un-interupted) run of the MCMC.

The code below is entirely copied from code above.

###########################
### Start new R session ###
###########################

library(nimble)

code <- nimbleCode({
    mu ~ dnorm(0, 0.0001)
    sigma ~ dunif(0, 1000)
    for(i in 1:2) {
        beta[i] ~ dnorm(0, 0.0001)
    }
    for(i in 1:N) {
        epsilon[i] ~ dnorm(mu, sd = sigma)
        log(lambda[i]) <- beta[1] + beta[2] * x[i] + epsilon[i]
        y[i] ~ dpois(lambda[i])
    }
})

set.seed(0)
N <- 5
x <- rnorm(N, 0, 2)
lambda <- exp(4 + .3*x)
y <- rpois(N, lambda)

constants <- list(N = N, x = x)
data <- list(y = y)
inits <- list(mu = 0, sigma = 1, beta = rep(0,2), epsilon = rep(0, N))

Rmodel <- nimbleModel(code, constants, data, inits)

Rmodel$calculate()

## [1] -1724.11

conf <- configureMCMC(Rmodel, enableWAIC = useWAIC)

## ===== Monitors =====
## thin = 1: beta, mu, sigma
## ===== Samplers =====
## RW sampler (8)
##   - sigma
##   - beta[]  (2 elements)
##   - epsilon[]  (5 elements)
## conjugate sampler (1)
##   - mu

conf$removeSamplers("sigma")
conf$addSampler("sigma", "slice")
conf$printSamplers()

## [1] conjugate_dnorm_dnorm_identity sampler: mu
## [2] RW sampler: beta[1]
## [3] RW sampler: beta[2]
## [4] RW sampler: epsilon[1]
## [5] RW sampler: epsilon[2]
## [6] RW sampler: epsilon[3]
## [7] RW sampler: epsilon[4]
## [8] RW sampler: epsilon[5]
## [9] slice sampler: sigma

Rmcmc <- buildMCMC(conf)

Cmodel <- compileNimble(Rmodel)
Cmcmc <- compileNimble(Rmcmc, project = Rmodel)

set.seed(0)
Cmcmc$run(20000)

## |-------------|-------------|-------------|-------------|
## |-------------------------------------------------------|

## NULL

samples_all <- as.matrix(Cmcmc$mvSamples)
dim(samples_all)

## [1] 20000     4

samples_all[9995:10005,]

##        beta[1]   beta[2]       mu      sigma
##  [1,] 1.394315 0.2578996 2.709307 0.03237251
##  [2,] 1.399039 0.2502140 2.724426 0.05516978
##  [3,] 1.400995 0.2494896 2.728632 0.03437194
##  [4,] 1.400995 0.2494896 2.719170 0.03925602
##  [5,] 1.381995 0.2288096 2.741263 0.05120657
##  [6,] 1.381995 0.2780219 2.716613 0.05292085
##  [7,] 1.381995 0.2220797 2.700665 0.04997873
##  [8,] 1.381995 0.2433144 2.690273 0.08012158
##  [9,] 1.381995 0.2410102 2.669353 0.07388723
## [10,] 1.381995 0.2388559 2.696224 0.04560240
## [11,] 1.331974 0.2388559 2.694611 0.03243129

if(useWAIC)
    Cmcmc$getWAIC()

##   [Warning] There are 5 individual pWAIC values that are greater than 0.4. This may indicate that the WAIC estimate is unstable (Vehtari et al., 2017), at least in cases without grouping of data nodes or multivariate data nodes.

## nimbleList object of type waicNimbleList
## Field "WAIC":
## [1] 190.8474
## Field "lppd":
## [1] -17.50426
## Field "pWAIC":
## [1] 77.91944

The results look good.

Let me know if anything here is unclear.

Good luck!

Daniel

Restarting NIMBLE MCMC

Daniel Turek

23 January 2018

Creating Model and MCMC

What to Save

Samples

Model Variables

MCMC Sampler State Variables

R’s random number seed

WAIC State Variables

New R Session: Restoring your MCMC

Restore the model variables

Restore the sampler state variables

Restore the sampler state variables

Restore WAIC state variables

Continue Running MCMC

Verify this is the same as one long run