chalk-lab
diff --git a/‎Project.toml‎
Lines changed: 2 additions & 0 deletions b/‎Project.toml‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎src/MCMCTempering.jl‎
Lines changed: 191 additions & 3 deletions b/‎src/MCMCTempering.jl‎
Lines changed: 191 additions & 3 deletions
diff --git a/‎src/abstractmcmc.jl‎
Lines changed: 106 additions & 0 deletions b/‎src/abstractmcmc.jl‎
Lines changed: 106 additions & 0 deletions
diff --git a/‎src/adaptation.jl‎
Lines changed: 1 addition & 1 deletion b/‎src/adaptation.jl‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎src/ladders.jl‎
Lines changed: 1 addition & 3 deletions b/‎src/ladders.jl‎
Lines changed: 1 addition & 3 deletions
@@ -9,7 +9,9 @@ ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 
@@ -9,19 +9,24 @@ using ProgressLogging: ProgressLogging
 using ConcreteStructs: @concrete
 using Setfield: @set, @set!
 
+using MCMCChains: MCMCChains
+
 using InverseFunctions
 
 using DocStringExtensions
 
 include("logdensityproblems.jl")
+include("abstractmcmc.jl")
 include("adaptation.jl")
 include("swapping.jl")
 include("state.jl")
+include("swapsampler.jl")
 include("sampler.jl")
 include("sampling.jl")
 include("ladders.jl")
 include("stepping.jl")
 include("model.jl")
+include("utils.jl")
 
 export tempered,
     tempered_sample,
@@ -39,16 +44,199 @@ implements_logdensity(x) = LogDensityProblems.capabilities(x) !== nothing
 maybe_wrap_model(model) = implements_logdensity(model) ? AbstractMCMC.LogDensityModel(model) : model
 maybe_wrap_model(model::AbstractMCMC.LogDensityModel) = model
 
+# Bundling.
+# Bundling of non-tempered samples.
+function bundle_nontempered_samples(
+    ts::AbstractVector{<:TemperedTransition{<:SwapTransition,<:MultipleTransitions}},
+    model::AbstractMCMC.AbstractModel,
+    sampler::TemperedSampler,
+    state::TemperedState,
+    ::Type{T};
+    kwargs...
+) where {T}
+    # Create the same model and sampler as we do in the initial step for `TemperedSampler`.
+    multimodel = MultiModel([
+        make_tempered_model(sampler, model, sampler.chain_to_beta[i])
+        for i in 1:numtemps(sampler)
+    ])
+    multisampler = MultiSampler([getsampler(sampler, i) for i in 1:numtemps(sampler)])
+    multitransitions = [
+        MultipleTransitions(sort_by_chain(ProcessOrder(), t.swaptransition, t.transition.transitions))
+        for t in ts
+    ]
+
+    return AbstractMCMC.bundle_samples(
+        multitransitions,
+        multimodel,
+        multisampler,
+        MultipleStates(sort_by_chain(ProcessOrder(), state.swapstate, state.state.states)),
+        T
+    )
+end
+
 function AbstractMCMC.bundle_samples(
-    ts::AbstractVector,
+    ts::Vector{<:MultipleTransitions},
+    model::MultiModel,
+    sampler::MultiSampler,
+    state::MultipleStates,
+    # TODO: Generalize for any eltype `T`? Then need to overload for `Real`, etc.?
+    ::Type{Vector{MCMCChains.Chains}};
+    kwargs...
+)
+    return map(1:length(model), model.models, sampler.samplers, state.states) do i, model, sampler, state
+        AbstractMCMC.bundle_samples([t.transitions[i] for t in ts], model, sampler, state, MCMCChains.Chains; kwargs...)
+    end
+end
+
+# HACK: https://github.com/TuringLang/AbstractMCMC.jl/issues/118
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:TemperedTransition{<:SwapTransition,<:MultipleTransitions}},
     model::AbstractMCMC.AbstractModel,
     sampler::TemperedSampler,
     state::TemperedState,
-    chain_type::Type;
+    ::Type{Vector{T}};
+    bundle_resolve_swaps::Bool=false,
+    kwargs...
+) where {T}
+    if bundle_resolve_swaps
+        return bundle_nontempered_samples(ts, model, sampler, state, Vector{T}; kwargs...)
+    end
+
+    # TODO: Do better?
+    return ts
+end
+
+function AbstractMCMC.bundle_samples(
+    ts::AbstractVector{<:TemperedTransition{<:SwapTransition,<:MultipleTransitions}},
+    model::AbstractMCMC.AbstractModel,
+    sampler::TemperedSampler,
+    state::TemperedState,
+    ::Type{MCMCChains.Chains};
     kwargs...
 )
+    # Extract the transitions ordered, which are ordered according to processes, according to the chains.
+    ts_actual = [t.transition.transitions[first(t.swaptransition.chain_to_process)] for t in ts]
+    return AbstractMCMC.bundle_samples(
+        ts_actual,
+        model,
+        sampler_for_chain(sampler, state, 1),
+        state_for_chain(state, 1),
+        MCMCChains.Chains;
+        kwargs...
+    )
+end
+
+function AbstractMCMC.bundle_samples(
+    ts::AbstractVector,
+    model::AbstractMCMC.AbstractModel,
+    sampler::CompositionSampler,
+    state::CompositionState,
+    ::Type{T};
+    kwargs...
+) where {T}
+    # In the case of `!saveall(sampler)`, the state is not a `CompositionTransition` so we just propagate
+    # the transitions to the `bundle_samples` for the outer stuff. Otherwise, we flatten the transitions.
+    ts_actual = saveall(sampler) ? mapreduce(t -> [inner_transition(t), outer_transition(t)], vcat, ts) : ts
+    # TODO: Should we really always default to outer sampler?
+    return AbstractMCMC.bundle_samples(
+        ts_actual, model, sampler.sampler_outer, state.state_outer, T;
+        kwargs...
+    )
+end
+
+# HACK: https://github.com/TuringLang/AbstractMCMC.jl/issues/118
+function AbstractMCMC.bundle_samples(
+    ts::Vector,
+    model::AbstractMCMC.AbstractModel,
+    sampler::CompositionSampler,
+    state::CompositionState,
+    ::Type{Vector{T}};
+    kwargs...
+) where {T}
+    if !saveall(sampler)
+        # In this case, we just use the `outer` for everything since this is the only
+        # transitions we're keeping around.
+        return AbstractMCMC.bundle_samples(
+            ts, model, sampler.sampler_outer, state.state_outer, Vector{T};
+            kwargs...
+        )
+    end
+
+    # Otherwise, we don't know what to do.
+    return ts
+end
+
+function AbstractMCMC.bundle_samples(
+    ts::AbstractVector{<:CompositionTransition{<:MultipleTransitions,<:SwapTransition}},
+    model::AbstractMCMC.AbstractModel,
+    sampler::CompositionSampler{<:MultiSampler,<:SwapSampler},
+    state::CompositionState{<:MultipleStates,<:SwapState},
+    ::Type{T};
+    bundle_resolve_swaps::Bool=false,
+    kwargs...
+) where {T}
+    !bundle_resolve_swaps && return ts
+
+    # Resolve the swaps.
+    sampler_without_saveall = @set sampler.sampler_inner.saveall = Val(false)
+    ts_actual = map(ts) do t
+        composition_transition(sampler_without_saveall, inner_transition(t), outer_transition(t))
+    end
+
     AbstractMCMC.bundle_samples(
-        ts, maybe_wrap_model(model), sampler_for_chain(sampler, state, 1), state_for_chain(state, 1), chain_type;
+        ts_actual, model, sampler.sampler_outer, state.state_outer, T;
+        kwargs...
+    )
+end
+
+# HACK: https://github.com/TuringLang/AbstractMCMC.jl/issues/118
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:CompositionTransition{<:MultipleTransitions,<:SwapTransition}},
+    model::AbstractMCMC.AbstractModel,
+    sampler::CompositionSampler{<:MultiSampler,<:SwapSampler},
+    state::CompositionState{<:MultipleStates,<:SwapState},
+    ::Type{Vector{T}};
+    bundle_resolve_swaps::Bool=false,
+    kwargs...
+) where {T}
+    !bundle_resolve_swaps && return ts
+
+    # Resolve the swaps (using the already implemented resolution in `composition_transition`
+    # for this particular sampler but without `saveall`).
+    sampler_without_saveall = @set sampler.saveall = Val(false)
+    ts_actual = map(ts) do t
+        composition_transition(sampler_without_saveall, inner_transition(t), outer_transition(t))
+    end
+
+    return AbstractMCMC.bundle_samples(
+        ts_actual, model, sampler.sampler_outer, state.state_outer, Vector{T};
+        kwargs...
+    )
+end
+
+function AbstractMCMC.bundle_samples(
+    ts::AbstractVector,
+    model::AbstractMCMC.AbstractModel,
+    sampler::RepeatedSampler,
+    state,
+    ::Type{MCMCChains.Chains};
+    kwargs...
+)
+    return AbstractMCMC.bundle_samples(ts, model, sampler.sampler, state, MCMCChains.Chains; kwargs...)
+end
+
+# Unflatten in the case of `SequentialTransitions`.
+function AbstractMCMC.bundle_samples(
+    ts::AbstractVector{<:SequentialTransitions},
+    model::AbstractMCMC.AbstractModel,
+    sampler::RepeatedSampler,
+    state::SequentialStates,
+    ::Type{MCMCChains.Chains};
+    kwargs...
+)
+    ts_actual = [t for tseq in ts for t in tseq.transitions]
+    return AbstractMCMC.bundle_samples(
+        ts_actual, model, sampler.sampler, state.states[end], MCMCChains.Chains;
         kwargs...
     )
 end
 
@@ -0,0 +1,106 @@
+using Setfield
+using AbstractMCMC: AbstractMCMC
+
+import LinearAlgebra: ×
+
+"""
+    getparams([model, ]state)
+
+Get the parameters from the `state`.
+
+Default implementation uses [`getparams_and_logprob`](@ref).
+"""
+getparams(state) = first(getparams_and_logprob(state))
+getparams(model, state) = first(getparams_and_logprob(model, state))
+
+"""
+    getlogprob([model, ]state)
+
+Get the log probability of the `state`.
+
+Default implementation uses [`getparams_and_logprob`](@ref).
+"""
+getlogprob(state) = last(getparams_and_logprob(state))
+getlogprob(model, state) = last(getparams_and_logprob(model, state))
+
+"""
+    getparams_and_logprob([model, ]state)
+
+Return a vector of parameters from the `state`.
+
+See also: [`setparams_and_logprob!!`](@ref).
+"""
+getparams_and_logprob(model, state) = getparams_and_logprob(state)
+
+"""
+    setparams_and_logprob!!([model, ]state, params)
+
+Set the parameters in the state to `params`, possibly mutating if it makes sense.
+
+See also: [`getparams_and_logprob`](@ref).
+"""
+setparams_and_logprob!!(model, state, params, logprob) = setparams_and_logprob!!(state, params, logprob)
+
+"""
+    state_from(model, state_target, state_source[, transition_source, transition_target])
+
+Return a new state similar to `state_target` but updated from `state_source`, which could be
+a different type of state.
+"""
+function state_from(model, state_target, state_source, transition_target, transition_source)
+    return state_from(model, state_target, state_source)
+end
+function state_from(model, state_target, state_source)
+    params, logp = getparams_and_logprob(model, state_source)
+    return setparams_and_logprob!!(model, state_target, params, logp)
+end
+
+"""
+    SequentialTransitions
+
+A `SequentialTransitions` object is a container for a sequence of transitions.
+"""
+struct SequentialTransitions{A}
+    transitions::A
+end
+
+# Since it's a _sequence_ of transitions, the parameters and logprobs are the ones of the
+# last transition/state.
+getparams_and_logprob(transitions::SequentialTransitions) = getparams_and_logprob(transitions.transitions[end])
+function getparams_and_logprob(model, transitions::SequentialTransitions)
+    return getparams_and_logprob(model, transitions.transitions[end])
+end
+
+function setparams_and_logprob!!(transitions::SequentialTransitions, params, logprob)
+    return @set transitions.transitions[end] = setparams_and_logprob!!(transitions.transitions[end], params, logprob)
+end
+function setparams_and_logprob!!(model, transitions::SequentialTransitions, params, logprob)
+    return @set transitions.transitions[end] = setparams_and_logprob!!(model, transitions.transitions[end], params, logprob)
+end
+
+"""
+    SequentialStates
+
+A `SequentialStates` object is a container for a sequence of states.
+"""
+struct SequentialStates{A}
+    states::A
+end
+
+# Since it's a _sequence_ of transitions, the parameters and logprobs are the ones of the
+# last transition/state.
+getparams_and_logprob(state::SequentialStates) = getparams_and_logprob(state.states[end])
+getparams_and_logprob(model, state::SequentialStates) = getparams_and_logprob(model, state.states[end])
+
+function setparams_and_logprob!!(state::SequentialStates, params, logprob)
+    return @set state.states[end] = setparams_and_logprob!!(state.states[end], params, logprob)
+end
+function setparams_and_logprob!!(model, state::SequentialStates, params, logprob)
+    return @set state.states[end] = setparams_and_logprob!!(model, state.states[end], params, logprob)
+end
+
+# Includes.
+include("samplers/composition.jl")
+include("samplers/repeated.jl")
+include("samplers/multi.jl")
+
@@ -19,7 +19,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
 
@@ -37,9 +37,7 @@ 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
+    !isempty(Δ) || error("Inverse temperatures array is empty.")
     if !all(zero.(Δ) .≤ Δ .≤ one.(Δ))
         error("The temperature ladder provided has values outside of the acceptable range, ensure all values are in [0, 1].")
     end