-
Notifications
You must be signed in to change notification settings - Fork 4
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 compositions and products of samplers and models #151
Conversation
additional methods for meta-type samplers
bundle_samples to a more appropriate place
args to be more reasonable
methods for interacting with the tempered state, etc.
Great stuff, small comment first of all wrt your write up above, could you add a bullet for repeated samplers ( |
LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" | ||
MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are we happy with adding the MCMCChains dependency, we have avoided it until now, not strongly opposed just wanting to make sure this is considered
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
How about me make them optional dependencies? I.e. using extensions on Julia >1.8 and Requires.jl on older versions?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, I think we maybe should keep MCMCChains as a dep for now because it means that we can just always return a Chains
object in sample
which will generally be much nicer to work with for a user rather than these nested transition types you see if you don't use MCMCchains 😕
We can look at making it optional later 👍
function AbstractMCMC.bundle_samples( | ||
ts::AbstractVector, | ||
ts::AbstractVector{<:TemperedTransition}, | ||
model::AbstractMCMC.AbstractModel, | ||
sampler::TemperedSampler, | ||
state::TemperedState, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
state::TemperedState, | |
state, |
I think we should drop the type signature on state, as it means a user needs to somehow access the states to call bundle_samples
manually, this was in my other PR
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
But we are using methods that expects this type, no?
The reason behind being restrictive with types here are:
bundle_samples
should really just be compatible with whatever thestep
can return for that particular combination of inputs, hence it makes sense to specialize on the same types as we do instep
.- Not specializing on the
state
here can potentially lead to unnecessary method ambiguities since this is heavily overloaded method.
ts::AbstractVector, | ||
model::AbstractMCMC.AbstractModel, | ||
sampler::CompositionSampler, | ||
state::CompositionState, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
state::CompositionState, | |
state, |
Same as above
src/MCMCTempering.jl
Outdated
ts::AbstractVector{<:SequentialTransitions}, | ||
model::AbstractMCMC.AbstractModel, | ||
sampler::CompositionSampler, | ||
state::SequentialStates, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
state::SequentialStates, | |
state, |
Same as above
ts::AbstractVector{<:SequentialTransitions}, | ||
model::AbstractMCMC.AbstractModel, | ||
sampler::RepeatedSampler, | ||
state::SequentialStates, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
state::SequentialStates, | |
state, |
Same as above
saveall::SaveAll | ||
end | ||
|
||
RepeatedSampler(sampler, num_repeat) = RepeatedSampler(sampler, num_repeat, Val(true)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do you wrap booleans with Val
, I am curious as haven't really encountered it
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Val
converts a value into a typed instance, which means that the saveall
below can be compiled away:) You can read more about it here: https://docs.julialang.org/en/v1/manual/types/#%22Value-types%22
src/state.jl
Outdated
@@ -115,6 +115,9 @@ transition_for_chain(state::TemperedState, I...) = transition_for_process(state, | |||
Return the transition corresponding to the process indexed by `I...`. | |||
""" | |||
transition_for_process(state::TemperedState, I...) = state.transitions_and_states[I...][1] | |||
function transition_for_process(state::TemperedState{<:Tuple{<:MultipleTransitions,<:MultipleStates}}, I...) | |||
return state.transitions_and_states[1].transitions[I...] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looking over these functions, I definitely think we should separate transitions_and_states as its becoming a bit messy. I would name them inner_states
and inner_transitions
. In the future when we want to make this code distributable, I believe it might be necessary to actually have one sampler and one state per worker right, is there a way to do this now, i.e. have a shared state that can be partitioned into different parts of memory or workers?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agreed:) This has been dealt with in #151 👍
src/stepping.jl
Outdated
@@ -89,7 +93,8 @@ function AbstractMCMC.step( | |||
# Reset state | |||
@set! state.swap_acceptance_ratios = empty(state.swap_acceptance_ratios) | |||
|
|||
if should_swap(sampler, state) | |||
isswap = should_swap(sampler, state) | |||
if isswap |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This will all be redundant once we have the CompositionalSampler
of SwapSampler
and InnerSampler
right?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yep!
src/stepping.jl
Outdated
for i in 1:numtemps(sampler) | ||
] | ||
) | ||
multisampler = MultiSampler(getsampler(sampler, i) for i in 1:numtemps(sampler)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe this should happen in tempered
as I have pointed out, rather than making one at every step, should we not instead swap the entries of a single persistent MultiSampler
, or use the reference mappings on them directly?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This applies to the non-initial step
below too
Just had a proper go through and tested it out a bit, all seems to work and all looks great, happy to merge pending the suggestions and comments above. |
Co-authored-by: Harrison Wilde <harrisondwilde@outlook.com>
…n 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>
As we've spoken about before, it should be possible to make things much more composable than it currently is.
Much of this functionality is useful beyond just MCMCTempering, and so this package will provide a great "testing-ground" for functionality that we can eventually move to something like AbstractMCMC proper.
To achieve the above we need the following:
sampler_outer ∘ sampler_inner
whosestep
first stepssampler_inner
and then, using the resulting state from steppingsampler_inner
, takes a step withsampler_outer
.sampler_1 × sampler_2
which targetsmodel_1 × model_2
whosestep
simply runssampler_1
onmodel_1
andsampler_2
onmodel_2
in "parallel".SequentialState
) which represents a sequence of states so we can later "unroll" these to obtain the full chain, if so desired. E.g. in the case ofsampler_outer ∘ sampler_inner
it's not completely obvious if we only want the iterations corresponding to the full composition, or if we actually want the iterations of each of the samplers.All in all, it's now possible to do some neat thingies.
Neat thingies
Some setup.
Compositions:
Same with "repeated" samplers (useful to avoid insane compilation times):
Products:
And of course, combining all of these:
MultiSampler
also works withAbstractArray
, etc.Drawbacks
foldl(∘, repeat(sampler, 1000))
and cry. This is fixable though (see for example how Julia handlesFunction
without compositions blowing up compilation times), probably by sprinkling@nospecialize
in the right places.sample
and/orAbstractMCMC.bundle_samples
. In the case where we wantMCMCChains.Chains
as the return-type, it seems reasonable to "unroll" things likeSequentialTransition
, but in the case where a simpleVector
of transitions is wanted, then it seems less so, i.e. in this case we probably want to return aVector{<:SequentialTransition}
which can later be "unrolled" if so desired.All in all, it's worth pointing out that we're not expecting users to nest these "meta"-samplers very often; in practice we'll probably look at a nested depth of 2 or 3 at most (e.g. mixture of composition of repeated).
Next steps
This PR on its own doesn't bring us all the way to just doing
sampler_outer ∘ sampler_swapper
, but this should now be straight forward to implement by just overloadingstep
forCompositionSampler{<:AbstractMCMC.AbstractSampler,<:TemperedSampler}
to performstep
forsampler_outer
and then a swap-step, finally returning aSequentialTransition
andSequentialState
, or just aCompositionState
withtransition
fromsampler_outer
(in the case where we don't want to further inspect the swapping statistics, etc.).I'll do this in another PR, as it will change
TemperedSampler
quite a bit.