R/NUTS.R

In petedodd/MCIR: Monte Carlo Inference Routines: A Collection of Generic Algorithms for Bayesian Inference

## This is an R translation by P.J.Dodd 2015 of the matlab NUTS sampler
## see below for original license


## % Copyright (c) 2011, Matthew D. Hoffman
## % All rights reserved.
## % 
## % Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
## % 
## % Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
## % Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
## % THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

##' Function implementing the No-U-Turn Sampler (NUTS)
##'
##' todo: more detail on arguments and provenance...epsilon
##' @title nuts_da
##' @param f a function returning list with \code{list(logp=logp,grad=grad)}, \code{logp} being the loglikelihood and \code{grad} being its gradient
##' @param M the number of samples to generate
##' @param Madapt the number of steps of burn-in/how long to run the dual
##' @param theta0 is a vector with the desired initial setting of the parameters.
##' @param delta hould be between 0 and 1, and is a target HMC acceptance probability. Defaults to 0.6 if unspecified.
##' @return \code{list(samples=samples,epsilon=epsilon)}, samples being a matrix with columns parameters and rows samples
##' @author Pete Dodd
##' @examples
##' rosenbrock <- function(x){
##'   f <- (1-x[1])^2 + 100*(x[2] - x[1]^2)^2
##'   g <- c(100*2*(x[2] - x[1]^2)*(-2)*x[1] - 2*(1-x[1]),
##'   200*(x[2] - x[1]^2) )
##'   return(list(logp=-f,grad=-g))
##' }
##'
##' run <- nuts_da(rosenbrock,10*1e3,1e3,runif(2))
##' 
##' corplot(run$samples,labels=c('x','y'))
nuts_da <- function(f, M, Madapt, theta0, delta = 0.6){
    ## % function [samples, epsilon] = nuts_da(f, M, Madapt, theta0, delta)
    ## % 
    ## % Implements the No-U-Turn Sampler (NUTS), specifically, algorithm 6
    ## % from the NUTS paper (Hoffman & Gelman, 2011). Runs Madapt steps of
    ## % burn-in, during which it adapts the step size parameter epsilon, then
    ## % starts generating samples to return.
    ## % 
    ## % epsilon is a step size parameter.
    ## % f(theta) should be a function that returns the log probability its
    ## % gradient evaluated at theta. I.e., you should be able to call 
    ## % [logp grad] = f(theta).
    ## % M is the number of samples to generate.
    ## % Madapt is the number of steps of burn-in/how long to run the dual
    ## % averaging algorithm to fit the step size epsilon. 
    ## % theta0 is a 1-by-D vector with the desired initial setting of the parameters.
    ## % delta should be between 0 and 1, and is a target HMC acceptance
    ## % probability. Defaults to 0.6 if unspecified.
    ## % 
    ## % The returned variable "samples" is an M-by-D matrix of post-burn-in samples 
    ## % generated by NUTS.
    ## % The returned variable "epsilon" is the step size that was fit using dual

    D <- length(theta0)
    samples <- matrix(nrow=M+Madapt,ncol=D)

    fval <- f(theta0)
    logp <- fval$logp
    grad <- fval$grad
    samples[1,] <- theta0

    
    ## % Choose a reasonable first epsilon by a simple heuristic.
    epsilon <- find_reasonable_epsilon(theta0, grad, logp, f)

    ## % Parameters to the dual averaging algorithm.
    gamma <- 0.05
    t0 <- 10
    kappa <- 0.75
    mu <- log(10*epsilon)
    ## % Initialize dual averaging algorithm.
    epsilonbar <- 1
    Hbar <- 0

    pb <- txtProgressBar(min=1,max=(M+Madapt),char='.',style=3) #progress bar
    for(m in 2:(M+Madapt)){
        setTxtProgressBar(pb, m)        #progress
    
        ## % Resample momenta.
        r0 <- rnorm(D)
        ## % Joint log-probability of theta and momenta r.
        joint <- logp - 0.5 * sum(r0^2) 
        ## % Resample u ~ uniform([0, exp(joint)]).
        ## % Equivalent to (log(u) - joint) ~ exponential(1).
        logu <- joint - rexp(1)
        ## % Initialize tree.
        thetaminus <- samples[m-1,]
        thetaplus <- samples[m-1,]
        rminus <- r0
        rplus <- r0
        gradminus <- grad
        gradplus <- grad
        ## % Initial height j = 0.
        j <- 0
        ## % If all else fails, the next sample is the previous sample.
        samples[m,] <- samples[m-1,]
        ## % Initially the only valid point is the initial point.
        n <- 1

        ## % Main loop---keep going until the criterion s == 0.
        s <- 1
        while(s == 1){
            ## % Choose a direction. -1=backwards, 1=forwards.
            v <- 2*(runif(1) < 0.5)-1;
            ## % Double the size of the tree.
            if ( v<0 ){                 #v == -1
                ## [thetaminus, rminus, gradminus, ~, ~, ~, thetaprime, gradprime, logpprime, nprime, sprime, alpha, nalpha] = build_tree(thetaminus, rminus, gradminus, logu, v, j, epsilon, f, joint);
                tree <- build_tree(thetaminus, rminus, gradminus, logu, v, j, epsilon, f, joint)
                thetaminus <- tree$thetaminus
                rminus <- tree$rminus
                gradminus <- tree$gradminus
                thetaprime <- tree$thetaprime
                gradprime <- tree$gradprime
                logpprime <- tree$logpprime
                nprime <- tree$nprime
                sprime <- tree$sprime
                alpha <- tree$alpha
                nalpha <- tree$nalpha
            } else {
                ## [~, ~, ~, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alpha, nalpha] = build_tree(thetaplus, rplus, gradplus, logu, v, j, epsilon, f, joint);
                tree <- build_tree(thetaplus, rplus, gradplus, logu, v, j, epsilon, f, joint)
                thetaplus <- tree$thetaplus
                rplus <- tree$rplus
                gradplus <- tree$gradplus
                thetaprime <- tree$thetaprime
                gradprime <- tree$gradprime
                logpprime <- tree$logpprime
                nprime <- tree$nprime
                sprime <- tree$sprime
                alpha <- tree$alpha
                nalpha <- tree$nalpha
            }
            ## % Use Metropolis-Hastings to decide whether or not to move to a
            ## % point from the half-tree we just generated.
            if ((sprime == 1) & (runif(1) < nprime/n)){
                samples[m,] <- thetaprime
                logp <- logpprime
                grad <- gradprime
            }
            ## % Update number of valid points we've seen.
            n <- n + nprime;
            ## % Decide if it's time to stop.
            s <- sprime && stop_criterion(thetaminus, thetaplus, rminus, rplus)
            ## % Increment depth.
            j <- j + 1
        }

        
        ## % Do adaptation of epsilon if we're still doing burn-in.
        eta <- 1 / (m - 1 + t0)
        Hbar <- (1 - eta) * Hbar + eta * (delta - alpha / nalpha)
        if (m <= Madapt) {
            epsilon <- exp(mu - sqrt(m-1)/gamma * Hbar)
            eta <- (m-1)^{-kappa}       
            epsilonbar <- exp((1 - eta) * log(epsilonbar) + eta * log(epsilon))
        } else {
            epsilon <- epsilonbar
        }
    }                                   #end adapt loop
    close(pb)                           #close progress bar
    
    samples <- samples[(Madapt+1):nrow(samples),] #(Madapt+1:end, :);
    ## sprintf('Took %d gradient evaluations.\n', nfevals) #todo:
    return(list(samples=samples,epsilon=epsilon))
}                                       #end of nuts_da

find_reasonable_epsilon <- function(theta0, grad0, logp0, f){
    epsilon <- 1
    r0 <- runif(length(theta0))
    ## % Figure out what direction we should be moving epsilon.
    LFval <- leapfrog(theta0, r0, grad0, epsilon, f )#, nfevals)
    rprime <- LFval$rprime
    logpprime <- LFval$logpprime
    ## nfevals <- LFval$nfevals
    acceptprob <- exp(logpprime - logp0 - 0.5 * (sum(rprime^2) - sum(r0^2))) 
    a <- 2 * (acceptprob > 0.5) - 1
    ## % Keep moving epsilon in that direction until acceptprob crosses 0.5.
    while(acceptprob^a > 2^(-a)){
        epsilon <- epsilon * 2^a
        LFval <- leapfrog(theta0, r0, grad0, epsilon, f)#, nfevals)
        rprime <- LFval$rprime
        logpprime <- LFval$logpprime
        ## nfevals <- LFval$nfevals
        acceptprob <- exp(logpprime - logp0 - 0.5 * (sum(rprime^2) - sum(r0^2))) 
    }
    return(epsilon)    
} 

## nfevals <- 1                            #a counter for recording function evaluations

leapfrog <- function(theta, r, grad, epsilon, f){#, nfevals){
    rprime <- r + 0.5 * epsilon * grad
    thetaprime <- theta + epsilon * rprime
    fval <- f(thetaprime)
    logpprime <- fval$logp
    gradprime <- fval$grad
    rprime <- rprime + 0.5 * epsilon * gradprime
    ## global nfevals                      #todo - this might need changing
    ## nfevals <<- nfevals + 1
    return(list(thetaprime=thetaprime, rprime=rprime, gradprime=gradprime,
                logpprime=logpprime))#, nfevals=nfevals))
}

stop_criterion <- function(thetaminus, thetaplus, rminus, rplus){
    thetavec <- thetaplus - thetaminus
    criterion <- (sum(thetavec * rminus) >= 0) & (sum(thetavec * rplus) >= 0) 
    return(criterion)
} 

## % The main recursion.
## function [thetaminus, rminus, gradminus, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alphaprime, nalphaprime] = ...
##                 build_tree(theta, r, grad, logu, v, j, epsilon, f, joint0)
build_tree <- function(theta, r, grad, logu, v, j, epsilon, f, joint0){
    ## TODO: remove all capital Ls...
    if(j==0){
        ##     % Base case: Take a single leapfrog step in the direction v.
        LFval <- leapfrog(theta, r, grad, v*epsilon, f)#, nfevals)
        thetaprime <- LFval$thetaprime
        rprime <- LFval$rprime
        gradprime <- LFval$gradprime   
        logpprime <- LFval$logpprime
        ## nfevals <- nfevals + LFval$nfevals
        joint <- logpprime - 0.5 * sum(rprime^2)
        ##     % Is the new point in the slice?
        nprime <- logu < joint
        ##     % Is the simulation wildly inaccurate?
        sprime <- (logu - 1000) < joint
        ##     % Set the return values---minus=plus for all things here, since the
        ##     % "tree" is of depth 0.
        thetaminus <- thetaprime       #just making sure these are local
        thetaplus <- thetaprime
        rminus <- rprime
        rplus <- rprime
        gradminus <- gradprime
        gradplus <- gradprime
        ##     % Compute the acceptance probability.
        alphaprime <- min(1, exp(logpprime - 0.5 * sum(rprime^2) - joint0))
        nalphaprime <- 1        
    } else {                            #j!=0
        ##     % Recursion: Implicitly build the height j-1 left and right subtrees.
        ##     [thetaminus, rminus, gradminus, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alphaprime, nalphaprime] = ...
        ##                 build_tree(theta, r, grad, logu, v, j-1, epsilon, f, joint0);
        tree <- build_tree(theta, r, grad, logu, v, j-1, epsilon, f, joint0)
        thetaminus <- tree$thetaminus
        rminus <- tree$rminus
        gradminus <- tree$gradminus
        thetaplus <- tree$thetaplus
        rplus <- tree$rplus
        gradplus <- tree$gradplus
        thetaprime <- tree$thetaprime
        gradprime <- tree$gradprime
        logpprime <- tree$logpprime
        ## nfevals <- LFval$nfevals
        nprime <- tree$nprime
        sprime <- tree$sprime
        alphaprime <- tree$alphaprime
        nalphaprime <- tree$nalphaprime
        ##     % No need to keep going if the stopping criteria were met in the first
        ##     % subtree.
        if(sprime == 1){
            if (v<0){                   #v == -1
                ##             [thetaminus, rminus, gradminus, ~, ~, ~, thetaprime2, gradprime2, logpprime2, nprime2, sprime2, alphaprime2, nalphaprime2] = ...
                ##                 build_tree(thetaminus, rminus, gradminus, logu, v, j-1, epsilon, f, joint0);
                tree <- build_tree(thetaminus, rminus, gradminus, logu, v, j-1, epsilon, f, joint0)
                thetaminus <- tree$thetaminus
                rminus <- tree$rminus
                gradminus <- tree$gradminus
                thetaprime2 <- tree$thetaprime
                gradprime2 <- tree$gradprime
                logpprime2 <- tree$logpprime
                ## nfevals2 <- tree$nfevals
                nprime2 <- tree$nprime
                sprime2 <- tree$sprime
                alphaprime2 <- tree$alphaprime
                nalphaprime2 <- tree$nalphaprime
            } else {
                ##         [~, ~, ~, thetaplus, rplus, gradplus, thetaprime2, gradprime2, logpprime2, nprime2, sprime2, alphaprime2, nalphaprime2] = ...
                ## ##                 build_tree(thetaplus, rplus, gradplus, logu, v, j-1, epsilon, f, joint0);
                tree <- build_tree(thetaplus, rplus, gradplus, logu, v, j-1, epsilon, f, joint0)
                thetaplus <- tree$thetaplus
                rplus <- tree$rplus
                gradplus <- tree$gradplus
                thetaprime2 <- tree$thetaprime
                gradprime2 <- tree$gradprime
                logpprime2 <- tree$logpprime
                ## nfevals2 <- tree$nfevals
                nprime2 <- tree$nprime
                sprime2 <- tree$sprime
                alphaprime2 <- tree$alphaprime
                nalphaprime2 <- tree$nalphaprime
            }                           #end v==-1
            ##         % Choose which subtree to propagate a sample up from.
            if (runif(1) * (nprime + nprime2) < nprime2 ){ #changed away 0/0
               thetaprime <- thetaprime2
               gradprime <- gradprime2
               logpprime <- logpprime2 
            }
            ##         % Update the number of valid points.
            nprime <- nprime + nprime2
            ##         % Update the stopping criterion.
            sprime <- sprime & sprime2 & stop_criterion(thetaminus, thetaplus, rminus, rplus)
            ##         % Update the acceptance probability statistics.
            alphaprime <- alphaprime + alphaprime2
            nalphaprime <- nalphaprime + nalphaprime2
        }                               #end sprime==1
    }                                   #end if j

    return(list(thetaminus=thetaminus, rminus=rminus, gradminus=gradminus,
                thetaplus=thetaplus, rplus=rplus, gradplus=gradplus,
                thetaprime=thetaprime, gradprime=gradprime,
                logpprime=logpprime, nprime=nprime, sprime=sprime,
                alphaprime=alphaprime, nalphaprime=nalphaprime))
}