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

Introduction of SwapSampler + make TemperedSampler a fancy version of CompositionSampler #152

Merged
merged 90 commits into from
Mar 11, 2023

Conversation

torfjelde
Copy link
Member

This PR is quite a mouthful and I'm still not quite done with it (one of the tests are failing we need more tests), but the gist is as follows:

  1. TemperedSampler now becomes a simple (on-the-fly) wrapper of a MultiSampler and SwapSampler.
  2. Swapping is separated into a SwapSampler which acts on MultiModel. This means that we can now also apply swaps to basically an arbitrary collection of models; it no longer has any ties to the "tempering". This can be done by simply composing a SwapSampler with a MultiSampler, e.g. (sampler × sampler) ∘ swapsampler (and, of course, providing a model that is MultiModel).

I will add more description and tests soon.

@torfjelde
Copy link
Member Author

I have some more changes coming here in the morrow 👍

src/MCMCTempering.jl Outdated Show resolved Hide resolved
src/ladders.jl Show resolved Hide resolved
src/adaptation.jl Show resolved Hide resolved
src/sampler.jl Show resolved Hide resolved
src/sampler.jl Show resolved Hide resolved
src/state.jl Outdated Show resolved Hide resolved
src/stepping.jl Outdated Show resolved Hide resolved
src/swapsampler.jl Show resolved Hide resolved
test/runtests.jl Outdated Show resolved Hide resolved
test/runtests.jl Show resolved Hide resolved
@torfjelde
Copy link
Member Author

Ah shit, I think I destroyed the commit history in an attempt to rebase here 😕 It looks like the diff is still okay though, so should be fine 👍

@HarrisonWilde
Copy link
Member

Haha no worries, all looks great now, go ahead and merge!

using MCSE to get tolerances, etc. and small improvements to the rest
of the tests
@torfjelde
Copy link
Member Author

If tests pass now, I'm merging.

In an attempt to make the tolerances used in the tests a bit more appropriate, I've gone with using the MCSE and basically doing hypothesis testing. The acutal significance levels don't mean much, but it provides an intuitive knob to tune wrt. how relaxed we want to be with the tests that adapts to the "problem" at hand. We should probably also use this approach more in the tests present in runtests.jl, but I'll leave this for now.

I've also just relaxed the ESS comparisons quite a bit globally as there seems to be a bit too much variance.

@torfjelde torfjelde merged commit 180a928 into torfjelde/abstractmcmc Mar 11, 2023
@torfjelde torfjelde deleted the torfjelde/tempered-sampler-rewamp branch March 11, 2023 11:42
torfjelde added a commit that referenced this pull request Mar 11, 2023
* added CompositionSampler, RepeatedSampler, MultiSampler together with
additional methods for meta-type samplers

* added LinearAlgebra as dep

* big update but now everything finally works

* added additional pass-on-methods for meta-samplers and moved the
bundle_samples to a more appropriate place

* renamed state_from_state to state_from and changed the ordering of the
args to be more reasonable

* added some missing methods and fixed a typo

* added model_for_chain and model_for_process similar to other utility
methods for interacting with the tempered state, etc.

* added todo

* moved bundling back to ordering of defintions

* added missing test dep

* increase number of steps for one of the tests

* specialize step for combination of RepeatedSampler and MultiSampler

* Update src/sampler.jl

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* Introduction of `SwapSampler` + make `TemperedSampler`  a fancy version of `CompositionSampler` (#152)

* split the transitions and states field in TemperedState

* improved internals of CompositionSampler

* ongoing work

* added swap sampler

* added ordering specification and a TemperedComposition

* integrated work on TemperedComposition into TemperedSampler and
removed the former

* reorederd stuff so it actually works

* fixed bug in swapping computation

* added length implementation for MultiModel

* improved construct for TemperedSampler and added some convenience methods

* fixed bundle_samples for Chains and TemperedTransition

* fixed breaking bug in setparams_and_logprob!! for SwapState

* remove usage of adapted HMC in tests

* remove doubling of iterations when testing tempering

* fixed bugs with MALA and tempering

* relax atol a bit for HMC

* relax another atol

* TemperedComposition is now truly just a wrapper around a CompositionSampler

* added method for computing roundtrips

* fixed testing + added test for roundtrips

* added docs for roundtrips method

* added some tests for SwapSampler without tempering

* remove ordering from SwapSampler since it should only interact with ProcessOrdering

* simplified the sorting according to chains and processes

* added some comments

* some minor refactoring

* some refactoring + TemperedSampler now orders the samplers correctly

* remove expected_ordering and make ordering assumptions more explicit

* relax type-constraints in state_for_chain so it also works with TemperedState

* removed redundant implementations of swap_attempt

* rename swap_betas! to swap!

* moved swap_attempt as it now requires definition of SwapSampler

* removed unnecessary setparams_and_logprob!! that should never be hit
with the current codebase

* removed expected_order

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed unnecessary variable in tests

* Update src/sampler.jl

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed burn-in from step in prep for AbstractMCMC improvements

* remove getparams_and_logprob implementation for SwapState as it's
unclear what is the right approach

* split the transitions and states field in TemperedState

* improved internals of CompositionSampler

* ongoing work

* added swap sampler

* added ordering specification and a TemperedComposition

* integrated work on TemperedComposition into TemperedSampler and
removed the former

* reorederd stuff so it actually works

* fixed bug in swapping computation

* added length implementation for MultiModel

* improved construct for TemperedSampler and added some convenience methods

* fixed bundle_samples for Chains and TemperedTransition

* fixed breaking bug in setparams_and_logprob!! for SwapState

* remove usage of adapted HMC in tests

* remove doubling of iterations when testing tempering

* fixed bugs with MALA and tempering

* relax atol a bit for HMC

* relax another atol

* TemperedComposition is now truly just a wrapper around a CompositionSampler

* added method for computing roundtrips

* fixed testing + added test for roundtrips

* added docs for roundtrips method

* added some tests for SwapSampler without tempering

* remove ordering from SwapSampler since it should only interact with ProcessOrdering

* simplified the sorting according to chains and processes

* added some comments

* some minor refactoring

* some refactoring + TemperedSampler now orders the samplers correctly

* remove expected_ordering and make ordering assumptions more explicit

* relax type-constraints in state_for_chain so it also works with TemperedState

* removed redundant implementations of swap_attempt

* rename swap_betas! to swap!

* moved swap_attempt as it now requires definition of SwapSampler

* removed unnecessary setparams_and_logprob!! that should never be hit
with the current codebase

* removed expected_order

* removed unnecessary variable in tests

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed burn-in from step in prep for AbstractMCMC improvements

* remove getparams_and_logprob implementation for SwapState as it's
unclear what is the right approach

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* added CompositionTransition + quite a few bundle_samples with a
`bundle_resolve_swaps` kwarg to allow converting into chains more easily

* more samples

* reduce requirement for ess comparison for AHMC a bit

* significant improvements to the simple Gaussian example, now testing
using MCSE to get tolerances, etc. and small improvements to the rest
of the tests

* trying to debug these tests

* more debug

* fixed typy

* reduce significance even further

---------

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

---------

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>
torfjelde added a commit that referenced this pull request Mar 15, 2023
* added CompositionSampler, RepeatedSampler, MultiSampler together with
additional methods for meta-type samplers

* added LinearAlgebra as dep

* big update but now everything finally works

* added additional pass-on-methods for meta-samplers and moved the
bundle_samples to a more appropriate place

* renamed state_from_state to state_from and changed the ordering of the
args to be more reasonable

* added some missing methods and fixed a typo

* added model_for_chain and model_for_process similar to other utility
methods for interacting with the tempered state, etc.

* added todo

* moved bundling back to ordering of defintions

* added missing test dep

* increase number of steps for one of the tests

* specialize step for combination of RepeatedSampler and MultiSampler

* Update src/sampler.jl

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* Introduction of `SwapSampler` + make `TemperedSampler`  a fancy version of `CompositionSampler` (#152)

* split the transitions and states field in TemperedState

* improved internals of CompositionSampler

* ongoing work

* added swap sampler

* added ordering specification and a TemperedComposition

* integrated work on TemperedComposition into TemperedSampler and
removed the former

* reorederd stuff so it actually works

* fixed bug in swapping computation

* added length implementation for MultiModel

* improved construct for TemperedSampler and added some convenience methods

* fixed bundle_samples for Chains and TemperedTransition

* fixed breaking bug in setparams_and_logprob!! for SwapState

* remove usage of adapted HMC in tests

* remove doubling of iterations when testing tempering

* fixed bugs with MALA and tempering

* relax atol a bit for HMC

* relax another atol

* TemperedComposition is now truly just a wrapper around a CompositionSampler

* added method for computing roundtrips

* fixed testing + added test for roundtrips

* added docs for roundtrips method

* added some tests for SwapSampler without tempering

* remove ordering from SwapSampler since it should only interact with ProcessOrdering

* simplified the sorting according to chains and processes

* added some comments

* some minor refactoring

* some refactoring + TemperedSampler now orders the samplers correctly

* remove expected_ordering and make ordering assumptions more explicit

* relax type-constraints in state_for_chain so it also works with TemperedState

* removed redundant implementations of swap_attempt

* rename swap_betas! to swap!

* moved swap_attempt as it now requires definition of SwapSampler

* removed unnecessary setparams_and_logprob!! that should never be hit
with the current codebase

* removed expected_order

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed unnecessary variable in tests

* Update src/sampler.jl

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed burn-in from step in prep for AbstractMCMC improvements

* remove getparams_and_logprob implementation for SwapState as it's
unclear what is the right approach

* split the transitions and states field in TemperedState

* improved internals of CompositionSampler

* ongoing work

* added swap sampler

* added ordering specification and a TemperedComposition

* integrated work on TemperedComposition into TemperedSampler and
removed the former

* reorederd stuff so it actually works

* fixed bug in swapping computation

* added length implementation for MultiModel

* improved construct for TemperedSampler and added some convenience methods

* fixed bundle_samples for Chains and TemperedTransition

* fixed breaking bug in setparams_and_logprob!! for SwapState

* remove usage of adapted HMC in tests

* remove doubling of iterations when testing tempering

* fixed bugs with MALA and tempering

* relax atol a bit for HMC

* relax another atol

* TemperedComposition is now truly just a wrapper around a CompositionSampler

* added method for computing roundtrips

* fixed testing + added test for roundtrips

* added docs for roundtrips method

* added some tests for SwapSampler without tempering

* remove ordering from SwapSampler since it should only interact with ProcessOrdering

* simplified the sorting according to chains and processes

* added some comments

* some minor refactoring

* some refactoring + TemperedSampler now orders the samplers correctly

* remove expected_ordering and make ordering assumptions more explicit

* relax type-constraints in state_for_chain so it also works with TemperedState

* removed redundant implementations of swap_attempt

* rename swap_betas! to swap!

* moved swap_attempt as it now requires definition of SwapSampler

* removed unnecessary setparams_and_logprob!! that should never be hit
with the current codebase

* removed expected_order

* removed unnecessary variable in tests

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* removed burn-in from step in prep for AbstractMCMC improvements

* remove getparams_and_logprob implementation for SwapState as it's
unclear what is the right approach

* Apply suggestions from code review

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* added CompositionTransition + quite a few bundle_samples with a
`bundle_resolve_swaps` kwarg to allow converting into chains more easily

* more samples

* reduce requirement for ess comparison for AHMC a bit

* significant improvements to the simple Gaussian example, now testing
using MCSE to get tolerances, etc. and small improvements to the rest
of the tests

* trying to debug these tests

* more debug

* fixed typy

* reduce significance even further

---------

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>

* added docs

* added a getting started example with a simple GMM + added another
natural impl of bundle_samples

* removed now redundant compute_tempered_logdensities + added some docs
on different meta-samplers

* minor improvements to docstrings + removed reference to non-existent method

* fixed typo

* added deployment and actions for doc deployment

* fixed issue with GR plotting and headless

* fixed missing renamings

* defer design docs

* added TODO for later on

---------

Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants