From 766738be460edee7ac927464f5e14afca8623fc9 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Thu, 17 Dec 2020 01:05:13 +0100 Subject: [PATCH 1/4] Separate cache from model --- README.md | 6 --- src/EllipticalSliceSampling.jl | 1 - src/abstractmcmc.jl | 29 ++++++++++++--- src/distributions.jl | 54 --------------------------- src/interface.jl | 68 +++++++++++++++++++++++----------- src/model.jl | 57 +++------------------------- 6 files changed, 74 insertions(+), 141 deletions(-) delete mode 100644 src/distributions.jl diff --git a/README.md b/README.md index 2a6ee73..2c356b5 100644 --- a/README.md +++ b/README.md @@ -74,12 +74,6 @@ Statistics.mean(dist::GaussianPrior) = ... # - otherwise only `rand!(rng, dist, sample)` is required Base.rand(rng::AbstractRNG, dist::GaussianPrior) = ... Random.rand!(rng::AbstractRNG, dist::GaussianPrior, sample) = ... - -# specify the type of a sample from the distribution -Base.eltype(::Type{<:GaussianPrior}) = ... - -# in the case of mutable samples, specify the array size of the samples -Base.size(dist::GaussianPrior) = ... ``` ### Log likelihood diff --git a/src/EllipticalSliceSampling.jl b/src/EllipticalSliceSampling.jl index 093a1ef..ff4f2f4 100644 --- a/src/EllipticalSliceSampling.jl +++ b/src/EllipticalSliceSampling.jl @@ -11,7 +11,6 @@ export sample, ESSModel, ESS include("abstractmcmc.jl") include("model.jl") -include("distributions.jl") include("interface.jl") end # module diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 0ce8444..aa81e15 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -2,11 +2,19 @@ struct ESS <: AbstractMCMC.AbstractSampler end # state of the elliptical slice sampler -struct ESSState{S,L} +struct ESSState{S,L,C} "Sample of the elliptical slice sampler." sample::S "Log-likelihood of the sample." loglikelihood::L + "Cache used for in-place sampling." + cache::C +end + +function ESSState(sample, loglikelihood) + # create cache since it was not provided (initial sampling step) + cache = ArrayInterface.ismutable(sample) ? similar(sample) : nothing + return ESSState(sample, loglikelihood, cache) end # first step of the elliptical slice sampler @@ -33,8 +41,17 @@ function AbstractMCMC.step( state::ESSState; kwargs... ) + # obtain the prior + prior = EllipticalSliceSampling.prior(model) + # sample from Gaussian prior - ν = sample_prior(rng, model) + cache = state.cache + if cache === nothing + ν = Random.rand(rng, prior) + else + Random.rand!(rng, prior, cache) + ν = cache + end # sample log-likelihood threshold loglikelihood = state.loglikelihood @@ -47,7 +64,7 @@ function AbstractMCMC.step( # compute the proposal f = state.sample - fnext = proposal(model, f, ν, θ) + fnext = proposal(prior, f, ν, θ) # compute the log-likelihood of the proposal loglikelihood = Distributions.loglikelihood(model, fnext) @@ -66,14 +83,14 @@ function AbstractMCMC.step( # recompute the proposal if ArrayInterface.ismutable(fnext) - proposal!(fnext, model, f, ν, θ) + proposal!(fnext, prior, f, ν, θ) else - fnext = proposal(model, f, ν, θ) + fnext = proposal(prior, f, ν, θ) end # compute the log-likelihood of the proposal loglikelihood = Distributions.loglikelihood(model, fnext) end - return fnext, ESSState(fnext, loglikelihood) + return fnext, ESSState(fnext, loglikelihood, cache) end diff --git a/src/distributions.jl b/src/distributions.jl deleted file mode 100644 index b2c8e71..0000000 --- a/src/distributions.jl +++ /dev/null @@ -1,54 +0,0 @@ -# define the element type of the samples -randtype(::Type{D}) where {D<:Distributions.MultivariateDistribution} = Vector{eltype(D)} -randtype(::Type{D}) where {D<:Distributions.MatrixDistribution} = Matrix{eltype(D)} -function randtype( - ::Type{D} -) where {D<:Distributions.Sampleable{Distributions.Multivariate}} - return Vector{eltype(D)} -end -function randtype( - ::Type{D} -) where {D<:Distributions.Sampleable{Distributions.Matrixvariate}} - return Matrix{eltype(D)} -end - -# define trait for Gaussian distributions -isgaussian(::Type{<:Distributions.Normal}) = true -isgaussian(::Type{<:Distributions.NormalCanon}) = true -isgaussian(::Type{<:Distributions.AbstractMvNormal}) = true - -# compute the proposal of the next sample -function proposal(prior::Distributions.Normal, f::Real, ν::Real, θ) - sinθ, cosθ = sincos(θ) - μ = prior.μ - if iszero(μ) - return cosθ * f + sinθ * ν - else - a = 1 - (sinθ + cosθ) - return cosθ * f + sinθ * ν + a * μ - end -end - -function proposal( - prior::Distributions.MvNormal, - f::AbstractVector{<:Real}, - ν::AbstractVector{<:Real}, - θ -) - sinθ, cosθ = sincos(θ) - a = 1 - (sinθ + cosθ) - return @. cosθ * f + sinθ * ν + a * prior.μ -end - -function proposal!( - out::AbstractVector{<:Real}, - prior::Distributions.MvNormal, - f::AbstractVector{<:Real}, - ν::AbstractVector{<:Real}, - θ -) - sinθ, cosθ = sincos(θ) - a = 1 - (sinθ + cosθ) - @. out = cosθ * f + sinθ * ν + a * prior.μ - return out -end diff --git a/src/interface.jl b/src/interface.jl index 3809a9a..4e52459 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -1,49 +1,73 @@ # private interface """ - initial_sample(rng, model) + isgaussian(dist) -Return the initial sample for the `model` using the random number generator `rng`. +Check if distribution `dist` is a Gaussian distribution. +""" +isgaussian(dist) = false +isgaussian(::Type{<:Distributions.Normal}) = true +isgaussian(::Type{<:Distributions.NormalCanon}) = true +isgaussian(::Type{<:Distributions.AbstractMvNormal}) = true -By default, sample from the prior by calling [`sample_prior(rng, model)`](@ref). """ -function initial_sample(rng::Random.AbstractRNG, model::AbstractMCMC.AbstractModel) - return sample_prior(rng, model) -end + prior(model) + +Return the prior distribution of the `model`. +""" +function prior(::AbstractMCMC.AbstractModel) end """ - sample_prior(rng, model) + initial_sample(rng, model) -Sample from the prior of the `model` using the random number generator `rng`. +Return the initial sample for the `model` using the random number generator `rng`. + +By default, sample from [`prior(model)`](@ref). """ -function sample_prior(::Random.AbstractRNG, ::AbstractMCMC.AbstractModel) end +function initial_sample(rng::Random.AbstractRNG, model::AbstractMCMC.AbstractModel) + return Random.rand(rng, prior(model)) +end """ - proposal(model, f, ν, θ) + proposal(prior, f, ν, θ) -Compute the proposal for the next sample in the elliptical slice sampling algorithm for the -`model` from the previous sample `f`, the sample `ν` from the Gaussian prior, and the angle -`θ`. +Compute the proposal for the next sample in the elliptical slice sampling algorithm. Mathematically, the proposal can be computed as ```math -\\cos θ f + ν \\sin θ ν + μ (1 - \\sin θ + \\cos θ), +f \\cos θ + ν \\sin θ + μ (1 - (\\sin θ + \\cos θ)), ``` -where ``μ`` is the mean of the Gaussian prior. +where ``μ`` is the mean of the Gaussian `prior`, `f` is the previous sample, and `ν` is a +sample from the Gaussian `prior`. + +See also: [`proposal!`](@ref) """ -function proposal(model::AbstractMCMC.AbstractModel, f, ν, θ) end +function proposal(prior, f, ν, θ) + sinθ, cosθ = sincos(θ) + a = 1 - (sinθ + cosθ) + μ = Statistics.mean(prior) + return @. cosθ * f + sinθ * ν + a * μ +end """ proposal!(out, model, f, ν, θ) -Compute the proposal for the next sample in the elliptical slice sampling algorithm for the -`model` from the previous sample `f`, the sample `ν` from the Gaussian prior, and the angle -`θ`, and save it to `out`. +Compute the proposal for the next sample in the elliptical slice sampling algorithm, and +save it to `out`. Mathematically, the proposal can be computed as ```math -\\cos θ f + ν \\sin θ ν + μ (1 - \\sin θ + \\cos θ), +f \\cos θ + ν \\sin θ + μ (1 - (\\sin θ + \\cos θ)), ``` -where ``μ`` is the mean of the Gaussian prior. +where ``μ`` is the mean of the Gaussian `prior`, `f` is the previous sample, and `ν` is a +sample from the Gaussian `prior`. + +See also: [`proposal`](@ref) """ -function proposal!(out, model::AbstractMCMC.AbstractModel, f, ν, θ) end +function proposal!(out, prior, f, ν, θ) + sinθ, cosθ = sincos(θ) + a = 1 - (sinθ + cosθ) + μ = Statistics.mean(prior) + @. out = cosθ * f + sinθ * ν + a * μ + return out +end diff --git a/src/model.jl b/src/model.jl index 8cbefbf..32b5119 100644 --- a/src/model.jl +++ b/src/model.jl @@ -1,70 +1,23 @@ -# internal model structure consisting of prior, log-likelihood function, and a cache +# internal model structure consisting of prior and log-likelihood function -struct ESSModel{P,L,C} <: AbstractMCMC.AbstractModel +struct ESSModel{P,L} <: AbstractMCMC.AbstractModel "Gaussian prior." prior::P "Log likelihood function." loglikelihood::L - "Cache." - cache::C function ESSModel{P,L}(prior::P, loglikelihood::L) where {P,L} isgaussian(P) || error("prior distribution has to be a Gaussian distribution") - - # create cache - c = cache(prior) - - new{P,L,typeof(c)}(prior, loglikelihood, c) + new{P,L}(prior, loglikelihood) end end ESSModel(prior, loglikelihood) = ESSModel{typeof(prior),typeof(loglikelihood)}(prior, loglikelihood) -# cache for high-dimensional samplers -function cache(dist) - T = randtype(typeof(dist)) - - # only create a cache if the distribution produces mutable samples - ArrayInterface.ismutable(T) || return nothing - - similar(T, size(dist)) -end - -# test if a distribution is Gaussian -isgaussian(dist) = false - -# unify element type of samplers -randtype(dist) = eltype(dist) +# obtain prior +prior(model::ESSModel) = model.prior # evaluate the loglikelihood of a sample Distributions.loglikelihood(model::ESSModel, f) = model.loglikelihood(f) - -# sample from the prior -initial_sample(rng::Random.AbstractRNG, model::ESSModel) = rand(rng, model.prior) -function sample_prior(rng::Random.AbstractRNG, model::ESSModel) - cache = model.cache - - if cache === nothing - return rand(rng, model.prior) - else - Random.rand!(rng, model.prior, model.cache) - return model.cache - end -end - -# compute the proposal -proposal(model::ESSModel, f, ν, θ) = proposal(model.prior, f, ν, θ) -proposal!(out, model::ESSModel, f, ν, θ) = proposal!(out, model.prior, f, ν, θ) - -# default out-of-place implementation -function proposal(prior, f, ν, θ) - sinθ, cosθ = sincos(θ) - a = 1 - (sinθ + cosθ) - μ = Statistics.mean(prior) - return @. cosθ * f + sinθ * ν + a * μ -end - -# default in-place implementation -proposal!(out, prior, f, ν, θ) = copyto!(out, proposal(prior, f, ν, θ)) From f465f8f70b010e9533810bdf5c69f77363fecd19 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Thu, 17 Dec 2020 01:08:57 +0100 Subject: [PATCH 2/4] Bump version --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 4456638..287232b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "EllipticalSliceSampling" uuid = "cad2338a-1db2-11e9-3401-43bc07c9ede2" -authors = ["David Widmann "] -version = "0.3.1" +authors = ["David Widmann "] +version = "0.4.0" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" From 6a6cbe8575f9d87cfe597cce6363b5e89df685f8 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Thu, 17 Dec 2020 01:50:44 +0100 Subject: [PATCH 3/4] Export and test MCMCThreads and MCMCDistributed --- Project.toml | 4 +- README.md | 13 +++++ src/EllipticalSliceSampling.jl | 6 ++- test/regression.jl | 17 ++----- test/runtests.jl | 30 ++++++++++-- test/simple.jl | 87 +++++++++++++++++++++++++++------- 6 files changed, 120 insertions(+), 37 deletions(-) diff --git a/Project.toml b/Project.toml index 287232b..70691af 100644 --- a/Project.toml +++ b/Project.toml @@ -18,11 +18,11 @@ julia = "1" [extras] Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" +Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Distances", "FillArrays", "LinearAlgebra", "SafeTestsets", "Statistics", "Test"] +test = ["Distances", "Distributed", "FillArrays", "LinearAlgebra", "Statistics", "Test"] diff --git a/README.md b/README.md index 2c356b5..2bdfbf0 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,16 @@ which returns a vector of `N` samples for approximating the posterior of a model with a Gaussian prior that allows sampling from the `prior` and evaluation of the log likelihood `loglikelihood`. +You can sample multiple chains in parallel with multiple threads or processes +by running +```julia +sample([rng, ]ESSModel(prior, loglikelihood), ESS(), MCMCThreads(), N, nchains[; kwargs...]) +``` +or +```julia +sample([rng, ]ESSModel(prior, loglikelihood), ESS(), MCMCDistributed(), N, nchains[; kwargs...]) +``` + If you want to have more control about the sampling procedure (e.g., if you only want to save a subset of samples or want to use another stopping criterion), the function @@ -44,6 +54,9 @@ AbstractMCMC.steps( gives you access to an iterator from which you can generate an unlimited number of samples. +For more details regarding `sample` and `steps` please check the documentation of +[AbstractMCMC.jl](https://github.com/TuringLang/AbstractMCMC.jl). + ### Prior You may specify Gaussian priors with arbitrary means. EllipticalSliceSampling.jl diff --git a/src/EllipticalSliceSampling.jl b/src/EllipticalSliceSampling.jl index ff4f2f4..e2c48c8 100644 --- a/src/EllipticalSliceSampling.jl +++ b/src/EllipticalSliceSampling.jl @@ -7,7 +7,11 @@ import Distributions import Random import Statistics -export sample, ESSModel, ESS +export ESSModel, ESS + +# reexports +using AbstractMCMC: sample, MCMCThreads, MCMCDistributed +export sample, MCMCThreads, MCMCDistributed include("abstractmcmc.jl") include("model.jl") diff --git a/test/regression.jl b/test/regression.jl index 4286ace..0930294 100644 --- a/test/regression.jl +++ b/test/regression.jl @@ -1,19 +1,9 @@ -using EllipticalSliceSampling - -using Distances -using Distributions - -using LinearAlgebra -using Random -using Statistics -using Test - -Random.seed!(0) +@testset "regression.jl" begin # number of input features -const N = 200 +N = 200 # homoscedastic observation noise -const σ = 0.3 +σ = 0.3 # experiment of Murray et al. @testset "GP regression" begin @@ -74,3 +64,4 @@ end # compare with empirical estimates @test mean(samples) ≈ posterior_μ rtol = 0.02 end +end diff --git a/test/runtests.jl b/test/runtests.jl index dcb7172..a9b6b95 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,28 @@ -using SafeTestsets +using EllipticalSliceSampling -@safetestset "Simple tests" begin include("simple.jl") end -@safetestset "GP regression tests" begin include("regression.jl") end +using Distances +using Distributions +using FillArrays + +using Distributed +using LinearAlgebra +using Random +using Statistics +using Test + +Random.seed!(0) + +# add additional processes +addprocs(2) +@everywhere begin + using EllipticalSliceSampling + using Distributions +end + +@testset "EllipticalSliceSampling" begin + println("Simple tests") + @testset "Simple tests" begin include("simple.jl") end + + println("GP regression tests") + @testset "GP regression tests" begin include("regression.jl") end +end diff --git a/test/simple.jl b/test/simple.jl index 8ee12e7..3b90e88 100644 --- a/test/simple.jl +++ b/test/simple.jl @@ -1,27 +1,42 @@ -using EllipticalSliceSampling +@testset "simple.jl" begin -using Distributions -using FillArrays - -using Random -using Statistics -using Test +# Load all required packages and define likelihood functions +ℓ(x) = logpdf(Normal(x, 0.5), 1.0) +ℓvec(x) = logpdf(MvNormal(x, 0.5), [1.0]) +@everywhere begin + ℓ(x) = logpdf(Normal(x, 0.5), 1.0) + ℓvec(x) = logpdf(MvNormal(x, 0.5), [1.0]) +end @testset "Scalar model" begin Random.seed!(1) # model prior = Normal(0, 1) - ℓ(x) = logpdf(Normal(x, 0.5), 1.0) # true posterior μ = 0.8 σ² = 0.2 - samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false) - + # regular sampling + model = let prior=prior, ℓ=ℓ + ESSModel(prior, ℓ) + end + samples = sample(model, ESS(), 2_000; progress = false) + @test samples isa Vector{Float64} + @test length(samples) == 2_000 @test mean(samples) ≈ μ atol=0.05 @test var(samples) ≈ σ² atol=0.05 + + # parallel sampling + for alg in (MCMCThreads(), MCMCDistributed()) + samples = sample(model, ESS(), alg, 2_000, 5; progress = false) + @test samples isa Vector{Vector{Float64}} + @test length(samples) == 5 + @test all(length(x) == 2_000 for x in samples) + @test mean(mean, samples) ≈ μ atol=0.05 + @test mean(var, samples) ≈ σ² atol=0.05 + end end @testset "Scalar model with nonzero mean" begin @@ -29,16 +44,27 @@ end # model prior = Normal(0.5, 1) - ℓ(x) = logpdf(Normal(x, 0.5), 1.0) # true posterior μ = 0.9 σ² = 0.2 + # regular sampling samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false) - + @test samples isa Vector{Float64} + @test length(samples) == 2_000 @test mean(samples) ≈ μ atol=0.05 @test var(samples) ≈ σ² atol=0.05 + + # parallel sampling + for alg in (MCMCThreads(), MCMCDistributed()) + samples = sample(ESSModel(prior, ℓ), ESS(), alg, 2_000, 5; progress = false) + @test samples isa Vector{Vector{Float64}} + @test length(samples) == 5 + @test all(length(x) == 2_000 for x in samples) + @test mean(mean, samples) ≈ μ atol=0.05 + @test mean(var, samples) ≈ σ² atol=0.05 + end end @@ -47,16 +73,28 @@ end # model prior = MvNormal([0.0], 1.0) - ℓ(x) = logpdf(MvNormal(x, 0.5), [1.0]) # true posterior μ = [0.8] σ² = [0.2] - samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false) - + # regular sampling + samples = sample(ESSModel(prior, ℓvec), ESS(), 2_000; progress = false) + @test samples isa Vector{Vector{Float64}} + @test length(samples) == 2_000 + @test all(length(x) == 1 for x in samples) @test mean(samples) ≈ μ atol=0.05 @test var(samples) ≈ σ² atol=0.05 + + # parallel sampling + for alg in (MCMCThreads(), MCMCDistributed()) + samples = sample(ESSModel(prior, ℓvec), ESS(), alg, 2_000, 5; progress = false) + @test samples isa Vector{Vector{Vector{Float64}}} + @test length(samples) == 5 + @test all(length(x) == 2_000 for x in samples) + @test mean(mean, samples) ≈ μ atol=0.05 + @test mean(var, samples) ≈ σ² atol=0.05 + end end @testset "Scalar model with nonzero mean (vectorized)" begin @@ -64,14 +102,27 @@ end # model prior = MvNormal([0.5], 1.0) - ℓ(x) = logpdf(MvNormal(x, 0.5), [1.0]) # true posterior μ = [0.9] σ² = [0.2] - samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false) - + # regular sampling + samples = sample(ESSModel(prior, ℓvec), ESS(), 2_000; progress = false) + @test samples isa Vector{Vector{Float64}} + @test length(samples) == 2_000 + @test all(length(x) == 1 for x in samples) @test mean(samples) ≈ μ atol=0.05 @test var(samples) ≈ σ² atol=0.05 + + # parallel sampling + for alg in (MCMCThreads(), MCMCDistributed()) + samples = sample(ESSModel(prior, ℓvec), ESS(), alg, 2_000, 5; progress = false) + @test samples isa Vector{Vector{Vector{Float64}}} + @test length(samples) == 5 + @test all(length(x) == 2_000 for x in samples) + @test mean(mean, samples) ≈ μ atol=0.05 + @test mean(var, samples) ≈ σ² atol=0.05 + end +end end From 2a1ec680982af89f142c6bd3f170514954deae5f Mon Sep 17 00:00:00 2001 From: David Widmann Date: Thu, 17 Dec 2020 02:09:11 +0100 Subject: [PATCH 4/4] Use 2 threads in CI tests --- .github/workflows/CI.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 3ca2401..8eb6efd 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -50,6 +50,8 @@ jobs: ${{ runner.os }}- - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 + env: + JULIA_NUM_THREADS: 2 - uses: julia-actions/julia-processcoverage@v1 if: matrix.coverage - uses: codecov/codecov-action@v1