Add ChEES-HMC adaptation

[sethaxen]

Hoffman et al introduced ChEES-HMC in http://proceedings.mlr.press/v130/hoffman21a/hoffman21a.pdf. It's one answer to the problem of the wastefulness of NUTS when sampling many chains discussed in #228 (comment). When sampling many chains on the GPU with NUTS, the number of steps taken for all chains is equal to the maximum number of steps taken on any one chain before termination (even though for the chains that already terminated, those steps are never used).

ChEES-HMC uses static HMC to sample all chains, where the trajectory length is jittered. The main contribution of the paper is an adaptation routine to tune the step size and trajectory length to try to minimize the autocorrelation of the chains. For the models they checked, they saw higher ESS/evaluation for many chains than with NUTS.

See in particular Algorithm 1.

[sethaxen]

Perhaps the only potentially complicated part of this is that the transition has a jittered integration time, where the jitter could be with an RNG or a pre-determined Halton sequence (they recommend Halton). So each transition call begins by jittering the integration time (even if warm-up has already finished). That is, the termination criterion has the same relationship to FixedIntegrationTime that JitteredLeapfrog has to Leapfrog. It wasn't obvious to me after a few minutes of thought a good way to do this with the current API.

[xukai92]

I guess in principle we can implement extend the jitter operation to support different noise and implement a specific RNG for Halton sequence. But I feel it might be simpler and neater that we just implement a new termination type for ChEES.

@torfjelde told me that he's going to implement this in AHMC. Maybe he has some thoughts on this as well.

If you need a low-discrepancy sequence, you can try GoldenSequences.jl. Generalized golden sequences are easy to calculate (just add φ each time) and get you very good dispersion. Halton sequences also require fine-tuning of parameters. See here for more.