Separate cache from model (#13)

Author: devmotion

.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

Project.toml

@@ -1,7 +1,7 @@
 name = "EllipticalSliceSampling"
 uuid = "cad2338a-1db2-11e9-3401-43bc07c9ede2"
-authors = ["David Widmann <dev+git@devmotion.de>"]
-version = "0.3.1"
+authors = ["David Widmann <david.widmann@it.uu.se>"]
+version = "0.4.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -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"]

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
@@ -74,12 +87,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

src/EllipticalSliceSampling.jl

@@ -7,11 +7,14 @@ 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")
-include("distributions.jl")
 include("interface.jl")
 
 end # module

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

src/distributions.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

src/interface.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, ν, θ))