Introduce RepeatSampler

Author: mhauru

HISTORY.md

@@ -12,7 +12,7 @@ may be accidental breakage that we haven't anticipated. Please report any you fi
 
 The old Gibbs constructor relied on being called with several subsamplers, and each of the constructors of the subsamplers would take as arguments the symbols for the variables that they are to sample, e.g. `Gibbs(HMC(:x), MH(:y))`. This constructor has been deprecated, and will be removed in the future. The new constructor works by assigning samplers to either symbols or `VarNames`, e.g. `Gibbs(; x=HMC(), y=MH())` or `Gibbs(@varname(x) => HMC(), @varname(y) => MH())`. This allows more granular specification of which sampler to use for which variable.
 
-Likewise, the old constructor for calling one subsampler more often than another, `Gibbs((HMC(:x), 2), (MH(:y), 1))` has been deprecated. The new way to achieve this effect is to list the same sampler multiple times, e.g. as `hmc = HMC(); mh = MH(); Gibbs(@varname(x) => hmc, @varname(x) => hmc, @varname(y) => mh)`.
+Likewise, the old constructor for calling one subsampler more often than another, `Gibbs((HMC(0.01, 4, :x), 2), (MH(:y), 1))` has been deprecated. The new way to do this is to use `RepeatSampler`, also introduced at this version: `Gibbs(@varname(x) => RepeatSampler(HMC(0.01, 4), 2), @varname(y) => MH())`.
 
 # Release 0.35.0
 

src/Turing.jl

@@ -95,6 +95,7 @@ export @model,                 # modelling
     SMC,
     CSMC,
     PG,
+    RepeatSampler,
     vi,                     # variational inference
     ADVI,
     sample,                 # inference

src/mcmc/Inference.jl

@@ -74,6 +74,7 @@ export InferenceAlgorithm,
     SMC,
     CSMC,
     PG,
+    RepeatSampler,
     Prior,
     assume,
     dot_assume,
@@ -100,6 +101,12 @@ Return an `InferenceAlgorithm` like `alg`, but with all space information remove
 """
 function drop_space end
 
+function drop_space(sampler::Sampler)
+    return Sampler(drop_space(sampler.alg), sampler.selector)
+end
+
+include("repeat_sampler.jl")
+
 """
     ExternalSampler{S<:AbstractSampler,AD<:ADTypes.AbstractADType,Unconstrained}
 
@@ -348,7 +355,7 @@ end
 function AbstractMCMC.sample(
     rng::AbstractRNG,
     model::AbstractModel,
-    sampler::Sampler{<:InferenceAlgorithm},
+    sampler::Union{Sampler{<:InferenceAlgorithm},RepeatSampler},
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
@@ -460,7 +467,7 @@ getlogevidence(transitions, sampler, state) = missing
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Union{AbstractTransition,AbstractVarInfo}},
     model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
+    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
     state,
     chain_type::Type{MCMCChains.Chains};
     save_state=false,
@@ -523,7 +530,7 @@ end
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Union{AbstractTransition,AbstractVarInfo}},
     model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
+    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
     state,
     chain_type::Type{Vector{NamedTuple}};
     kwargs...,

src/mcmc/gibbs.jl

@@ -15,6 +15,8 @@ isgibbscomponent(::NUTS) = true
 isgibbscomponent(::MH) = true
 isgibbscomponent(::PG) = true
 
+isgibbscomponent(spl::RepeatSampler) = isgibbscomponent(spl.sampler)
+
 isgibbscomponent(spl::ExternalSampler) = isgibbscomponent(spl.sampler)
 isgibbscomponent(::AdvancedHMC.HMC) = true
 isgibbscomponent(::AdvancedMH.MetropolisHastings) = true
@@ -364,7 +366,7 @@ function Gibbs(algs::InferenceAlgorithm...)
         "`Gibbs(NUTS(:x), MH(:y))` is deprecated and will be removed in the future. " *
         "Please use `Gibbs(; x=NUTS(), y=MH())` instead. If you want different iteration " *
         "counts for different subsamplers, use e.g. " *
-        "`Gibbs(@varname(x) => NUTS(), @varname(x) => NUTS(), @varname(y) => MH())`"
+        "`Gibbs(@varname(x) => RepeatSampler(NUTS(), 2), @varname(y) => MH())`"
     )
     Base.depwarn(msg, :Gibbs)
     return Gibbs(varnames, map(wrap_algorithm_maybe ∘ drop_space, algs))

src/mcmc/repeat_sampler.jl

@@ -0,0 +1,62 @@
+"""
+    RepeatSampler <: AbstractMCMC.AbstractSampler
+
+A `RepeatSampler` is a container for a sampler and a number of times to repeat it.
+
+# Fields
+$(FIELDS)
+
+# Examples
+```julia
+repeated_sampler = RepeatSampler(sampler, 10)
+AbstractMCMC.step(rng, model, repeated_sampler) # take 10 steps of `sampler`
+```
+"""
+struct RepeatSampler{S<:AbstractMCMC.AbstractSampler} <: AbstractMCMC.AbstractSampler
+    "The sampler to repeat"
+    sampler::S
+    "The number of times to repeat the sampler"
+    num_repeat::Int
+
+    function RepeatSampler(sampler::S, num_repeat::Int) where {S}
+        @assert num_repeat > 0
+        return new{S}(sampler, num_repeat)
+    end
+end
+
+function RepeatSampler(alg::InferenceAlgorithm, num_repeat::Int)
+    return RepeatSampler(Sampler(alg), num_repeat)
+end
+
+drop_space(rs::RepeatSampler) = RepeatSampler(drop_space(rs.sampler), rs.num_repeat)
+getADType(spl::RepeatSampler) = getADType(spl.sampler)
+DynamicPPL.default_chain_type(sampler::RepeatSampler) = default_chain_type(sampler.sampler)
+DynamicPPL.getspace(spl::RepeatSampler) = getspace(spl.sampler)
+DynamicPPL.inspace(vn::VarName, spl::RepeatSampler) = inspace(vn, spl.sampler)
+
+function setparams_varinfo!!(model::DynamicPPL.Model, sampler::RepeatSampler, state, params)
+    return setparams_varinfo!!(model, sampler.sampler, state, params)
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.AbstractModel,
+    sampler::RepeatSampler;
+    kwargs...,
+)
+    return AbstractMCMC.step(rng, model, sampler.sampler; kwargs...)
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.AbstractModel,
+    sampler::RepeatSampler,
+    state;
+    kwargs...,
+)
+    transition, state = AbstractMCMC.step(rng, model, sampler.sampler, state; kwargs...)
+    for _ in 2:(sampler.num_repeat)
+        transition, state = AbstractMCMC.step(rng, model, sampler.sampler, state; kwargs...)
+    end
+    return transition, state
+end

test/mcmc/gibbs.jl

@@ -251,6 +251,18 @@ end
     )
 end
 
+@testset "Equivalence of RepeatSampler and repeating Sampler" begin
+    sampler1 = Gibbs(@varname(s) => RepeatSampler(MH(), 3), @varname(m) => ESS())
+    sampler2 = Gibbs(
+        @varname(s) => MH(), @varname(s) => MH(), @varname(s) => MH(), @varname(m) => ESS()
+    )
+    Random.seed!(23)
+    chain1 = sample(gdemo_default, sampler1, 10)
+    Random.seed!(23)
+    chain2 = sample(gdemo_default, sampler1, 10)
+    @test chain1.value == chain2.value
+end
+
 @testset "Testing gibbs.jl with $adbackend" for adbackend in ADUtils.adbackends
     @testset "Deprecated Gibbs constructors" begin
         N = 10
@@ -302,7 +314,12 @@ end
             vnm = @varname(m)
             Gibbs(vns => hmc, vns => hmc, vns => hmc, vnm => pg, vnm => pg)
         end
-        for s in (s1, s2, s3, s4, s5, s6, s7, s8)
+        # Same thing but using RepeatSampler.
+        s9 = Gibbs(
+            @varname(s) => RepeatSampler(HMC(0.1, 5; adtype=adbackend), 3),
+            @varname(m) => RepeatSampler(PG(10), 2),
+        )
+        for s in (s1, s2, s3, s4, s5, s6, s7, s8, s9)
             @test DynamicPPL.alg_str(Turing.Sampler(s, gdemo_default)) == "Gibbs"
         end
 
@@ -314,6 +331,7 @@ end
         sample(gdemo_default, s6, N)
         sample(gdemo_default, s7, N)
         sample(gdemo_default, s8, N)
+        sample(gdemo_default, s9, N)
 
         g = Turing.Sampler(s3, gdemo_default)
         @test sample(gdemo_default, g, N) isa MCMCChains.Chains
@@ -355,7 +373,7 @@ end
             @varname(s) => MH(),
             (@varname(s), @varname(m)) => MH(),
             @varname(m) => ESS(),
-            @varname(s) => MH(),
+            @varname(s) => RepeatSampler(MH(), 3),
             @varname(m) => HMC(0.2, 4; adtype=adbackend),
             (@varname(m), @varname(s)) => HMC(0.2, 4; adtype=adbackend),
         )
@@ -367,7 +385,7 @@ end
             (@varname(z1), @varname(z2), @varname(z3), @varname(z4)) => PG(15),
             (@varname(z1), @varname(z2)) => PG(15),
             (@varname(mu1), @varname(mu2)) => HMC(0.15, 3; adtype=adbackend),
-            (@varname(z3), @varname(z4)) => PG(15),
+            (@varname(z3), @varname(z4)) => RepeatSampler(PG(15), 2),
             (@varname(mu1)) => ESS(),
             (@varname(mu2)) => ESS(),
             (@varname(z1), @varname(z2)) => PG(15),

test/mcmc/repeat_sampler.jl

@@ -0,0 +1,47 @@
+module HMCTests
+
+using ..Models: gdemo_default
+using ..ADUtils: ADTypeCheckContext
+using ..NumericalTests: check_gdemo, check_numerical
+import ..ADUtils
+using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample
+import DynamicPPL
+using DynamicPPL: Sampler
+import ForwardDiff
+using HypothesisTests: ApproximateTwoSampleKSTest, pvalue
+import ReverseDiff
+using LinearAlgebra: I, dot, vec
+import Random
+using StableRNGs: StableRNG
+using StatsFuns: logistic
+import Mooncake
+using Test: @test, @test_logs, @testset, @test_throws
+using Turing
+
+# RepeatedSampler only really makes sense as a component sampler of Gibbs.
+# Here we just check that running it by itself is equivalent to thinning.
+@testset "RepeatedSampler" begin
+    num_repeats = 17
+    num_samples = 10
+    num_chains = 2
+
+    rng = StableRNG(0)
+    for sampler in [MH(), Sampler(HMC(0.01, 4))]
+        chn1 = sample(
+            copy(rng),
+            gdemo_default,
+            sampler,
+            MCMCThreads(),
+            num_samples,
+            num_chains;
+            thinning=num_repeats,
+        )
+        repeat_sampler = RepeatSampler(sampler, num_repeats)
+        chn2 = sample(
+            copy(rng), gdemo_default, repeat_sampler, MCMCThreads(), num_samples, num_chains
+        )
+        @test chn1.value == chn2.value
+    end
+end
+
+end

test/runtests.jl

@@ -60,6 +60,7 @@ end
             @timeit_include("mcmc/abstractmcmc.jl")
             @timeit_include("mcmc/mh.jl")
             @timeit_include("ext/dynamichmc.jl")
+            @timeit_include("mcmc/repeat_sampler.jl")
         end
 
         @testset "variational algorithms" begin