Support Turing's externalsampler interface, return Transition instead of raw samples.

.github/workflows/CI.yml

@@ -19,8 +19,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
-          - '1.9'
+          - '1.7'
+          - '1.10'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "SliceSampling"
 uuid = "43f4d3e8-9711-4a8c-bd1b-03ac73a255cf"
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -9,20 +9,30 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a"
 
+[weakdeps]
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
+
+[extensions]
+SliceSamplingTuringExt = ["Turing"]
+
 [compat]
-AbstractMCMC = "5"
+AbstractMCMC = "4, 5"
 Accessors = "0.1"
 Distributions = "0.25"
 FillArrays = "1"
 LogDensityProblems = "2"
 Random = "1"
+Requires = "1"
 SimpleUnPack = "1"
-julia = "1.6"
+Turing = "0.31, 0.32"
+julia = "1.7"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [targets]
 test = ["Test"]

docs/Project.toml

@@ -11,3 +11,4 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SliceSampling = "43f4d3e8-9711-4a8c-bd1b-03ac73a255cf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"

docs/make.jl

@@ -17,7 +17,6 @@ makedocs(;
     pages=[
         "Home"                      => "index.md",
         "General Usage"             => "general.md",
-        #"Benchmarks"                => "benchmarks.md",
         "Univariate Slice Sampling" => "univariate_slice.md",
         "Latent Slice Sampling"     => "latent_slice.md"
     ],

docs/src/general.md

@@ -4,6 +4,26 @@
 This package implements the `AbstractMCMC` [interface](https://github.com/TuringLang/AbstractMCMC.jl).
 `AbstractMCMC` provides a unifying interface for MCMC algorithms applied to [LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl).
 
+## Drawing Samples From `Turing` Models
+`SliceSampling.jl` can be used to sample from [Turing](https://github.com/TuringLang/Turing.jl) models through `Turing`'s `externalsampler` interface.
+See the following example of using the [latent slice sampler](@ref latent):
+
+```@example turing
+using Distributions
+using Turing
+using SliceSampling
+
+@model function demo()
+    s ~ InverseGamma(3, 3)
+    m ~ Normal(0, sqrt(s))
+end
+
+sampler   = LatentSlice(2)
+n_samples = 10000
+model     = demo()
+sample(model, externalsampler(sampler), n_samples; initial_params=[1.0, 0.0])
+```
+
 ## Drawing Samples
 
 For drawing samples using the algorithms provided by `SliceSampling`, the user only needs to call:
@@ -14,6 +34,12 @@ sample([rng,] model, slice, N; initial_params)
 - `model`: A model implementing the `LogDensityProblems` interface.
 - `N`: The number of samples
 
+The output is a `SliceSampling.Transition` object, which contains the following:
+```@docs
+SliceSampling.Transition
+```
+
+
 For the keyword arguments, `SliceSampling` allows:
 - `initial_params`: The intial state of the Markov chain (default: `nothing`).
 

docs/src/latent_slice.md

@@ -1,5 +1,5 @@
 
-# Latent Slice Sampling 
+# [Latent Slice Sampling](@id latent)
 
 ## Introduction
 Latent slice sampling is a recent vector-valued slice sampling algorithm proposed by Li and Walker[^LW2023].

ext/SliceSamplingTuringExt.jl

@@ -0,0 +1,14 @@
+
+module SliceSamplingTuringExt
+
+if isdefined(Base, :get_extension)
+    using SliceSampling: Transition
+    using Turing: Turing
+else
+    using ..SliceSampling: Transition
+    using ..Turing: Turing
+end
+
+Turing.Inference.getparams(::Turing.DynamicPPL.Model, sample::Transition) = sample.params
+
+end

src/SliceSampling.jl

@@ -16,6 +16,27 @@ export sample, MCMCThreads, MCMCDistributed, MCMCSerial
 # Interfaces
 abstract type AbstractSliceSampling <: AbstractMCMC.AbstractSampler end
 
+"""
+    struct Transition
+
+Struct containing the results of the transition.
+
+# Fields
+- `params`: Samples generated by the transition.
+- `lp::Real`: Log-target density of the samples.
+- `info::NamedTuple`: Named tuple containing information about the transition. 
+"""
+struct Transition{P, L <: Real, I <: NamedTuple}
+    "current state of the slice sampling chain"
+    params::P
+
+    "log density of the current state"
+    lp::L
+
+    "information generated from the sampler"
+    info::I
+end
+
 """
     initial_sample(rng, model)
 
@@ -46,4 +67,18 @@ export LatentSlice
 
 include("latent.jl")
 
+# Turing Compatibility
+
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
+@static if !isdefined(Base, :get_extension)
+    function __init__()
+        @require Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" include(
+            "../ext/SliceSamplingTuringExt.jl"
+        )
+    end
+end
+
 end

src/gibbs.jl

@@ -1,4 +1,9 @@
 
+struct GibbsSliceState{T <: Transition}
+    "Current [`Transition`](@ref)."
+    transition::T
+end
+
 struct GibbsObjective{Model, Idx <: Integer, Vec <: AbstractVector}
     model::Model
     idx  ::Idx
@@ -10,17 +15,6 @@ function LogDensityProblems.logdensity(gibbs::GibbsObjective, θi)
     LogDensityProblems.logdensity(model, (@set θ[idx] = θi))
 end
 
-struct SliceState{P, L <: Real, I <: NamedTuple}
-    "current state of the slice sampling chain"
-    params::P
-
-    "log density of the current state"
-    lp::L
-
-    "information generated from the sampler"
-    info::I
-end
-
 function AbstractMCMC.step(rng    ::Random.AbstractRNG,
                            model  ::AbstractMCMC.LogDensityModel,
                            sampler::AbstractGibbsSliceSampling;
@@ -33,14 +27,15 @@ function AbstractMCMC.step(rng    ::Random.AbstractRNG,
     end
     θ  = isnothing(initial_params) ? initial_sample(rng, logdensitymodel) : initial_params
     lp = LogDensityProblems.logdensity(logdensitymodel, θ)
-    return θ, SliceState(θ, lp, NamedTuple())
+    t  = Transition(θ, lp, NamedTuple())
+    return t, GibbsSliceState(t)
 end
 
 function AbstractMCMC.step(
     rng    ::Random.AbstractRNG,
     model  ::AbstractMCMC.LogDensityModel, 
     sampler::AbstractGibbsSliceSampling,
-    state  ::SliceState;
+    state  ::GibbsSliceState;
     kwargs...,
 )
     logdensitymodel = model.logdensity
@@ -49,8 +44,8 @@ function AbstractMCMC.step(
     else
         sampler.window
     end
-    ℓp = state.lp
-    θ  = copy(state.params)
+    ℓp = state.transition.lp
+    θ  = copy(state.transition.params)
 
     total_props = 0
     for idx in shuffle(rng, 1:length(θ))
@@ -61,6 +56,6 @@ function AbstractMCMC.step(
         total_props += props
         θ[idx] = θ′idx
     end
-
-    θ, SliceState(θ, ℓp, (num_proposals=total_props,))
+    t = Transition(θ, ℓp, (num_proposals=total_props,))
+    t, GibbsSliceState(t)
 end

src/latent.jl

@@ -11,11 +11,12 @@ struct LatentSlice{B <: Real} <: AbstractSliceSampling
     beta::B
 end
 
-struct LatentSliceState{V <: AbstractVector, L <: Real, I <: NamedTuple}
-    y   ::V
-    s   ::V
-    lp  ::L
-    info::I
+struct LatentSliceState{T <: Transition, S <: AbstractVector}
+    "Current [`Transition`](@ref)."
+    transition ::T
+
+    "Auxiliary variables for adapting the slice window (\$s\$ in the original paper[^LW2023])"
+    sliceparams::S
 end
 
 function AbstractMCMC.step(rng    ::Random.AbstractRNG,
@@ -29,7 +30,8 @@ function AbstractMCMC.step(rng    ::Random.AbstractRNG,
     d  = length(y)
     lp = LogDensityProblems.logdensity(logdensitymodel, y)
     s  = convert(Vector{eltype(y)}, rand(rng, Gamma(2, 1/β), d))
-    return y, LatentSliceState(y, s, lp, NamedTuple())
+    t  = Transition(y, lp, NamedTuple())
+    return t, LatentSliceState(t, s)
 end
 
 function AbstractMCMC.step(
@@ -42,9 +44,9 @@ function AbstractMCMC.step(
     logdensitymodel = model.logdensity
 
     β  = sampler.beta
-    ℓp = state.lp
-    y  = copy(state.y)
-    s  = copy(state.s)
+    ℓp = state.transition.lp
+    y  = state.transition.params
+    s  = state.sliceparams
     d  = length(y)
     ℓw = ℓp - Random.randexp(rng, eltype(y))
 
@@ -76,5 +78,6 @@ function AbstractMCMC.step(
         end
     end
     s = β*randexp(rng, eltype(y), d) + 2*abs.(l - y)
-    y, LatentSliceState(y, s, ℓp, NamedTuple())
+    t = Transition(y, ℓp, NamedTuple())
+    t, LatentSliceState(t, s)
 end

test/Project.toml

@@ -7,14 +7,16 @@ MCMCTesting = "9963b6a1-5d46-439c-8efc-3a487843c7fa"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
-AbstractMCMC = "5"
+AbstractMCMC = "4, 5"
 Accessors = "0.1"
 Distributions = "0.25"
 LogDensityProblems = "2"
 MCMCTesting = "0.3"
 Random = "1"
 StableRNGs = "1"
 Test = "1"
+Turing = "0.31"
 julia = "1.6"

test/runtests.jl

@@ -6,6 +6,7 @@ using LogDensityProblems
 using MCMCTesting
 using Random 
 using Test
+using Turing
 using StableRNGs
 
 using SliceSampling
@@ -37,8 +38,8 @@ function MCMCTesting.markovchain_transition(
 )
     model′ = AbstractMCMC.LogDensityModel(@set model.y = y)
     _, init_state = AbstractMCMC.step(rng, model′, sampler; initial_params=copy(θ))
-    θ′, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
-    θ′
+    transition, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
+    transition.params
 end
 
 function LogDensityProblems.logdensity(model::Model{F, V}, θ) where {F <: Real, V}
@@ -91,11 +92,13 @@ end
 
             rng           = StableRNG(1)
             _, init_state = AbstractMCMC.step(rng, model′, sampler; initial_params=copy(θ))
-            θ′, _         = AbstractMCMC.step(rng, model′, sampler, init_state)
+            transition, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
+            θ′            = transition.params
 
             rng           = StableRNG(1)
             _, init_state = AbstractMCMC.step(rng, model′, sampler; initial_params=copy(θ))
-            θ′′, _         = AbstractMCMC.step(rng, model′, sampler, init_state)
+            transition, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
+            θ′′            = transition.params
             @test θ′ == θ′′
         end
 
@@ -109,7 +112,8 @@ end
             @test eltype(y) == type
 
             _, init_state = AbstractMCMC.step(rng, model′, sampler; initial_params=copy(θ))
-            θ′, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
+            transition, _ = AbstractMCMC.step(rng, model′, sampler, init_state)
+            θ′             = transition.params
 
             @test eltype(θ′) == type
         end
@@ -127,3 +131,30 @@ end
         end
     end
 end
+
+@testset "turing compatibility" begin
+    @model function demo()
+        s ~ InverseGamma(2, 3)
+        m ~ Normal(0, sqrt(s))
+        1.5 ~ Normal(m, sqrt(s))
+        2.0 ~ Normal(m, sqrt(s))
+    end
+
+    n_samples = 1000
+    model     = demo()
+
+    @testset for sampler in [
+        Slice(1),
+        SliceSteppingOut(1),
+        SliceDoublingOut(1),
+        LatentSlice(5),
+    ]
+        chain = sample(
+            model,
+            externalsampler(sampler),
+            n_samples;
+            initial_params=[1.0, 0.0],
+            progress=false,
+        )
+    end
+end