Update to AbstractMCMC 2 and fully commit to its interface (#8)

Author: devmotion

Project.toml

@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-AbstractMCMC = "0.5, 1"
+AbstractMCMC = "2"
 ArrayInterface = "2"
 Distributions = "0.22, 0.23"
 julia = "1"

README.md

@@ -22,9 +22,10 @@ priors with non-zero means and handle the change of variables internally.
 
 ## Usage
 
-Probably most users would like to use the exported function
+Probably most users would like to generate a MC Markov chain of samples from
+the posterior distribution by calling
 ```julia
-ESS_mcmc([rng, ]prior, loglikelihood, N[; kwargs...])
+sample([rng, ]ESSModel(prior, loglikelihood), ESS(), N[; kwargs...])
 ```
 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
@@ -34,10 +35,10 @@ 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
 ```julia
-AbstractMCMC.steps!(
+AbstractMCMC.steps(
     [rng,]
-    EllipticalSliceSampling.Model(prior, loglikelihood),
-    EllipticalSliceSampling.EllipticalSliceSampler();
+    ESSModel(prior, loglikelihood),
+    ESS();
     kwargs...
 )
 ```

src/EllipticalSliceSampling.jl

@@ -7,7 +7,7 @@ import Distributions
 import Random
 import Statistics
 
-export ESS_mcmc
+export sample, ESSModel, ESS
 
 include("abstractmcmc.jl")
 include("model.jl")

src/abstractmcmc.jl

@@ -1,21 +1,19 @@
 # elliptical slice sampler
-struct EllipticalSliceSampler <: AbstractMCMC.AbstractSampler end
+struct ESS <: AbstractMCMC.AbstractSampler end
 
 # state of the elliptical slice sampler
-struct EllipticalSliceSamplerState{S,L}
+struct ESSState{S,L}
     "Sample of the elliptical slice sampler."
     sample::S
     "Log-likelihood of the sample."
     loglikelihood::L
 end
 
 # first step of the elliptical slice sampler
-function AbstractMCMC.step!(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.AbstractModel,
-    ::EllipticalSliceSampler,
-    N::Integer,
-    ::Nothing;
+    ::ESS;
     kwargs...
 )
     # initial sample from the Gaussian prior
@@ -24,16 +22,15 @@ function AbstractMCMC.step!(
     # compute log-likelihood of the initial sample
     loglikelihood = Distributions.loglikelihood(model, f)
 
-    return EllipticalSliceSamplerState(f, loglikelihood)
+    return f, ESSState(f, loglikelihood)
 end
 
 # subsequent steps of the elliptical slice sampler
-function AbstractMCMC.step!(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.AbstractModel,
-    ::EllipticalSliceSampler,
-    N::Integer,
-    state::EllipticalSliceSamplerState;
+    ::ESS,
+    state::ESSState;
     kwargs...
 )
     # sample from Gaussian prior
@@ -78,29 +75,5 @@ function AbstractMCMC.step!(
         loglikelihood = Distributions.loglikelihood(model, fnext)
     end
 
-    return EllipticalSliceSamplerState(fnext, loglikelihood)
-end
-
-# only save the samples by default
-function AbstractMCMC.transitions_init(
-    state::EllipticalSliceSamplerState,
-    model::AbstractMCMC.AbstractModel,
-    ::EllipticalSliceSampler,
-    N::Integer;
-    kwargs...
-)
-    return Vector{typeof(state.sample)}(undef, N)
-end
-
-function AbstractMCMC.transitions_save!(
-    samples::AbstractVector{S},
-    iteration::Integer,
-    state::EllipticalSliceSamplerState{S},
-    model::AbstractMCMC.AbstractModel,
-    ::EllipticalSliceSampler,
-    N::Integer;
-    kwargs...
-) where S
-    samples[iteration] = state.sample
-    return
+    return fnext, ESSState(fnext, loglikelihood)
 end

src/interface.jl

@@ -1,26 +1,3 @@
-# public interface
-
-"""
-    ESS_mcmc([rng, ]prior, loglikelihood, N; kwargs...)
-
-Create a Markov chain of `N` samples for a model with given `prior` and `loglikelihood`
-functions using the elliptical slice sampling algorithm.
-"""
-function ESS_mcmc(
-    rng::Random.AbstractRNG,
-    prior,
-    loglikelihood,
-    N::Integer;
-    kwargs...
-)
-    model = Model(prior, loglikelihood)
-    return AbstractMCMC.sample(rng, model, EllipticalSliceSampler(), N; kwargs...)
-end
-
-function ESS_mcmc(prior, loglikelihood, N::Integer; kwargs...)
-    return ESS_mcmc(Random.GLOBAL_RNG, prior, loglikelihood, N; kwargs...)
-end
-
 # private interface
 
 """

src/model.jl

@@ -1,14 +1,14 @@
 # internal model structure consisting of prior, log-likelihood function, and a cache
 
-struct Model{P,L,C} <: AbstractMCMC.AbstractModel
+struct ESSModel{P,L,C} <: AbstractMCMC.AbstractModel
     "Gaussian prior."
     prior::P
     "Log likelihood function."
     loglikelihood::L
     "Cache."
     cache::C
 
-    function Model{P,L}(prior::P, loglikelihood::L) where {P,L}
+    function ESSModel{P,L}(prior::P, loglikelihood::L) where {P,L}
         isgaussian(P) ||
             error("prior distribution has to be a Gaussian distribution")
 
@@ -19,8 +19,8 @@ struct Model{P,L,C} <: AbstractMCMC.AbstractModel
     end
 end
 
-Model(prior, loglikelihood) =
-    Model{typeof(prior),typeof(loglikelihood)}(prior, loglikelihood)
+ESSModel(prior, loglikelihood) =
+    ESSModel{typeof(prior),typeof(loglikelihood)}(prior, loglikelihood)
 
 # cache for high-dimensional samplers
 function cache(dist)
@@ -39,11 +39,11 @@ isgaussian(dist) = false
 randtype(dist) = eltype(dist)
 
 # evaluate the loglikelihood of a sample
-Distributions.loglikelihood(model::Model, f) = model.loglikelihood(f)
+Distributions.loglikelihood(model::ESSModel, f) = model.loglikelihood(f)
 
 # sample from the prior
-initial_sample(rng::Random.AbstractRNG, model::Model) = rand(rng, model.prior)
-function sample_prior(rng::Random.AbstractRNG, model::Model)
+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
@@ -55,8 +55,8 @@ function sample_prior(rng::Random.AbstractRNG, model::Model)
 end
 
 # compute the proposal
-proposal(model::Model, f, ν, θ) = proposal(model.prior, f, ν, θ)
-proposal!(out, model::Model, f, ν, θ) = proposal!(out, model.prior, f, ν, θ)
+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, ν, θ)

test/regression.jl

@@ -41,7 +41,7 @@ const σ = 0.3
         end
 
         # run elliptical slice sampler for 100 000 time steps
-        samples = ESS_mcmc(prior, ℓ, 100_000; progress = false)
+        samples = sample(ESSModel(prior, ℓ), ESS(), 100_000; progress = false)
 
         # compute analytical posterior of GP
         posterior_Σ = prior_Σ * (I - (prior_Σ + σ^2 * I) \ prior_Σ)
@@ -66,7 +66,7 @@ end
     end
 
     # run elliptical slice sampling for 100 000 time steps
-    samples = ESS_mcmc(prior, ℓ, 100_000; progress = false)
+    samples = sample(ESSModel(prior, ℓ), ESS(), 100_000; progress = false)
 
     # compute analytical posterior
     posterior_μ = observations / (1 + σ^2)

test/simple.jl

@@ -18,7 +18,7 @@ using Test
     μ = 0.8
     σ² = 0.2
 
-    samples = ESS_mcmc(prior, ℓ, 2_000; progress = false)
+    samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false)
 
     @test mean(samples) ≈ μ atol=0.05
     @test var(samples) ≈ σ² atol=0.05
@@ -35,7 +35,7 @@ end
     μ = 0.9
     σ² = 0.2
 
-    samples = ESS_mcmc(prior, ℓ, 2_000; progress = false)
+    samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false)
 
     @test mean(samples) ≈ μ atol=0.05
     @test var(samples) ≈ σ² atol=0.05
@@ -53,7 +53,7 @@ end
     μ = [0.8]
     σ² = [0.2]
 
-    samples = ESS_mcmc(prior, ℓ, 2_000; progress = false)
+    samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false)
 
     @test mean(samples) ≈ μ atol=0.05
     @test var(samples) ≈ σ² atol=0.05
@@ -70,7 +70,7 @@ end
     μ = [0.9]
     σ² = [0.2]
 
-    samples = ESS_mcmc(prior, ℓ, 2_000; progress = false)
+    samples = sample(ESSModel(prior, ℓ), ESS(), 2_000; progress = false)
 
     @test mean(samples) ≈ μ atol=0.05
     @test var(samples) ≈ σ² atol=0.05