search

CDCgov / Rt-without-renewal

https://cdcgov.github.io/Rt-without-renewal/
Apache License 2.0
19 stars 3 forks source link

Add metric from `pathfinder` to pre-heat NUTS #238

Open SamuelBrand1 opened 5 months ago

SamuelBrand1 commented 5 months ago

We are by default using pathfinder.jl to pre-heat our NUTS sampling by providing an (approximate) sample from the posterior distribution as the initial point for NUTS warm-up, as per here.

The reason we didn't also include an estimate of the underlying kinetic energy matrix/inverse curvature metric was because of the API to Turing not being straightforward.

As of https://turinglang.org/docs/tutorials/docs-16-using-turing-external-samplers/ this appears to be fixed/made easy to do, and so we should have metric pre-heating.

SamuelBrand1 commented 5 months ago

Unfortunately, AdvancedHMC.NUTS doesn't have an adtype option because its designed to interface directly with a LogDensityProblemAD.ADgradient object... which means that we can't direct port across.

I've flagged this with Turing https://github.com/TuringLang/Turing.jl/issues/2248

SamuelBrand1 commented 5 months ago

The direction of travel here:

https://github.com/TuringLang/Turing.jl/issues/2248#issuecomment-2148237328

SamuelBrand1 commented 5 months ago

This can go forwards after #262 because the latest version of Turing has a kwarg to externalsampler to pass the adtype choice.

SamuelBrand1 commented 4 months ago

Update: f2f chat with @dylanhmorris and @athowes (don't have his git handle) as part of STF meeting included discussion around when sometimes Pathfinder doesn't improve NUTS pre-heating.

Some experimentation on the associated branch to this issue is having mixed results with including the metric in pre-heat; I'm de-prioritizing investigating this for now but would like to keep this issue open.