Hi everyone. I wanted to understand the NUTS algorithm better, so I implemented the naive version from the Hoffman 2014 paper from scratch using a linear regression on the mtcars dataset as a toy example. I made a little 3D animation of the results which I thought I would share here.

The (messy) code is here:

I’ve seen the devs mention that what Stan does is more sophisticated than what is presented in that paper, but this was mostly just for the sake of self-education. There’s a lot of other ways the animation could be improved but anyway yeah enjoy.

edit: And here’s a second one using a spherical density that is Gamma(60,12) along any radius. It leads to some pleasingly wiggly iterations: