MCLMC adaptation num_steps and initial guess

[hsimonfroy]

Hello,
2 small things:

adaptation num_steps

mclmc_find_L_and_step_size now also returns an estimate of the number of steps performed, but does not take into account roundings and +1 (negligeables), and especially stepsize readjustment when diagonally preconditioning:

                # readjust the stepsize
                steps = num_steps2 // 3  # we do some small number of steps

adaptation initial guess

It may be worth to allow initial guess proposition for adapted parameters. Currently, they are initialized at

    params = MCLMCAdaptationState(sqrt(dim), jnp.sqrt(dim) * 0.25, inverse_mass_matrix=jnp.ones((dim,))

so if the optimal values are far from it, convergence can be slower than if starting at a better guess. This can simply be an optional MCLMCAdaptationState or dict replacing the initialization.

[reubenharry]

Hi,

Thanks for opening this issue!

Adaptation num_steps

Yes, the stepsize readjustment should be taken into account. Good point! Feel free to open a PR. Otherwise at some point I should fix this.

initial guess

Agreed. If you look as adjusted_mclmc_adaptation (the tuning for the MH adjusted version of the algorithm), I do exactly this. but apparently I didn't also add this to mclmc_adaptation. Again, feel free to submit a PR.