Introduction of SwapSampler + make TemperedSampler a fancy version of CompositionSampler

src/MCMCTempering.jl

@@ -20,6 +20,7 @@ include("abstractmcmc.jl")
 include("adaptation.jl")
 include("swapping.jl")
 include("state.jl")
+include("swapsampler.jl")
 include("sampler.jl")
 include("sampling.jl")
 include("ladders.jl")
@@ -46,18 +47,21 @@ maybe_wrap_model(model::AbstractMCMC.LogDensityModel) = model
 # TODO: Improve this, somehow.
 # TODO: Move this to an extension.
 function AbstractMCMC.bundle_samples(
-    ts::AbstractVector{<:TemperedTransition},
+    ts::AbstractVector{<:TemperedTransition{<:SwapTransition,<:MultipleTransitions}},
     model::AbstractMCMC.AbstractModel,
     sampler::TemperedSampler,
     state::TemperedState,
     ::Type{MCMCChains.Chains};
     kwargs...
 )
+    # Extract the transitions ordered according to the chains.
+    # TODO: Improve this.
+    ts_actual = [t.transition.transitions[first(t.swaptransition.chain_to_process)] for t in ts]
     return AbstractMCMC.bundle_samples(
-        map(Base.Fix2(getproperty, :transition), filter(!Base.Fix2(getproperty, :is_swap), ts)),  # Remove the swaps.
+        ts_actual,
         model,
-        sampler_for_chain(sampler, state),
-        state_for_chain(state),
+        sampler_for_chain(sampler, state, 1),
+        state_for_chain(state.swapstate),
         MCMCChains.Chains;
         kwargs...
     )

src/adaptation.jl

@@ -18,7 +18,7 @@ See also: [`AdaptiveState`](@ref), [`update_inverse_temperatures`](@ref), and
 """
 struct Geometric end
 
-defaultscale(::Geometric, Δ) = eltype(Δ)(0.9)
+defaultscale(::Geometric, Δ) = float(eltype(Δ))(0.9)
 
 """
     InverselyAdditive

src/ladders.jl

@@ -37,9 +37,6 @@ end
 Checks and returns a sorted `Δ` containing `{β₀, ..., βₙ}` conforming such that `1 = β₀ > β₁ > ... > βₙ ≥ 0`
 """
 function check_inverse_temperatures(Δ)
-    if length(Δ) <= 1
-        error("More than one inverse temperatures must be provided.")
-    end
     if !all(zero.(Δ) .≤ Δ .≤ one.(Δ))
         error("The temperature ladder provided has values outside of the acceptable range, ensure all values are in [0, 1].")
     end

src/sampler.jl

@@ -1,3 +1,20 @@
+"""
+    TemperedState
+
+A state for a tempered sampler.
+
+# Fields
+$(FIELDS)
+"""
+@concrete struct TemperedState
+    "state for swap-sampler"
+    swapstate
+    "state for the main sampler"
+    state
+    "inverse temperature for each of the chains"
+    chain_to_beta
+end
+
 """
     TemperedSampler <: AbstractMCMC.AbstractSampler
 
@@ -7,44 +24,39 @@ A `TemperedSampler` struct wraps a sampler upon which to apply the Parallel Temp
 
 $(FIELDS)
 """
-@concrete struct TemperedSampler <: AbstractMCMC.AbstractSampler
+Base.@kwdef struct TemperedSampler{SplT,A,SwapT,Adapt} <: AbstractMCMC.AbstractSampler
     "sampler(s) used to target the tempered distributions"
-    sampler
+    sampler::SplT
     "collection of inverse temperatures β; β[i] correponds i-th tempered model"
-    inverse_temperatures
-    "number of steps of `sampler` to take before proposing swaps"
-    swap_every
-    "the swap strategy that will be used when proposing swaps"
-    swap_strategy
-    # TODO: This should be replaced with `P` just being some `NoAdapt` type.
+    chain_to_beta::A
+    "strategy to use for swapping"
+    swapstrategy::SwapT=ReversibleSwap()
+    # TODO: Remove `adapt` and just consider `adaptation_states=nothing` as no adaptation.
     "boolean flag specifying whether or not to adapt"
-    adapt
+    adapt=false
     "adaptation parameters"
-    adaptation_states
+    adaptation_states::Adapt=nothing
 end
 
-swapstrategy(sampler::TemperedSampler) = sampler.swap_strategy
+TemperedSampler(sampler, chain_to_beta; kwargs...) = TemperedSampler(; sampler, chain_to_beta, kwargs...)
+
+swapsampler(sampler::TemperedSampler) = SwapSampler(sampler.swapstrategy, ProcessOrdering())
 
+# TODO: Do we need this now?
 getsampler(samplers, I...) = getindex(samplers, I...)
 getsampler(sampler::AbstractMCMC.AbstractSampler, I...) = sampler
 getsampler(sampler::TemperedSampler, I...) = getsampler(sampler.sampler, I...)
 
-"""
-    numsteps(sampler::TemperedSampler)
-
-Return number of inverse temperatures used by `sampler`.
-"""
-numtemps(sampler::TemperedSampler) = length(sampler.inverse_temperatures)
+chain_to_process(state::TemperedState, I...) = chain_to_process(state.swapstate, I...)
+process_to_chain(state::TemperedState, I...) = process_to_chain(state.swapstate, I...)
 
 """
-    sampler_for_chain(sampler::TemperedSampler, state::TemperedState[, I...])
+    sampler_for_chain(sampler::TemperedSampler, state::TemperedState, I...)
 
 Return the sampler corresponding to the chain indexed by `I...`.
-If `I...` is not specified, the sampler corresponding to `β=1.0` will be returned.
 """
-sampler_for_chain(sampler::TemperedSampler, state::TemperedState) = sampler_for_chain(sampler, state, 1)
 function sampler_for_chain(sampler::TemperedSampler, state::TemperedState, I...)
-    return getsampler(sampler.sampler, chain_to_process(state, I...))
+    return sampler_for_process(sampler, state, chain_to_process(state, I...))
 end
 
 """
@@ -53,9 +65,51 @@ end
 Return the sampler corresponding to the process indexed by `I...`.
 """
 function sampler_for_process(sampler::TemperedSampler, state::TemperedState, I...)
-    return getsampler(sampler.sampler, I...)
+    return _sampler_for_process_temper(sampler.sampler, state, I...)
 end
 
+# If `sampler` is a `MultiSampler`, we assume it's ordered according to chains.
+_sampler_for_process_temper(sampler::MultiSampler, state, I...) = sampler.samplers[process_to_chain(state, I...)]
+# Otherwise, we just use the same sampler for everything.
+_sampler_for_process_temper(sampler, state, I...) = sampler
+
+# Defer extracting the corresponding state to the `swapstate`.
+state_for_process(state::TemperedState, I...) = state_for_process(state.swapstate, I...)
+
+# Here we make the model(s) using the temperatures.
+function model_for_process(sampler::TemperedSampler, model, state::TemperedState, I...)
+    return make_tempered_model(sampler, model, beta_for_process(state, I...))
+end
+
+"""
+    beta_for_chain(state[, I...])
+
+Return the β corresponding to the chain indexed by `I...`.
+If `I...` is not specified, the β corresponding to `β=1.0` will be returned.
+"""
+beta_for_chain(state::TemperedState) = beta_for_chain(state, 1)
+beta_for_chain(state::TemperedState, I...) = beta_for_chain(state.chain_to_beta, I...)
+# NOTE: Array impl. is useful for testing.
+beta_for_chain(chain_to_beta::AbstractArray, I...) = chain_to_beta[I...] 
+
+"""
+    beta_for_process(state, I...)
+
+Return the β corresponding to the process indexed by `I...`.
+"""
+beta_for_process(state::TemperedState, I...) = beta_for_process(state.chain_to_beta, state.swapstate.process_to_chain, I...)
+# NOTE: Array impl. is useful for testing.
+function beta_for_process(chain_to_beta::AbstractArray, proc2chain::AbstractArray, I...)
+    return beta_for_chain(chain_to_beta, process_to_chain(proc2chain, I...))
+end
+
+"""
+    numsteps(sampler::TemperedSampler)
+
+Return number of inverse temperatures used by `sampler`.
+"""
+numtemps(sampler::TemperedSampler) = length(sampler.chain_to_beta)
+
 """
     tempered(sampler, inverse_temperatures; kwargs...)
     OR
@@ -118,5 +172,5 @@ function tempered(
     # TODO: Generalize. Allow passing in a `MultiSampler`, etc.
     sampler_inner = sampler^swap_every
     # FIXME: Remove the hard-coded `2` for swap-every, and change `should_swap` acoordingly.
-    return TemperedSampler(sampler_inner, inverse_temperatures, 2, swap_strategy, adapt, adaptation_states)
+    return TemperedSampler(sampler_inner, inverse_temperatures, swap_strategy, adapt, adaptation_states)
 end

src/samplers/composition.jl

@@ -58,40 +58,57 @@ function setparams_and_logprob!!(model, state::CompositionState, params, logprob
     return @set state.state_outer = setparams_and_logprob!!(model, state.state_outer, params, logprob)
 end
 
+# Useful functions for interacting with composition sampler and states.
+inner_sampler(sampler::CompositionSampler) = sampler.sampler_inner
+outer_sampler(sampler::CompositionSampler) = sampler.sampler_outer
+
+inner_state(state::CompositionState) = state.state_inner
+outer_state(state::CompositionState) = state.state_outer
+
+inner_state(state::SequentialStates) = first(state.states)
+outer_state(state::SequentialStates) = last(state.states)
+
+function composition_state(sampler, state_inner, state_outer)
+    return if saveall(sampler)
+        SequentialStates((state_inner, state_outer))
+    else
+        CompositionState(state_outer, state_inner)
+    end
+end
+function composition_transition(sampler, transition_inner, transition_outer)
+    return if saveall(sampler)
+        SequentialTransitions((transition_inner, transition_outer))
+    else
+        transition_outer
+    end
+end
+
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.AbstractModel,
     sampler::CompositionSampler;
     kwargs...
 )
-    state_inner_initial = last(AbstractMCMC.step(rng, model, sampler.sampler_inner; kwargs...))
-    state_outer_initial = last(AbstractMCMC.step(rng, model, sampler.sampler_outer; kwargs...))
+    state_inner_initial = last(AbstractMCMC.step(rng, model, inner_sampler(sampler); kwargs...))
+    state_outer_initial = last(AbstractMCMC.step(rng, model, outer_sampler(sampler); kwargs...))
 
     # Create the composition state, and take a full step.
-    state = if saveall(sampler)
-        SequentialStates((state_inner_initial, state_outer_initial))
-    else
-        CompositionState(state_outer_initial, state_inner_initial)
-    end
+    state = composition_state(sampler, state_inner_initial, state_outer_initial)
     return AbstractMCMC.step(rng, model, sampler, state; kwargs...)
 end
 
 # TODO: Do we even need two versions? We could technically use `SequentialStates`
 # in place of `CompositionState` and just have one version.
 # The annoying part here is that we'll have to check `saveall` on every `step`
 # rather than just for the initial step.
-
-# NOTE: Version which does keep track of all transitions and states.
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.AbstractModel,
     sampler::CompositionSampler,
-    state::SequentialStates;
+    state;
     kwargs...
 )
-    @assert length(state.states) == 2 "Composition samplers only support SequentialStates with two states."
-
-    state_inner_prev, state_outer_prev = state.states
+    state_inner_prev, state_outer_prev = inner_state(state), outer_state(state)
 
     # Update the inner state.
     current_state_inner = state_from(model, state_inner_prev, state_outer_prev)
@@ -103,29 +120,8 @@ function AbstractMCMC.step(
     current_state_outer = state_from(model, state_outer_prev, state_inner)
     transition_outer, state_outer = AbstractMCMC.step(rng, model, sampler.sampler_outer, current_state_outer; kwargs...)
 
-    return SequentialTransitions((transition_inner, transition_outer)), SequentialStates((state_inner, state_outer))
-end
-
-# NOTE: Version which does NOT keep track of all transitions and states.
-function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    model::AbstractMCMC.AbstractModel,
-    sampler::CompositionSampler,
-    state::CompositionState;
-    kwargs...
-)
-    # Update the inner state.
-    current_state_inner = state_from(model, state.state_inner, state.state_outer)
-
-    # Take a step in the inner sampler.
-    state_inner = last(AbstractMCMC.step(rng, model, sampler.sampler_inner, current_state_inner; kwargs...))
-
-    # Take a step in the outer sampler.
-    current_state_outer = state_from(model, state.state_outer, state_inner)
-    transition_outer, state_outer = AbstractMCMC.step(rng, model, sampler.sampler_outer, current_state_outer; kwargs...)
-
-    # Create the composition state.
-    state = CompositionState(state_outer, state_inner)
-
-    return transition_outer, state
+    return (
+        composition_transition(sampler, transition_inner, transition_outer),
+        composition_state(sampler, state_inner, state_outer)
+    )
 end

src/samplers/multi.jl

@@ -57,6 +57,8 @@ end
 ×(model1::AbstractMCMC.AbstractModel, model2::MultiModel) = MultiModel(combine(model1, model2.models))
 ×(model1::MultiModel, model2::MultiModel) = MultiModel(combine(model1.models, model2.models))
 
+Base.length(model::MultiModel) = length(model.models)
+
 # TODO: Make these subtypes of `AbstractVector`?
 """
     MultipleTransitions
@@ -108,15 +110,23 @@ function getparams_and_logprob(model::MultiModel, state::MultipleStates)
     return map(first, params_and_logprobs), map(last, params_and_logprobs)
 end
 
-function setparams_and_logprob!!(state::MultipleStates, params, logprob)
-    @assert length(params) == length(logprob) == length(state.states) "The number of parameters and log probabilities must match the number of states."
-    return @set state.states = map(setparams_and_logprob!!, state.states, params, logprob)
+function setparams_and_logprob!!(state::MultipleStates, params, logprobs)
+    @assert length(params) == length(logprobs) == length(state.states) "The number of parameters and log probabilities must match the number of states."
+    return @set state.states = map(setparams_and_logprob!!, state.states, params, logprobs)
+end
+function setparams_and_logprob!!(model::MultiModel, state::MultipleStates, params, logprobs)
+    @assert length(model.models) == length(params) == length(logprobs) == length(state.states) "The number of models, states, parameters, and log probabilities must match."
+    return @set state.states = map(setparams_and_logprob!!, model.models, state.states, params, logprobs)
 end
-function setparams_and_logprob!!(model::MultiModel, state::MultipleStates, params, logprob)
-    @assert length(params) == length(logprob) == length(state.states) "The number of parameters and log probabilities must match the number of states."
-    return @set state.states = map(setparams_and_logprob!!, model.models, state.states, params, logprob)
+# NOTE: If we're not working with a `MultiModel`, we assume we just have to pass it on.
+function setparams_and_logprob!!(model, state::MultipleStates, params, logprobs)
+    @assert length(params) == length(logprobs) == length(state.states) "The number of parameters and log probabilities must match the number of states."
+    return @set state.states = map(state.states, params, logprobs) do state, param, logprob
+        setparams_and_logprob!!(model, state, param, logprob)
+    end
 end
 
+
 # TODO: Clean this up.
 initparams(model::MultiModel, init_params) = map(Base.Fix1(get_init_params, init_params), 1:length(model.models))
 initparams(model::MultiModel{<:Tuple}, init_params) =  ntuple(length(model.models)) do i