Writing user samplers

class: center, middle, inverse, title-slide

.title[
# Writing user samplers
]
.subtitle[
## NIMBLE NESS 2022 short course
]
.author[
### NIMBLE Development Team
]
.date[
### May 2022
]

---

### Introduction

- NIMBLE allows you to write your own sampler and use it in combination with NIMBLE's built-in samplers.

- If you've developed a new MCMC sampling algorithm or want to try out an algorithm from the literature, you may want to combine it with standard samplers.

- Or for someone with a specific model, you might want to do some very specific sampling approach for one or a few parameters in your model and rely on NIMBLE's built-in samplers for the rest of the model.

- Here we'll see an illustration of how to do this.

---
### The reflection sampler

Suppose you have a parameter with a finite domain, in particular a fixed lower bound, such as a gamma distribution, a uniform distribution, or a lognormal distribution.

A standard Metropolis sampler could propose a value that is below the lower bound. This would give a probability density for the proposed value that is zero so the proposal would be rejected. That's fine, but it wastes the computation involved in proposing the value and determining that it should be rejected. If the current value of the parameter under consideration is near the bound, this will happen nearly 50% of the time.

Instead, we can use **reflection**. If the proposed `$\theta^\prime < b$` where `$b$` is the bound, simply set `$\theta^\prime$` to `$b + (b-\theta^\prime)$`

---

---

---
### Writing a nimbleFunction for the reflection sampler

The *run* function for the reflection sampler needs to check the proposed value against the distribution bounds and modify the proposal as needed.

However, we first need to modify the *setup* function to check if the distribution has finite lower or upper bounds and only consider scalar parameters, thereby avoiding some computation at run-time.

.scroll-box-16[

```r
target = "tau"

RW_reflect <- nimbleFunction(
    contains = sampler_BASE,
    setup = function(model, mvSaved, target, control) {
        
        ## Some checking...
        targetComponents <- model$expandNodeNames(target, returnScalarComponents = TRUE)
        if(length(targetComponents) > 1)
            stop("RW_reflect: cannot use univariate RW sampler on multivariate target, try RW_block sampler.")

## Is the reflection sampler appropriate?
        dist <- model$getDistribution(target)
        rg <- getDistributionInfo(dist)$range
        if(rg[[1]] > -Inf || rg[[2]] < Inf)
            reflect <- TRUE else reflect <- FALSE

calcNodes  <- model$getDependencies(target)
    },
    run = function() {
        propValue <- rnorm(1, mean = model[[target]], sd = scale)

if(reflect) {
            lower <- model$getBound(target, 'lower')
            upper <- model$getBound(target, 'upper')
            if(propValue < lower) propValue <- 2*lower - propValue
            if(propValue > upper) propValue <- 2*upper - propValue
        }

model[[target]] <<- propValue
        ## Calculate difference of log posterior relative to previous value (which has been saved).
        logMHR <- calculateDiff(model, calcNodes) 
        jump <- decide(logMHR)
        if(jump)
            nimCopy(from = model, to = mvSaved, row = 1, nodes = calcNodes, 
                    logProb = TRUE)
        else
            nimCopy(from = mvSaved, to = model, row = 1, nodes = calcNodes, 
                    logProb = TRUE)
    },
    methods = list(
        reset = function () {}
    )
)
```
]

---

# NIMBLE's Metropolis sampler

- Often it's easiest when writing a new sampler that is similar to an existing NIMBLE sampler to just modify the code for the existing sampler.

- In this case, NIMBLE's existing random walk sampler has some nice additional functionality that we can include in our reflection sampler, specifically the ability to adapt the proposal variance.

- You can see all of our sampler nimbleFunctions in the file *R/MCMC_samplers.R* in the [NIMBLE source package](https://cran.r-project.org/src/contrib/nimble_0.12.1.tar.gz), also available as [MCMC_samplers.R](MCMC_samplers.R) in this repository.

- In the next slide there is the full new reflection sampler, building on NIMBLE's baseline random walk sampler to include adaptation.

---
.scroll-box-20[

```r
RW_reflect <- nimbleFunction(
    contains = sampler_BASE,
    setup = function(model, mvSaved, target, control) {
        ## control list extraction
        logScale      <- if(!is.null(control$log))
                      control$log else FALSE
        reflective    <- if(!is.null(control$reflective))
                      control$reflective else FALSE
        adaptive      <- if(!is.null(control$adaptive))
                      control$adaptive else TRUE
        adaptInterval <- if(!is.null(control$adaptInterval))
                      control$adaptInterval else 200
        scale         <- if(!is.null(control$scale))
                      control$scale else 1
        ###  node list generation  ###
        targetAsScalar <- model$expandNodeNames(target, 
                       returnScalarComponents = TRUE)
        if(length(targetAsScalar) > 1)     
                       stop('more than one target; cannot use RW_reflect sampler, try RW_block sampler')
        if(model$isDiscrete(target))
                       stop('cannot use RW_reflect sampler on discrete-valued target; try slice sampler')

### ADDED code ############################################
        dist <- model$getDistribution(target)
        rg <- getDistributionInfo(dist)$range
        if(rg[[1]] > -Inf || rg[[2]] < Inf)
                  reflect <- TRUE else reflect <- FALSE
        ###########################################################

calcNodes  <- model$getDependencies(target)
        ###  numeric value generation  ###
        scaleOriginal <- scale
        timesRan      <- 0
        timesAccepted <- 0
        timesAdapted  <- 0
        scaleHistory          <- c(0, 0)
        acceptanceRateHistory <- c(0, 0)
        optimalAR <- 0.44
        gamma1    <- 0
    },
    
    run = function() {
        propValue <- rnorm(1, mean = model[[target]], sd = scale)

### ADDED code ############################################
        if(reflect) {
            lower <- model$getBound(target, 'lower')
            upper <- model$getBound(target, 'upper')
             if(propValue < lower) propValue <- 2*lower - propValue
             if(propValue > upper) propValue <- 2*upper - propValue
        }
        ###########################################################

model[[target]] <<- propValue
        logMHR <- calculateDiff(model, calcNodes)
        jump <- decide(logMHR)
        if(jump)
            nimCopy(from = model, to = mvSaved, row = 1, nodes = calcNodes, 
                         logProb = TRUE)
        else
            nimCopy(from = mvSaved, to = model, row = 1, nodes = calcNodes, 
                         logProb = TRUE)
        if(adaptive)     adaptiveProcedure(jump)
    },
    
    methods = list(
        adaptiveProcedure = function(jump = logical()) {
            timesRan <<- timesRan + 1
            if(jump)     timesAccepted <<- timesAccepted + 1
            if(timesRan %% adaptInterval == 0) {
                acceptanceRate <- timesAccepted / timesRan
                timesAdapted <<- timesAdapted + 1
                setSize(scaleHistory,          timesAdapted)
                setSize(acceptanceRateHistory, timesAdapted)
                scaleHistory[timesAdapted] <<- scale
                acceptanceRateHistory[timesAdapted] <<- acceptanceRate
                gamma1 <<- 1/((timesAdapted + 3)^0.8)
                gamma2 <- 10 * gamma1
                adaptFactor <- exp(gamma2 * (acceptanceRate - optimalAR))
                scale <<- scale * adaptFactor
                timesRan <<- 0
                timesAccepted <<- 0
            }
        },
        
        reset = function() {
            scale <<- scaleOriginal
            timesRan      <<- 0
            timesAccepted <<- 0
            timesAdapted  <<- 0
            scaleHistory          <<- scaleHistory          * 0
            acceptanceRateHistory <<- acceptanceRateHistory * 0
            gamma1 <<- 0
        }
    )
)
```
]

---

### Using the new sampler: setup

Using the sampler is simple. Just modify the default MCMC configuration for a model to use the new sampler on a node of interest.

Let's try this with a simulated example where the true value of a variance component, named tau, is zero. We'll set up the model so that the default sampler for `tau` is a RW sampler so that the comparison with the reflection version is reasonable.

```r
set.seed(1)
m  <- nsub <- 10
id <- rep(1:m, each = nsub)
n <- length(id)
x <- runif(n)
x  <- x - mean(x)
b0 <- 2
b1 <- 3
y <- b0 + b1*x + rnorm(n)
y <- matrix(y, nrow = m)
x <- matrix(x, nrow = m)
```

---

And here's the model, with a default MCMC.

```r
code  <- nimbleCode({
    for(j in 1:m) {
        for(i in 1:nsub) {
            y[j,i] ~ dnorm(b0[j] + b1*x[j,i], sd = sigma)
        }
        b0[j] ~ dnorm(mu, sd = tau)
    }
    
    sigma ~ dunif(0, 20) 
    tau ~ dunif(0, 20)
    b1 ~ dflat()
    mu ~ dflat()
})
    
model <- nimbleModel(code, data = list(y = y),
                  constants = list(x = x, m = m, nsub = nsub),
                  inits = list(mu = 0, sigma = 0.5, tau = 0.5, b1 = 0))
cModel  <- compileNimble(model)
conf <- configureMCMC(cModel)
```

```
## ===== Monitors =====
## thin = 1: b1, mu, sigma, tau
## ===== Samplers =====
## RW sampler (2)
##   - sigma
##   - tau
## conjugate sampler (12)
##   - b1
##   - mu
##   - b0[]  (10 elements)
```

```r
MCMC  <- buildMCMC(conf)
cMCMC  <- compileNimble(MCMC)
smp1  <- runMCMC(cMCMC)
```

```
## |-------------|-------------|-------------|-------------|
## |-------------------------------------------------------|
```

```r
par(mfrow = c(1,2))
ts.plot(smp1[ , 'tau'], ylab = 'tau', xlab = 'iteration')
ts.plot(smp1[ , 'mu'], ylab = 'mu', xlab = 'iteration')
```

![](writing_user_sampler_files/figure-html/mcmc-1.png)

---

### Using the new sampler: modifying the MCMC configuration

Now we'll try the reflection sampler instead. We can just assign it as a sampler
in the way we've seen before.

```r
conf$removeSamplers('tau')
conf$addSampler('tau', type = 'RW_reflect')

MCMC2 <- buildMCMC(conf)
cMCMC2 <- compileNimble(MCMC2, resetFunctions = TRUE)
```

```
## Compiling
##   [Note] This may take a minute.
##   [Note] Use 'showCompilerOutput = TRUE' to see C++ compilation details.
```

```r
nimCopy(model, cModel)  # initialize 2nd MCMC same as 1st MCMC

smp2 <- runMCMC(cMCMC2)
```

```
## Running chain 1 ...
```

```
## |-------------|-------------|-------------|-------------|
## |-------------------------------------------------------|
```
---

![](writing_user_sampler_files/figure-html/unnamed-chunk-1-1.png)

So we see that the sampler mixes rather better, able to escape from near zero more quickly. 
Mixing is still not great, because of the posterior dependence between `tau` and the `b0` random effects. 
---
### Required arguments and methods for a sampler `nimbleFunction`

- It must contain the argument `contains = sampler_BASE`.
- (This is a simple class inheritance system that mimic's R's *contains* argument.)

- The `setup` function must take the arguments `model`, `mvSaved`, `target`, and `control`.

- `model` is the model being sampled.
  - `mvSaved` is a length-one `modelValues` object that keeps an up-to-date copy of all model values, including log probabilities.
  - `target` is a vector of node names to be sampled.
  - `control` is a list that can contain whatever elements you want.

- The `run` function (method) must execute the sampler.

- The `reset` method must reset the sampler if that means anything for the particular sampler.
  - Example: An adaptive sampler would reset its proposal distribution.

---
### Required behavior of a sampler:

- Upon entry, the sampler can assume that `mvSaved[[1]]` contains a complete copy of the model's variables, including logProb variables.

- The sampler may do whatever it wants (assuming it is valid for MCMC) in its `run` function, including modifying values of model variables, including logProb variables.

- Upon exiting the `run` function, `mvSaved[[1]]` must again contain a complete copy of the model's variables, including logProb variables.

- The `mvSaved[[1]]` is like the "current" state of the model.
  - The `run` function puts proposed values in the model and does appropriate calculations.
  - If the proposal is rejected: copy from `mvSaved[[1]]` to the model.
  - If the proposal is accepted: copy from the model to `mvSaved[[1]]`.