Hamiltonian as derived quantity in HMC

[Aaron Renggli]

Hi everyone

I am using HMC with the NUTS sampler. For diagnostics, it would be nice to record some quantities from the sampler, for example the value of the Hamiltonian at each step or if the sampler was divergent (recorded for each step, not just cumulatively at the end via numDivergences).
As I don't really want to touch the sampler itself, I thought the derived quantities framework is perfect for this as it has access to the mcmc object, so I tried to write a nimble function to record the Hamiltonian as a derived quantity (see repex below).
I am able to add it to the mcmc-conf as a derived quantity, and I can build an mcmc from it; but compilation fails with 'Error in getSymObj_recurse(firstArg, symTab, recurse = TRUE) : Problem (ii) working through mcmc$samplerFunctionsFALSE'. Unfortunately, I have no idea what that error means. 

What is going wrong?

Thanks
Aaron

############################################################
library(nimble)
library(nimbleHMC)

code <- nimbleCode({
    b0 ~ dnorm(0, 0.001)
    b1 ~ dnorm(0, 0.001)
    sigma ~ dunif(0, 10000)
    for(i in 1:N) {
        mu[i] <- b0 + b1 * x[i]
        y[i] ~ dnorm(mu[i], sd = sigma)
    }
})

set.seed(0)
N <- 100
x <- rnorm(N)
y <- 1 + 0.3*x + rnorm(N)
constants <- list(N = N, x = x)
data <- list(y = y)
inits <- list(b0 = 1, b1 = 0.1, sigma = 1)

Rmodel <- nimbleModel(code, constants, data, inits, buildDerivs = TRUE)

getEnergy_virtual <- nimbleFunctionVirtual(
  run = function() { returnType(double()) }
)
getEnergyNF <- nimbleFunction(
  contains = getEnergy_virtual,
  setup = function(mcmc) {},
  run = function() {
    returnType(double())
    return(mcmc$samplerFunctions$contentsList[[1]]$state_current$H)
  }
 
)

derived_Energy <- nimbleFunction(
  name = "derived_Energy",
  contains = derived_BASE,
  setup = function(model, mcmc, interval, control){
    names <- c("Energy", "", "")
    results <- array(0, c(1, 1))
    getCurEnergy <- getEnergyNF(mcmc)
  },
  run = function(timesRan = double()){
    results[timesRan, 1] <<- getCurEnergy$run()
  },
  methods = list(
    set_interval = function(newInterval = double()) {
      interval <<- newInterval
    },
    before_chain = function(niter = double(), nburnin = double(), thin = double(1), chain = double()) {
      results <<- nimArray(NA, c(niter, 1))
    },
    after_chain = function() {
    },
    get_results = function() {
      returnType(double(2))
      return(results)
    },
    get_names = function() {
      returnType(character(1))
      return(names)
    },
    reset = function() {
     
    }
  )
 
)

conf <- configureHMC(Rmodel)
conf$addDerivedQuantity('Energy', print = TRUE)

Rmcmc <- buildMCMC(conf) # this works
getEnergyNF(Rmcmc)$run() # this works too

Cmodel <- compileNimble(Rmodel) # model compiles
Cmcmc <- compileNimble(Rmcmc, showCompilerOutput = TRUE) # error

###############################################################