Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Autodiff compatable emulator and Barker gradient-based sampler #339

Open
wants to merge 30 commits into
base: main
Choose a base branch
from

Conversation

odunbar
Copy link
Collaborator

@odunbar odunbar commented Jan 16, 2025

Purpose

Closes #341
Closes #342
Closes #343
Closes #344
Closes #345

Content

From a private repos of @KotaOK-00 - additional changes to interface

  • Copy implementation of the gradient-based sampler proposals
  • Implemented the calculation of ESJD
  • Added esjd unit test.
  • Created a new system AutodiffProtocol containing derived types GradFreeProtocol ForwardDiffProtocol and ReverseDiffProtocol. autodiff options are created with the MCMCProtocol, e.g.,
BarkerSampling() # creates BarkerSampling{ForwardDiffProtocol}()
BarkerSampling{ReverseDiffProtocol}() # creates BarkerSampling{ReverseDiffProtocol}()
pCNMHSampling() # creates pCNMHSampling{GradFreeProtocol}()

within the MCMC propose method direct calls to

ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
Symmetric(ForwardDiff.hessian(x -> AdvancedMH.logdensity(model, x), current_state.params))

Are replaced with autodiff_gradient and autodiff_hessian

autodiff_gradient(model, current_state.params, sampler)
autodiff_hessian(model, current_state.params, sampler)

that dispatch off the pre-specified autodiff protocol in the sampler. This will make adding new autodiff options easier

  • Created a new type AGPJL for AbstractGPs. When building an emulator with AGPJL, As there is no optimizer of such kernels, we allow users to train e.g. with GPJL, then do the following:
opt_params_per_model = Emulators.get_params(gp_optimized)
kernel_params = [
Dict(
       "log_rbf_len" => params[1:end-2],
       "log_std_sqexp" => params[end-1],
       "log_std_noise" => params[end],
) 
for params in opt_params_per_model]
agp = GaussianProcess(AGPJL(), ...)
Emulator( agp, ... ; ..., kernel_params = kernel_params) # builds

If the user does not provide kernel_params, a useful error message appears to direct them what to do.

  • Unit tests in GaussianProcesses/runtests.jl for AGP interface as above, and that GPJL vs AGP gives very similar emulator predicitions
  • Unit tests in MarkovChainMonteCarlo/runtests.jl for AGP on RWM and pCN sampling - give similar posterior means to GPJL (NB it is much slower, 5-10x)
  • Unit testing in MarkovChainMonteCarlo/runtests.jl for Barker update (ForwardDiff and ReverseDiff) with comparable results to RWM.

MISC

added non-Barker implementation as new issue. I have removed these algorithms as they do not perform well in unit testing and require possibly significant improvements to the tuning mechanisms to work


  • I have read and checked the items on the review checklist.

@odunbar odunbar changed the title MCMC updates (gradient-based variants and ESJD) [WIP] MCMC updates (gradient-based variants and ESJD calculation with the chain) Jan 16, 2025
@odunbar odunbar changed the title [WIP] MCMC updates (gradient-based variants and ESJD calculation with the chain) [WIP] MCMC updates (gradient-based variants and ESJD) Jan 16, 2025
Copy link

codecov bot commented Jan 16, 2025

Codecov Report

Attention: Patch coverage is 86.36364% with 15 lines in your changes missing coverage. Please review.

Project coverage is 88.62%. Comparing base (d5a079b) to head (671c5ff).

Files with missing lines Patch % Lines
src/MarkovChainMonteCarlo.jl 81.03% 11 Missing ⚠️
src/GaussianProcess.jl 91.83% 4 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #339      +/-   ##
==========================================
- Coverage   88.82%   88.62%   -0.20%     
==========================================
  Files           7        7              
  Lines        1271     1363      +92     
==========================================
+ Hits         1129     1208      +79     
- Misses        142      155      +13     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@odunbar odunbar changed the title [WIP] MCMC updates (gradient-based variants and ESJD) Add Forward-diff emulator and Barker gradient-based sampler Jan 22, 2025
@odunbar odunbar changed the title Add Forward-diff emulator and Barker gradient-based sampler Add Autodiff compatable emulator and Barker gradient-based sampler Jan 22, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment