General improvements and fixes

Project.toml

@@ -5,16 +5,27 @@ version = "0.1.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 
 [compat]
 AbstractMCMC = "3.2"
+ConcreteStructs = "0.2"
 Distributions = "0.24, 0.25"
+DocStringExtensions = "0.8"
+Setfield = "0.7"
 julia = "1"
 
 [extras]
+AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "AdvancedMH", "MCMCChains", "Bijectors", "StatsPlots", "LinearAlgebra"]

src/MCMCTempering.jl

@@ -4,15 +4,24 @@ import AbstractMCMC
 import Distributions
 import Random
 
+using ConcreteStructs: @concrete
+using Setfield: @set, @set!
+
+using DocStringExtensions
+
 include("adaptation.jl")
+include("swapping.jl")
 include("tempered.jl")
 include("ladders.jl")
 include("stepping.jl")
 include("model.jl")
-include("swapping.jl")
-include("plotting.jl")
 
-export tempered, TemperedSampler, plot_swaps, plot_ladders, make_tempered_model, get_tempered_loglikelihoods_and_params, make_tempered_loglikelihood, get_params
+export tempered,
+    TemperedSampler,
+    make_tempered_model,
+    StandardSwap,
+    RandomPermutationSwap,
+    NonReversibleSwap
 
 function AbstractMCMC.bundle_samples(
     ts::Vector,
@@ -22,7 +31,10 @@ function AbstractMCMC.bundle_samples(
     chain_type::Type;
     kwargs...
 )
-    AbstractMCMC.bundle_samples(ts, model, sampler.internal_sampler, state, chain_type; kwargs...)
+    AbstractMCMC.bundle_samples(
+        ts, model, sampler_for_chain(sampler, state), state_for_chain(state), chain_type;
+        kwargs...
+    )
 end
 
 end

src/adaptation.jl

@@ -1,20 +1,16 @@
-
-struct PolynomialStep
-    η :: Real
-    c :: Real
+@concrete struct PolynomialStep
+    η
+    c
 end
 function get(step::PolynomialStep, k::Real)
     step.c * (k + 1.) ^ (-step.η)
 end
 
 
-struct AdaptiveState
-    swap_target_ar :: Real
-    scale          :: Base.RefValue{<:Real}
-    step           :: PolynomialStep
-end
-function AdaptiveState(swap_target::Real, scale::Real, step::PolynomialStep)
-    AdaptiveState(swap_target, Ref(log(scale)), step)
+struct AdaptiveState{T1<:Real,T2<:Real,P<:PolynomialStep}
+    swap_target_ar :: T1
+    logscale       :: T2
+    step           :: P
 end
 
 
@@ -26,15 +22,15 @@ function init_adaptation(
 )
     Nt = length(Δ)
     step = PolynomialStep(γ, Nt - 1)
-    Ρ = [AdaptiveState(swap_target, scale, step) for _ in 1:(Nt - 1)]
+    Ρ = [AdaptiveState(swap_target, log(scale), step) for _ in 1:(Nt - 1)]
     return Ρ
 end
 
 
 function rhos_to_ladder(Ρ, Δ)
     β′ = Δ[1]
     for i in 1:length(Ρ)
-        β′ += exp(Ρ[i].scale[])
+        β′ += exp(Ρ[i].logscale)
         Δ[i + 1] = Δ[1] / β′
     end
     return Δ
@@ -48,7 +44,9 @@ function adapt_rho(ρ::AdaptiveState, swap_ar, n)
 end
 
 
-function adapt_ladder(Ρ, Δ, k, swap_ar, n)
-    Ρ[k].scale[] += adapt_rho(Ρ[k], swap_ar, n)
-    return Ρ, rhos_to_ladder(Ρ, Δ)
-end
+function adapt_ladder(P, Δ, k, swap_ar, n)
+    P[k] = let Pk = P[k]
+        @set Pk.logscale += adapt_rho(Pk, swap_ar, n)
+    end
+    return P, rhos_to_ladder(P, Δ)
+end

src/ladders.jl

@@ -1,46 +1,37 @@
 """
     get_scaling_val(Nt, swap_strategy)
 
-Calculates the correct scaling factor for polynomial step size between temperatures
+Calculates the correct scaling factor for polynomial step size between temperatures.
 """
-function get_scaling_val(Nt, swap_strategy)
-    # Why these?
-    if swap_strategy == :standard
-        scaling_val = Nt - 1
-    elseif swap_strategy == :nonrev
-        scaling_val = 2
-    else
-        scaling_val = 1
-    end
-    return scaling_val
-end
-
+get_scaling_val(Nt, ::StandardSwap) = Nt - 1
+get_scaling_val(Nt, ::NonReversibleSwap) = 2
+get_scaling_val(Nt, ::RandomPermutationSwap) = 1
 
 """
-    generate_Δ(Nt, swap_strategy)
+    generate_inverse_temperatures(Nt, swap_strategy)
 
 Returns a temperature ladder `Δ` containing `Nt` temperatures,
-generated in accordance with the chosen `swap_strategy`
+generated in accordance with the chosen `swap_strategy`.
 """
-function generate_Δ(Nt, swap_strategy)
+function generate_inverse_temperatures(Nt, swap_strategy)
     scaling_val = get_scaling_val(Nt, swap_strategy)
-    Δ = zeros(Real, Nt)
+    Δ = zeros(Nt)
     Δ[1] = 1.0
     β′ = Δ[1]
     for i ∈ 1:(Nt - 1)
-        β′ += exp(scaling_val)
+        β′ += scaling_val
         Δ[i + 1] = Δ[1] / β′
     end
     return Δ
 end
 
 
 """
-    check_Δ(Δ)
+    check_inverse_temperatures(Δ)
 
 Checks and returns a sorted `Δ` containing `{β₀, ..., βₙ}` conforming such that `1 = β₀ > β₁ > ... > βₙ ≥ 0`
 """
-function check_Δ(Δ)
+function check_inverse_temperatures(Δ)
     if !all(zero.(Δ) .≤ Δ .≤ one.(Δ))
         error("Temperature schedule provided has values outside of the acceptable range, ensure all values are in [0, 1].")
     end

src/model.jl

@@ -1,8 +1,9 @@
+"""
+    make_tempered_model(sampler, model, args...)
 
-# struct TemperedModel <: AbstractPPL.AbstractProbabilisticProgram
-#     model :: DynamicPPL.Model
-#     β     :: AbstractFloat
-# end
+Return an instance representing a model.
 
+The return-type depends on its usage in [`compute_tempered_logdensities`](@ref).
+"""
 function make_tempered_model end