Retain reference particle properly (#1599)

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.15.18"
+version = "0.15.19"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/inference/AdvancedSMC.jl

@@ -281,17 +281,21 @@ function AbstractMCMC.step(
 )
     # Reset the VarInfo before new sweep.
     reset_num_produce!(vi)
-    set_retained_vns_del_by_spl!(vi, spl)
     resetlogp!(vi)
 
+    # Create reference particle for which the samples will be retained.
+    reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi))
+
+    # For all other particles, do not retain the variables but resample them.
+    set_retained_vns_del_by_spl!(vi, spl)
+
     # Create a new set of particles.
     num_particles = spl.alg.nparticles
     x = map(1:num_particles) do i
         if i != num_particles
             return AdvancedPS.Trace(model, spl, vi)
         else
-        # Create reference particle.
-            return AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi))
+            return reference
         end
     end
     particles = AdvancedPS.ParticleContainer(x)

test/inference/AdvancedSMC.jl

@@ -166,6 +166,13 @@ end
         @test all(isone, chains_pg[:x])
         @test chains_pg.logevidence ≈ -2 * log(2) atol = 0.01
     end
+
+    # https://github.com/TuringLang/Turing.jl/issues/1598
+    @turing_testset "reference particle" begin
+        c = sample(gdemo_default, PG(1), 1_000)
+        @test length(unique(c[:m])) == 1
+        @test length(unique(c[:s])) == 1
+    end
 end
 
 # @testset "pmmh.jl" begin

test/inference/Inference.jl

@@ -53,7 +53,7 @@
 
         alg1 = HMCDA(1000, 0.65, 0.15)
         alg2 = PG(20)
-        alg3 = Gibbs(PG(30, :s), HMCDA(500, 0.65, 0.05, :m))
+        alg3 = Gibbs(PG(30, :s), HMC(0.2, 4, :m))
 
         chn1 = sample(gdemo_default, alg1, 5000; save_state=true)
         check_gdemo(chn1)
@@ -73,31 +73,31 @@
         chn2_contd = sample(gdemo_default, alg2, 1000; resume_from=chn2)
         check_gdemo(chn2_contd)
 
-        chn3 = sample(gdemo_default, alg3, 1000; save_state=true)
+        chn3 = sample(gdemo_default, alg3, 5000; save_state=true)
         check_gdemo(chn3)
 
         chn3_contd = sample(gdemo_default, alg3, 1000; resume_from=chn3)
-        check_gdemo(chn3_contd; atol=0.25)
+        check_gdemo(chn3_contd)
     end
     @testset "Contexts" begin
         # Test LikelihoodContext
-        @model testmodel(x) = begin
+        @model function testmodel1(x)
             a ~ Beta()
-            lp1 = getlogp(_varinfo)
-            x[1] ~ Bernoulli(a)
-            global loglike = getlogp(_varinfo) - lp1
+            lp1 = getlogp(__varinfo__)
+            x[1] ~ Bernoulli(a)
+            global loglike = getlogp(__varinfo__) - lp1
         end
-        model = testmodel([1.0])
+        model = testmodel1([1.0])
         varinfo = Turing.VarInfo(model)
         model(varinfo, Turing.SampleFromPrior(), Turing.LikelihoodContext())
         @test getlogp(varinfo) == loglike
 
         # Test MiniBatchContext
-        @model testmodel(x) = begin
+        @model function testmodel2(x)
             a ~ Beta()
-            x[1] ~ Bernoulli(a)
+            x[1] ~ Bernoulli(a)
         end
-        model = testmodel([1.0])
+        model = testmodel2([1.0])
         varinfo1 = Turing.VarInfo(model)
         varinfo2 = deepcopy(varinfo1)
         model(varinfo1, Turing.SampleFromPrior(), Turing.LikelihoodContext())

test/inference/gibbs.jl

@@ -35,8 +35,8 @@
         alg = Gibbs(
             CSMC(10, :s),
             HMC(0.2, 4, :m))
-        chain = sample(gdemo(1.5, 2.0), alg, 3000)
-        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)
+        chain = sample(gdemo(1.5, 2.0), alg, 1_500)
+        check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.15)
 
         Random.seed!(100)