Fix everything (I _think_)

Author: penelopeysm

test/mcmc/ess.jl

@@ -2,7 +2,7 @@ module ESSTests
 
 using ..Models: MoGtest, MoGtest_default, gdemo, gdemo_default
 using ..NumericalTests: check_MoGtest_default, check_numerical
-using ..SamplerTestUtils: check_rng_respected
+using ..SamplerTestUtils: test_rng_respected
 using Distributions: Normal, sample
 using DynamicPPL: DynamicPPL
 using DynamicPPL: Sampler
@@ -40,9 +40,9 @@ using Turing
     end
 
     @testset "RNG is respected" begin
-        check_rng_respected(ESS())
-        check_rng_respected(Gibbs(:x => ESS(), :y => MH()))
-        check_rng_respected(Gibbs(:x => ESS(), :y => ESS()))
+        test_rng_respected(ESS())
+        test_rng_respected(Gibbs(:x => ESS(), :y => MH()))
+        test_rng_respected(Gibbs(:x => ESS(), :y => ESS()))
     end
 
     @testset "ESS inference" begin

test/mcmc/external_sampler.jl

@@ -174,7 +174,7 @@ end
 
 function test_initial_params(model, sampler; kwargs...)
     # Generate some parameters.
-    dict = DynamicPPL.values_as(VarInfo(model), Dict)
+    dict = DynamicPPL.values_as(DynamicPPL.VarInfo(model), Dict)
     init_strategy = DynamicPPL.InitFromParams(dict)
 
     # Execute the transition with two different RNGs and check that the resulting

test/mcmc/is.jl

@@ -1,63 +1,52 @@
 module ISTests
 
-using Distributions: Normal, sample
 using DynamicPPL: logpdf
 using Random: Random
+using StableRNGs: StableRNG
 using StatsFuns: logsumexp
 using Test: @test, @testset
 using Turing
 
 @testset "is.jl" begin
-    function reference(n)
-        as = Vector{Float64}(undef, n)
-        bs = Vector{Float64}(undef, n)
-        logps = Vector{Float64}(undef, n)
+    @testset "numerical accuracy" begin
+        function reference(n)
+            rng = StableRNG(468)
+            as = Vector{Float64}(undef, n)
+            bs = Vector{Float64}(undef, n)
 
-        for i in 1:n
-            as[i], bs[i], logps[i] = reference()
+            for i in 1:n
+                as[i] = rand(rng, Normal(4, 5))
+                bs[i] = rand(rng, Normal(as[i], 1))
+            end
+            # logevidence = logsumexp(logps) - log(n)
+            return (as=as, bs=bs)
         end
-        logevidence = logsumexp(logps) - log(n)
 
-        return (as=as, bs=bs, logps=logps, logevidence=logevidence)
-    end
-
-    function reference()
-        x = rand(Normal(4, 5))
-        y = rand(Normal(x, 1))
-        loglik = logpdf(Normal(x, 2), 3) + logpdf(Normal(y, 2), 1.5)
-        return x, y, loglik
-    end
-
-    @model function normal()
-        a ~ Normal(4, 5)
-        3 ~ Normal(a, 2)
-        b ~ Normal(a, 1)
-        1.5 ~ Normal(b, 2)
-        return a, b
-    end
-
-    alg = IS()
-    seed = 0
-    n = 10
+        @model function normal()
+            a ~ Normal(4, 5)
+            3 ~ Normal(a, 2)
+            b ~ Normal(a, 1)
+            1.5 ~ Normal(b, 2)
+            return a, b
+        end
 
-    model = normal()
-    for i in 1:100
-        Random.seed!(seed)
-        ref = reference(n)
+        function expected_loglikelihoods(as, bs)
+            return logpdf.(Normal.(as, 2), 3) .+ logpdf.(Normal.(bs, 2), 1.5)
+        end
 
-        Random.seed!(seed)
-        chain = sample(model, alg, n; check_model=false)
-        sampled = get(chain, [:a, :b, :loglikelihood])
+        alg = IS()
+        N = 1000
+        model = normal()
+        chain = sample(StableRNG(468), model, alg, N)
+        ref = reference(N)
 
-        @test vec(sampled.a) == ref.as
-        @test vec(sampled.b) == ref.bs
-        @test vec(sampled.loglikelihood) == ref.logps
-        @test chain.logevidence == ref.logevidence
+        @test isapprox(mean(chain[:a]), mean(ref.as); atol=0.1)
+        @test isapprox(mean(chain[:b]), mean(ref.bs); atol=0.1)
+        @test isapprox(chain[:loglikelihood], expected_loglikelihoods(chain[:a], chain[:b]))
+        @test isapprox(chain.logevidence, logsumexp(chain[:loglikelihood]) - log(N))
     end
 
     @testset "logevidence" begin
-        Random.seed!(100)
-
         @model function test()
             a ~ Normal(0, 1)
             x ~ Bernoulli(1)

test/test_utils/sampler.jl

@@ -25,6 +25,9 @@ function test_chain_logp_metadata(spl)
     @test chn[:lp] ≈ chn[:logprior] + chn[:loglikelihood]
 end
 
+"""
+Check that sampling is deterministic when using the same RNG seed.
+"""
 function test_rng_respected(spl)
     @model function f(z)
         # put at least two variables here so that we can meaningfully test Gibbs
@@ -35,7 +38,8 @@ function test_rng_respected(spl)
     model = f(2.0)
     chn1 = sample(Xoshiro(468), model, spl, 100)
     chn2 = sample(Xoshiro(468), model, spl, 100)
-    @test chn1 == chn2
+    @test isapprox(chn1[:x], chn2[:x])
+    @test isapprox(chn1[:y], chn2[:y])
 end
 
 end