initializing NUTS window adapatation with non-default mass matrix

[ismael-mendoza]

Hi, I was wondering if there was a supported way of initializing the window adaptation with the NUTS algorithm using a mass matrix that was not the identity? I know that you can pass in an initial_step_size, but I don't think there is a similar argument for the mass matrix?

[junpenglao]

It is not something easy to do:
You can modify the init in https://github.com/blackjax-devs/blackjax/blob/main/blackjax/adaptation/window_adaptation.py#L105-L125 to take an initial mass matrix. However, note that the initial mass matrix would be used to generate samples for taking sample for the mass matrix adaptation (started from the first slow window), which then get replace with the inverse mass matrix estimated from the first slow window at the end. If this is your intention, then great; but if you want to actually adapt the inverse mass matrix starting from the your matrix input, it is much more complicated because you need to modify the welford_algorithm https://github.com/blackjax-devs/blackjax/blob/main/blackjax/adaptation/mass_matrix.py#L159

[ismael-mendoza]

Thanks for the information!