Merge pull request #71 from TuringLang/csp/serial

Add MCMCSerial

Author: cpfiffer

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probablistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "3.1.0"
+version = "3.2.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

docs/src/api.md

@@ -28,17 +28,17 @@ AbstractMCMC.sample(
     ::AbstractRNG,
     ::AbstractMCMC.AbstractModel,
     ::AbstractMCMC.AbstractSampler,
-    ::AbstractMCMC.AbstractMCMCParallel,
+    ::AbstractMCMC.AbstractMCMCEnsemble,
     ::Integer,
     ::Integer,
 )
 ```
 
-Two algorithms are provided for parallel sampling with multiple threads and multiple processes,
-respectively:
+Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):
 ```@docs
 AbstractMCMC.MCMCThreads
 AbstractMCMC.MCMCDistributed
+AbstractMCMC.MCMCSerial
 ```
 
 ## Common keyword arguments

src/AbstractMCMC.jl

@@ -17,7 +17,7 @@ using StatsBase: sample
 export sample
 
 # Parallel sampling types
-export MCMCThreads, MCMCDistributed
+export MCMCThreads, MCMCDistributed, MCMCSerial
 
 """
     AbstractChains
@@ -48,34 +48,43 @@ An `AbstractModel` represents a generic model type that can be used to perform i
 abstract type AbstractModel end
 
 """
-    AbstractMCMCParallel
+    AbstractMCMCEnsemble
 
-An `AbstractMCMCParallel` algorithm represents a specific algorithm for sampling MCMC chains
+An `AbstractMCMCEnsemble` algorithm represents a specific algorithm for sampling MCMC chains
 in parallel.
 """
-abstract type AbstractMCMCParallel end
+abstract type AbstractMCMCEnsemble 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
+struct MCMCThreads <: AbstractMCMCEnsemble 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
+struct MCMCDistributed <: AbstractMCMCEnsemble end
+
+
+"""
+    MCMCSerial
+
+The `MCMCSerial` algorithm allows users to sample serially, with no thread or process parallelism.
+"""
+struct MCMCSerial <: AbstractMCMCEnsemble end
 
 include("samplingstats.jl")
 include("logging.jl")
 include("interface.jl")
 include("sample.jl")
 include("stepper.jl")
 include("transducer.jl")
+include("deprecations.jl")
 
 end # module AbstractMCMC

src/deprecations.jl

@@ -0,0 +1,2 @@
+# Deprecate the old name AbstractMCMCParallel in favor of AbstractMCMCEnsemble
+Base.@deprecate_binding AbstractMCMCParallel AbstractMCMCEnsemble false

src/sample.jl

@@ -61,7 +61,7 @@ end
 function StatsBase.sample(
     model::AbstractModel,
     sampler::AbstractSampler,
-    parallel::AbstractMCMCParallel,
+    parallel::AbstractMCMCEnsemble,
     N::Integer,
     nchains::Integer;
     kwargs...
@@ -80,7 +80,7 @@ function StatsBase.sample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
     sampler::AbstractSampler,
-    parallel::AbstractMCMCParallel,
+    parallel::AbstractMCMCEnsemble,
     N::Integer,
     nchains::Integer;
     kwargs...
@@ -444,5 +444,31 @@ 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
+
+    # Sample the chains.
+    chains = map(
+        i -> StatsBase.sample(rng, model, sampler, N; progressname = string(progressname, " (Chain ", i, " of ", nchains, ")"),
+        kwargs...),
+        1:nchains
+    )
+
+    # 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)