Implement AbstractMCMC interface for SampleFromPrior and SampleFromUniform

Author: devmotion

Project.toml

@@ -8,6 +8,7 @@ AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]

src/DynamicPPL.jl

@@ -4,6 +4,9 @@ using AbstractMCMC: AbstractSampler, AbstractChains, AbstractModel
 using Distributions
 using Bijectors
 using MacroTools
+
+import AbstractMCMC
+import Random
 import ZygoteRules
 
 import Base: Symbol,

src/sampler.jl

@@ -43,3 +43,18 @@ end
 Sampler(alg) = Sampler(alg, Selector())
 Sampler(alg, model::Model) = Sampler(alg, model, Selector())
 Sampler(alg, model::Model, s::Selector) = Sampler(alg, model, s)
+
+# AbstractMCMC interface for SampleFromUniform and SampleFromPrior
+
+function AbstractMCMC.step!(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::Union{SampleFromUniform,SampleFromPrior},
+    ::Integer,
+    transition;
+    kwargs...
+)
+    vi = VarInfo()
+    model(vi, sampler)
+    return vi
+end

test/runtests.jl

@@ -9,6 +9,7 @@ turnprogress(false)
     include("utils.jl")
     include("compiler.jl")
     include("varinfo.jl")
+    include("sampler.jl")
     include("prob_macro.jl")
     include("independence.jl")
 end

test/sampler.jl

@@ -0,0 +1,34 @@
+using DynamicPPL
+using Distributions
+using AbstractMCMC: sample
+
+using Random
+using Statistics
+using Test
+
+Random.seed!(1234)
+
+@testset "AbstractMCMC interface" begin
+    @model gdemo(x, y) = begin
+        s ~ InverseGamma(2, 3)
+        m ~ Normal(0.0, sqrt(s))
+        x ~ Normal(m, sqrt(s))
+        y ~ Normal(m, sqrt(s))
+    end
+
+    model = gdemo(1.0, 2.0)
+    N = 10_000
+
+    chains = sample(model, SampleFromPrior(), N; progress = false)
+    @test chains isa Vector{<:VarInfo}
+    @test length(chains) == N
+    @test mean(vi[@varname(m)] for vi in chains) ≈ 0 atol = 0.1
+    @test mean(vi[@varname(s)] for vi in chains) ≈ 3 atol = 0.1
+
+    chains = sample(model, SampleFromUniform(), N; progress = false)
+    @test chains isa Vector{<:VarInfo}
+    @test length(chains) == N
+    @test mean(vi[@varname(m)] for vi in chains) ≈ 1 atol = 0.1
+    @test mean(vi[@varname(s)] for vi in chains) ≈ 3.3 atol = 0.1
+end
+