Saving steps within HMC trajectory - for teaching / visualisation

General
1

Someone has possibly done this before. If so, please let me know! I would like to obtain a matrix of little steps taken along the HMC trajectory over 2 parameters, at least for a few transitions. This is just for teaching, so I can plot them and illustrate the algorithm. I only need the locations half of the symplectic vectors at each step (though maybe gradients would be cool too)

Best case scenario: Someone replies to ask why I don’t use some option already in Stan that I just didn’t know about.

OK case scenario: someone has a clever trick I would never have thought of, or some non-Stan option

Worst, i.e. not worth the effort and/or impenetrable for learners: I have to write some bespoke sampling program to get it.

There are some really nice illustrations out there by the way, e.g. in the new book by Fearnhead, Nemeth et al. But I’d like a way for learners to see it in action on their own posteriors. So to speak.

1 Like
2

Stan doesn’t currently support sub-iteration output like this. You could do something using BridgeStan, but that probably falls into your ‘worst’ category.

Do you know about the MCMC Interactive Gallery? It doesn’t let you plug in arbitrary posteriors, but it’s still a step above static figures

3 Likes
3

Thanks Brian. Yes, the interactive Gallery is very nice indeed. I sometimes show people that. I suppose what I want to particularly add is that they can use their own Stan code with a prior and likelihood and see some trajectories. then we can make it hard for HMC with an inapproriate prior to get a lot of flat plateau, ridge, or tiny variance, whatever, and get some intuition for divergences, step size etc.

bridgeStan is on my to-learn list for sure though!

4

I hadn’t heard of this despite having seen several of the authors on several occasions last summer. I tracked down the reference:

Amazon lists it as due July 31, 2025? I see there’s a July 2024 draft. Is there something more recent?

HMC is trivial. It’s less than half a page of code to write this in Python for an arbitrary Stan model using BridgeStan. I would find this clearer as a learner than if it were implemented inside of Stan through umpteen layers of interfaces and indirection.

1 Like
5

The code for the interactive gallery is quite readable, and I have some good experience by now compiling Stan models for use in browsers, so wouldn’t be that difficult I don’t think… it would still only really work for two dimensional distributions, though.

I’ll add this to my list of things to try on a rainy day sometime

1 Like
6

Agreed – it’s a rainy day kind of task. Please don’t divert your energy to it on my behalf. I don’t imagine I’ll have any rainy days for a good few months, but if I return to it, I’ll post an update here.

7

Not yet. CUP lists it as due in July 2025, after a luxurious copy editing and printing process. However, I think the PDF available online will be pretty close to the final book, from the look of it.

8

Wow, that MCMC Interactive Gallery is impressive! @robertgrant, McElreath showed some code for this in his second edition, chapter 9.

9

With the aid of ChatGPT, this kind of stuff is very easy to do these days. The following took less than 15m and a couple of those were getting the color gradient and writing this reply.

norm2d.stan

parameters {
  vector[2] x;
}
model {
  x ~ normal(0, 1);
}

trajectory.py

import numpy as np
import bridgestan
import pandas as pd
import plotnine as pn

def plot_trajectory(trajectory):
    df = pd.DataFrame({'x1': trajectory[:, 0], 'x2': trajectory[:, 1] })
    df['step'] = df.index
    plot = (
        pn.ggplot(df, pn.aes(x='x1', y='x2', color='step'))
        + pn.geom_path()
        + pn.scale_color_gradient(low='blue', high='red')
        + pn.labs(title="Hamiltonian Trajectory", x="x1", y="x2")
        + pn.theme_minimal()
    )
    return plot

def hamiltonian_trajectory(model, init_pos, init_mom, step_size, mass_matrix, steps):
    dim = init_pos.shape[0]
    trajectory = np.empty((steps + 1, dim))
    trajectory[0] = init_pos.copy()
    q = init_pos.copy()
    p = init_mom.copy()
    M_inv = np.linalg.inv(mass_matrix)

    def grad_U(position):
        _, grad = model.log_density_gradient(position)
        return -grad
    
    for i in range(steps):
        p -= 0.5 * step_size * grad_U(q)
        q += step_size * (M_inv @ p)
        p -= 0.5 * step_size * grad_U(q)
        trajectory[i + 1] = q.copy()

    return trajectory

norm2d = bridgestan.StanModel("norm2d.stan")
t = hamiltonian_trajectory(norm2d, init_pos=np.array([1.0, 0.0]), init_mom=np.array([0.0, 0.2]),
                               step_size=0.5, mass_matrix=np.array([[1.0, 0.0], [0.0, 1.0]]), steps=20)

for i in range(20):
    print(f"t[{i}] = {t[i]}")

plot = plot_trajectory(t)
plot.show()

You can run it with:

$ python trajectory.py
10

Vibe coding on Stan Discourse (shocked face)

11

[edit: expanded]

You mean there are still programmers who don’t use LLMs?

Is it still vibe coding when I’m fluent in the language and packages used and edit the results to do what I want? I can’t imagine coding without GPT now any more than I could’ve imagined coding without StackOverflow three years ago. But Mitzi says Claude 3.7 Sonnet is better at Stan.

I’m doing things like having ChatGPT o3-mini-high take a Stan data block, a CSV data format and some description and write code to generate the data in JSON format, plot it in various ways, etc. It’s exactly the kind of programming where I can never remember the names of the magic functions and their arguments in the data munging and plotting programs. I know how they work and can read the results, so it’s really time saving. And frankly, o3-mini-high is a much more accurate programmer than me if you can clearly articulate what you want in small pieces.

Here’s an example that took a second prompt to debug and then produced working code that I verified at the output level: https://chatgpt.com/share/67e6c867-67f0-8011-8c2e-825e79784366