Thanks for putting this online, it has been helpful while I have been trying to write my own implementation of NUTS.
I just wanted to draw your attention to a discrepancy between your R code and algorithm 6 in Hoffman & Gelman.
Line 81 of nuts.R has theta <- temp$theta
, in other words the accepted proposal is immediately stored in the variable theta
. If the loop of lines 66-91 repeats, this means that this new value of theta
will be passed to build_tree
as argument 7, which is called theta0
inside build_tree
.
But in algorithm 6 of Hoffman & Gelman, the equivalent calls to BuildTree are always passed $\theta^{m-1}$ for $\theta^0$, so that within BuildTree, $\theta^0$ will always be equal to the previous iteration's $\theta$ rather than any prospectively accepted proposals for $\theta$.
I'm not sure which is correct, the paper or your implementation. In any case, I think this only affects how the step size is calculated, so it wouldn't lead to incorrect inferences.