Update ESS to use LDF

Author: penelopeysm

src/mcmc/abstractmcmc.jl

@@ -21,8 +21,7 @@
 
 # Because this is a pain to implement all at once, we do it for one sampler at a time.
 # This type tells us which samplers have been 'updated' to the new interface.
-
-const LDFCompatibleSampler = Union{Hamiltonian}
+const LDFCompatibleSampler = Union{Hamiltonian,ESS}
 
 """
     sample(

src/mcmc/ess.jl

@@ -20,22 +20,26 @@ Mean
 │ 1   │ m          │ 0.824853 │
 ```
 """
-struct ESS <: InferenceAlgorithm end
+struct ESS <: AbstractSampler end
+
+DynamicPPL.initialsampler(::ESS) = DynamicPPL.SampleFromPrior()
+update_sample_kwargs(::ESS, ::Integer, kwargs) = kwargs
+get_adtype(::ESS) = nothing
+requires_unconstrained_space(::ESS) = false
 
 # always accept in the first step
-function DynamicPPL.initialstep(
-    rng::AbstractRNG, model::Model, spl::Sampler{<:ESS}, vi::AbstractVarInfo; kwargs...
-)
+function AbstractMCMC.step(rng::AbstractRNG, ldf::LogDensityFunction, spl::ESS; kwargs...)
+    vi = ldf.varinfo
     for vn in keys(vi)
         dist = getdist(vi, vn)
         EllipticalSliceSampling.isgaussian(typeof(dist)) ||
             error("ESS only supports Gaussian prior distributions")
     end
-    return Transition(model, vi), vi
+    return Transition(ldf.model, vi), vi
 end
 
 function AbstractMCMC.step(
-    rng::AbstractRNG, model::Model, spl::Sampler{<:ESS}, vi::AbstractVarInfo; kwargs...
+    rng::AbstractRNG, ldf::LogDensityFunction, spl::ESS, vi::AbstractVarInfo; kwargs...
 )
     # obtain previous sample
     f = vi[:]
@@ -45,14 +49,13 @@ function AbstractMCMC.step(
     oldstate = EllipticalSliceSampling.ESSState(f, getlogp(vi), nothing)
 
     # compute next state
+    # Note: `f_loglikelihood` effectively calculates the log-likelihood (not
+    # log-joint, despite the use of `LDP.logdensity`) because `tilde_assume` is
+    # overloaded on `SamplingContext(rng, ESS(), ...)` below.
+    f_loglikelihood = Base.Fix1(LogDensityProblems.logdensity, ldf)
     sample, state = AbstractMCMC.step(
         rng,
-        EllipticalSliceSampling.ESSModel(
-            ESSPrior(model, spl, vi),
-            DynamicPPL.LogDensityFunction(
-                model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
-            ),
-        ),
+        EllipticalSliceSampling.ESSModel(ESSPrior(ldf.model, spl, vi), f_loglikelihood),
         EllipticalSliceSampling.ESS(),
         oldstate,
     )
@@ -61,67 +64,57 @@ function AbstractMCMC.step(
     vi = DynamicPPL.unflatten(vi, sample)
     vi = setlogp!!(vi, state.loglikelihood)
 
-    return Transition(model, vi), vi
+    return Transition(ldf.model, vi), vi
 end
 
 # Prior distribution of considered random variable
-struct ESSPrior{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo,T}
+struct ESSPrior{M<:Model,V<:AbstractVarInfo,T}
     model::M
-    sampler::S
+    sampler::ESS
     varinfo::V
     μ::T
 
-    function ESSPrior{M,S,V}(
-        model::M, sampler::S, varinfo::V
-    ) where {M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo}
+    function ESSPrior{M,V}(
+        model::M, sampler::ESS, varinfo::V
+    ) where {M<:Model,V<:AbstractVarInfo}
         vns = keys(varinfo)
         μ = mapreduce(vcat, vns) do vn
             dist = getdist(varinfo, vn)
             EllipticalSliceSampling.isgaussian(typeof(dist)) ||
                 error("[ESS] only supports Gaussian prior distributions")
             DynamicPPL.tovec(mean(dist))
         end
-        return new{M,S,V,typeof(μ)}(model, sampler, varinfo, μ)
+        return new{M,V,typeof(μ)}(model, sampler, varinfo, μ)
     end
 end
 
-function ESSPrior(model::Model, sampler::Sampler{<:ESS}, varinfo::AbstractVarInfo)
-    return ESSPrior{typeof(model),typeof(sampler),typeof(varinfo)}(model, sampler, varinfo)
+function ESSPrior(model::Model, sampler::ESS, varinfo::AbstractVarInfo)
+    return ESSPrior{typeof(model),typeof(varinfo)}(model, sampler, varinfo)
 end
 
 # Ensure that the prior is a Gaussian distribution (checked in the constructor)
 EllipticalSliceSampling.isgaussian(::Type{<:ESSPrior}) = true
 
 # Only define out-of-place sampling
 function Base.rand(rng::Random.AbstractRNG, p::ESSPrior)
-    sampler = p.sampler
-    varinfo = p.varinfo
-    # TODO: Surely there's a better way of doing this now that we have `SamplingContext`?
-    vns = keys(varinfo)
-    for vn in vns
-        set_flag!(varinfo, vn, "del")
+    # TODO(penelopeysm): This is ugly -- need to set 'del' flag because
+    # otherwise DynamicPPL.SampleWithPrior will just use the existing
+    # parameters in the varinfo. In general SampleWithPrior etc. need to be
+    # reworked.
+    for vn in keys(p.varinfo)
+        set_flag!(p.varinfo, vn, "del")
     end
-    p.model(rng, varinfo, sampler)
-    return varinfo[:]
+    _, vi = DynamicPPL.evaluate!!(p.model, p.varinfo, SamplingContext(rng, p.sampler))
+    return vi[:]
 end
 
 # Mean of prior distribution
 Distributions.mean(p::ESSPrior) = p.μ
 
-# Evaluate log-likelihood of proposals
-const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} =
-    DynamicPPL.LogDensityFunction{M,V,<:DynamicPPL.SamplingContext{<:S},AD} where {AD}
-
-(ℓ::ESSLogLikelihood)(f::AbstractVector) = LogDensityProblems.logdensity(ℓ, f)
-
 function DynamicPPL.tilde_assume(
-    rng::Random.AbstractRNG, ::DefaultContext, ::Sampler{<:ESS}, right, vn, vi
+    rng::Random.AbstractRNG, ::DefaultContext, ::ESS, right, vn, vi
 )
     return DynamicPPL.tilde_assume(
         rng, LikelihoodContext(), SampleFromPrior(), right, vn, vi
     )
 end
-
-function DynamicPPL.tilde_observe(ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vi)
-    return DynamicPPL.tilde_observe(ctx, SampleFromPrior(), right, left, vi)
-end

test/mcmc/ess.jl

@@ -13,6 +13,41 @@ using Turing
 @testset "ESS" begin
     @info "Starting ESS tests"
 
+    @testset "InferenceAlgorithm interface" begin
+        alg = ESS()
+        @test Turing.Inference.get_adtype(alg) === nothing
+        @test !Turing.Inference.requires_unconstrained_space(alg)
+        kwargs = (; _foo="bar")
+        @test Turing.Inference.update_sample_kwargs(alg, 1000, kwargs) == kwargs
+    end
+
+    @testset "sample() interface" begin
+        @model function demo_normal(x)
+            a ~ Normal()
+            return x ~ Normal(a)
+        end
+        model = demo_normal(2.0)
+        ldf = LogDensityFunction(model)
+        sampling_objects = Dict("DynamicPPL.Model" => model, "LogDensityFunction" => ldf)
+        seed = 468
+
+        @testset "sampling with $name" for (name, model_or_ldf) in sampling_objects
+            spl = ESS()
+            # check sampling works without rng
+            @test sample(model_or_ldf, spl, 5) isa Chains
+            # check reproducibility with rng
+            chn1 = sample(Random.Xoshiro(seed), model_or_ldf, spl, 5)
+            chn2 = sample(Random.Xoshiro(seed), model_or_ldf, spl, 5)
+            @test mean(chn1[:a]) == mean(chn2[:a])
+        end
+
+        @testset "check that initial_params are respected" begin
+            a0 = 5.0
+            chn = sample(model, ESS(), 5; initial_params=[a0])
+            @test chn[:a][1] == a0
+        end
+    end
+
     @model function demo(x)
         m ~ Normal()
         return x ~ Normal(m, 0.5)