Adjusted MCLMC (#675)

* TESTS

* TESTS

* UPDATE DOCSTRING

* ADD STREAMING VERSION

* ADD PRECONDITIONING TO MCLMC

* ADD PRECONDITIONING TO TUNING FOR MCLMC

* UPDATE GITIGNORE

* UPDATE GITIGNORE

* UPDATE TESTS

* UPDATE TESTS

* ADD DOCSTRING

* ADD TEST

* STREAMING AVERAGE

* ADD TEST

* REFACTOR RUN_INFERENCE_ALGORITHM

* UPDATE DOCSTRING

* Precommit

* CLEAN TESTS

* FIX BAD MERGE

* ADJUSTED MCLMC

* REMOVE BENCHMARKS:

* ADD ADJUSTED MCLMC

* GITIGNORE

* PRECOMMIT CLEAN UP

* FIX SPELLING, ADD OMELYAN, EXPORT COEFFICIENTS

* TEMPORARILY ADD BENCHMARKS

* ADD ADJUSTED MCLMC TUNING

* CLEAN

* UNIFY ADJUSTED MCLMC AND MCHMC

* ADD INITIAL_POSITION

* FIX TEST

* CLEAN UP

* REMOVE BENCHMARKS

* ADD TEST

* REMOVE BENCHMARKS

* MODIFY WINDOW ADAPTATION TO TAKE INTEGRATOR

* MODIFY WINDOW ADAPTATION TO TAKE INTEGRATOR

* BUG FIX

* CHANGE PRECISION

* CHANGE PRECISION

* ADD OMELYAN TEST

* ADD ADJUSTED MCLMC TEST

* ADD ADJUSTED MCLMC TEST

* RENAME O

* UPDATE STREAMING AVG

* UPDATE STREAMING AVG

* FIX MERGE

* UPDATE PR

* RENAME STD_MAT

* RENAME STD_MAT

* RENAME STD_MAT

* MERGE MAIN

* REMOVE COEFFICIENT EXPORTS

* REMOVE COEFFICIENT EXPORTS

* RESOLVE MYPY ISSUE

* RESOLVE MYPY ISSUE

* RETURN EXPECTATION HISTORY

* FIX KWARG BUG

* FIX KWARG BUG

* FIX KWARG BUG IN ADJUSTED MCLMC

* MAKE WINDOW ADAPTATION TAKE INTEGRATOR AS ARGUMENT

* L_proposal_factor

* SPLIT TUNING FOR AMCLMC INTO SEPARATE FILE

* SPLIT TUNING FOR AMCLMC INTO SEPARATE FILE

* RENAME STREAMING_AVERAGE_UPDATE ARGS IN ADJUSTED MCLMC ADAPTATION

* diagnostics

* fix bugs

* FIX MINOR TUNING BUGS

* UPDATE TUNING

* UPDATE TUNING

* UPDATE TUNING

* names

* test

* tuning

* update

* ready for test

* ready for test

* ready for test

* Update blackjax/adaptation/adjusted_mclmc_adaptation.py

Co-authored-by: Junpeng Lao <[email protected]>

* edit

---------

Co-authored-by: Junpeng Lao <[email protected]>

blackjax/__init__.py

@@ -3,6 +3,7 @@
 
 from blackjax._version import __version__
 
+from .adaptation.adjusted_mclmc_adaptation import adjusted_mclmc_find_L_and_step_size
 from .adaptation.chees_adaptation import chees_adaptation
 from .adaptation.mclmc_adaptation import mclmc_find_L_and_step_size
 from .adaptation.meads_adaptation import meads_adaptation
@@ -11,6 +12,7 @@
 from .base import SamplingAlgorithm, VIAlgorithm
 from .diagnostics import effective_sample_size as ess
 from .diagnostics import potential_scale_reduction as rhat
+from .mcmc import adjusted_mclmc as _adjusted_mclmc
 from .mcmc import barker
 from .mcmc import dynamic_hmc as _dynamic_hmc
 from .mcmc import elliptical_slice as _elliptical_slice
@@ -110,6 +112,7 @@ def generate_top_level_api_from(module):
 additive_step_random_walk.register_factory("normal_random_walk", normal_random_walk)
 
 mclmc = generate_top_level_api_from(_mclmc)
+adjusted_mclmc = generate_top_level_api_from(_adjusted_mclmc)
 elliptical_slice = generate_top_level_api_from(_elliptical_slice)
 ghmc = generate_top_level_api_from(_ghmc)
 barker_proposal = generate_top_level_api_from(barker)
@@ -160,6 +163,7 @@ def generate_top_level_api_from(module):
     "chees_adaptation",
     "pathfinder_adaptation",
     "mclmc_find_L_and_step_size",  # mclmc adaptation
+    "adjusted_mclmc_find_L_and_step_size",  # adjusted mclmc adaptation
     "ess",  # diagnostics
     "rhat",
 ]