NUTS misses U-turns, runs in circles until max_treedepth

@nhuurre, this is the NUTS implementation that wraps around Stan models so you can experiment with them in R. You think this might be useful to you? If so give it a whirl and see if it works okay. I’m tagging @arya (he started this code) so he can see how it goes.

It’s broken up into two packages. First is this thing, hamwrapper, that wraps up Stan gradients. Install that with:

devtools::install_github("bbbales2/hamwrapper")

This could be done with just the gradients exposed by Rstan, but @arya started this working on a project that needed higher order autodiff which is why the extra code.

The next package is RNUTS. I think it’s better to just check this one out as an RStudio project. The git url is https://github.com/bbbales2/RNUTS.git. So make a project with that, and then control + shift + B to build and install it. The Hadley R book is a quick skim to see how to work with packages: http://r-pkgs.had.co.nz/. This is my first attempt at doing R packages so it’s probly pretty bad. I thought this would be a good way to organize stuff though.

There’s a file in the root of the package (I don’t think files are supposed to be here but whatever), rnuts_example.R, that shows how to use this along with a little bit of what I think might be helpful plotting code.

The main code is in R/nuts.R. Navigate down to “oneSampleNuts” and you can look through how it works. Anyway tear it up and try whatever.

There’s one unit test that makes sure that the algorithm is roughly sampling a single normal random variable correctly and with some tiny bit of efficiency. There are probably bugs here. I didn’t exhaustively test this.

Here are a couple plots produced from the example file:

scores

Here’s another. Check the color scale here. It’s kinda weird, but I wanted to be able to differentiate between something be close to uturning but not uturning and something that just barely uturned (so there’s kinda of a discontinuity in the color gradient near zero, and < 0.0 is the uturn condition).

checks

Hopefully that’s useful

3 Likes