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

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

[Daniel Turek]

Aaron, it's great to see that you're using both nimbleHMC, and also the new system for derived quantities.

You're very close with what you were trying.  It's actually a little simpler than the approach you tried, but requires a bit of an understanding about how different data structures work in nimble - specifically a "nimbleList" object.

You don't need either the "getEnergy_virtual" or the "getEnergyNF" nimbleFunctions, although those were nicely implemented.  Instead, in the setup code of the derived quantity definition ("derived_Energy"), I extracted the state_current variable from the NUTS HMC sampler.  This internal "state_current" variable is a nimbleList object, which is why I called the local version of it "state_current_NL".  Then, in the run function of derived_Energy, the current value of the Hamiltonian (H) is directly extracted from the state_current variable using "state_current_NL$H".  Easy peasy!

In the code below, I commented out some of the original code which is no longer necessary, and marked the new lines with "NEW".

Let me know how this works for you, and great job again!

Daniel

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)
        state_current_NL <- mcmc$samplerFunctions$contentsList[[1]]$state_current    ## NEW

    },
    run = function(timesRan = double()){

        ##### results[timesRan, 1] <<- getCurEnergy$run()
        results[timesRan, 1] <<- state_current_NL$H                                  ## NEW

    },
    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)

pppp

Rmcmc <- buildMCMC(conf) # this works

##### getEnergyNF(Rmcmc)$run() # no longer necessary

Cmodel <- compileNimble(Rmodel) # model compiles

Cmcmc <- compileNimble(Rmcmc, showCompilerOutput = TRUE) # works now

--
You received this message because you are subscribed to the Google Groups "nimble-users" group.
To unsubscribe from this group and stop receiving emails from it, send an email to nimble-users...@googlegroups.com.
To view this discussion visit https://groups.google.com/d/msgid/nimble-users/710ccdb8-fa2a-4bb2-ab86-222e13f27f34n%40googlegroups.com.