Add MCMCSerial

Author: cpfiffer

src/AbstractMCMC.jl

@@ -58,19 +58,27 @@ abstract type AbstractMCMCParallel end
 """
     MCMCThreads
 
-The `MCMCThreads` algorithm allows to sample MCMC chains in parallel using multiple
+The `MCMCThreads` algorithm allows users to sample MCMC chains in parallel using multiple
 threads.
 """
 struct MCMCThreads <: AbstractMCMCParallel end
 
 """
     MCMCDistributed
 
-The `MCMCDistributed` algorithm allows to sample MCMC chains in parallel using multiple
+The `MCMCDistributed` algorithm allows users to sample MCMC chains in parallel using multiple
 processes.
 """
 struct MCMCDistributed <: AbstractMCMCParallel end
 
+
+"""
+    MCMCSerial
+
+The `MCMCSerial` algorithm allows users to sample serially, with no thread or process parallelism.
+"""
+struct MCMCSerial <: AbstractMCMCParallel end
+
 include("samplingstats.jl")
 include("logging.jl")
 include("interface.jl")

src/sample.jl

@@ -444,5 +444,35 @@ function mcmcsample(
     return chainsstack(tighten_eltype(chains))
 end
 
+function mcmcsample(
+    rng::Random.AbstractRNG,
+    model::AbstractModel,
+    sampler::AbstractSampler,
+    ::MCMCSerial,
+    N::Integer,
+    nchains::Integer;
+    progressname = "Sampling",
+    kwargs...
+)
+    # Check if the number of chains is larger than the number of samples
+    if nchains > N
+        @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
+    end
+
+    # Set up a chains vector.
+    chains = Vector{Any}(undef, nchains)
+
+    # Sample each chain
+    for i in 1:nchains
+        # Sample a chain and save it to the vector.
+        chains[i] = StatsBase.sample(rng, model, sampler, N; 
+                                     progressname = string(progressname, " (Chain $i of $nchains)"),
+                                     kwargs...)
+    end
+
+    # Concatenate the chains together.
+    return chainsstack(tighten_eltype(chains))
+end
+
 tighten_eltype(x) = x
 tighten_eltype(x::Vector{Any}) = map(identity, x)

test/sample.jl

@@ -225,6 +225,52 @@
         @test all(l.level > Logging.LogLevel(-1) for l in logs)
     end
 
+    @testset "Serial sampling" begin
+        # No dedicated chains type
+        N = 10_000
+        chains = sample(MyModel(), MySampler(), MCMCSerial(), N, 1000)
+        @test chains isa Vector{<:Vector{<:MySample}}
+        @test length(chains) == 1000
+        @test all(length(x) == N for x in chains)
+
+        Random.seed!(1234)
+        chains = sample(MyModel(), MySampler(), MCMCSerial(), N, 1000;
+                        chain_type = MyChain)
+
+        # Test output type and size.
+        @test chains isa Vector{<:MyChain}
+        @test all(c.as[1] === missing for c in chains)
+        @test length(chains) == 1000
+        @test all(x -> length(x.as) == length(x.bs) == N, chains)
+
+        # Test some statistical properties.
+        @test all(x -> isapprox(mean(@view x.as[2:end]), 0.5; atol=5e-2), chains)
+        @test all(x -> isapprox(var(@view x.as[2:end]), 1 / 12; atol=5e-3), chains)
+        @test all(x -> isapprox(mean(@view x.bs[2:end]), 0; atol=5e-2), chains)
+        @test all(x -> isapprox(var(@view x.bs[2:end]), 1; atol=5e-2), chains)
+
+        # Test reproducibility.
+        Random.seed!(1234)
+        chains2 = sample(MyModel(), MySampler(), MCMCSerial(), N, 1000;
+                         chain_type = MyChain)
+
+        @test all(c1.as[i] === c2.as[i] for (c1, c2) in zip(chains, chains2), i in 1:N)
+        @test all(c1.bs[i] === c2.bs[i] for (c1, c2) in zip(chains, chains2), i in 1:N)
+
+        # Unexpected order of arguments.
+        str = "Number of chains (10) is greater than number of samples per chain (5)"
+        @test_logs (:warn, str) match_mode=:any sample(MyModel(), MySampler(),
+                                                       MCMCSerial(), 5, 10;
+                                                       chain_type = MyChain)
+
+        # Suppress output.
+        logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
+            sample(MyModel(), MySampler(), MCMCSerial(), 10_000, 100;
+                   progress = false, chain_type = MyChain)
+        end
+        @test all(l.level > Logging.LogLevel(-1) for l in logs)
+    end
+
     @testset "Chain constructors" begin
         chain1 = sample(MyModel(), MySampler(), 100; sleepy = true)
         chain2 = sample(MyModel(), MySampler(), 100; sleepy = true, chain_type = MyChain)