Add back (some kind of) dynamic model test

Author: penelopeysm

test/mcmc/gibbs.jl

@@ -565,6 +565,99 @@ end
         end
     end
 
+    @testset "PG with variable number of observations" begin
+        # When sampling from a model with Particle Gibbs, it is mandatory for
+        # the number of observations to be the same in all particles, since the
+        # observations trigger particle resampling.
+        #
+        # Up until Turing v0.39, `x ~ dist` statements where `x` was the
+        # responsibility of a different (non-PG) Gibbs subsampler used to not
+        # count as an observation. Instead, the log-probability `logpdf(dist, x)`
+        # would be manually added to the VarInfo's `logp` field and included in the
+        # weighting for the _following_ observation.
+        #
+        # In Turing v0.40, this is now changed: `x ~ dist` uses tilde_observe!!
+        # which thus triggers resampling. Thus, for example, the following model
+        # does not work any more:
+        #
+        #   @model function f()
+        #       a ~ Poisson(2.0)
+        #       x = Vector{Float64}(undef, a)
+        #       for i in eachindex(x)
+        #           x[i] ~ Normal()
+        #       end
+        #   end
+        #   sample(f(), Gibbs(:a => PG(10), :x => MH()), 1000)
+        # 
+        # because the number of observations in each particle depends on the value
+        # of `a`.
+        #
+        # This testset checks that ways of working around such a situation.
+
+        function test_dynamic_bernoulli(chain)
+            means = Dict(:b => 0.5, "x[1]" => 1.0, "x[2]" => 2.0)
+            stds = Dict(:b => 0.5, "x[1]" => 1.0, "x[2]" => 1.0)
+            for vn in keys(means)
+                @test isapprox(mean(skipmissing(chain[:, vn, 1])), means[vn]; atol=0.1)
+                @test isapprox(std(skipmissing(chain[:, vn, 1])), stds[vn]; atol=0.1)
+            end
+        end
+
+        @testset "Coalescing multiple observations into one" begin
+            # Instead of observing x[1] and x[2] separately, we lump them into a
+            # single distribution.
+            @model function dynamic_bernoulli()
+                b ~ Bernoulli()
+                if b
+                    dists = [Normal(1.0)]
+                else
+                    dists = [Normal(1.0), Normal(2.0)]
+                end
+                return x ~ product_distribution(dists)
+            end
+            model = dynamic_bernoulli()
+            # This currently fails because if the global varinfo has `x` with length 2,
+            # and the particle sampler has `b = true`, it attempts to calculate the
+            # log-likelihood of a length-2 vector with respect to a length-1
+            # distribution.
+            @test_throws DimensionMismatch chain = sample(
+                StableRNG(468),
+                model,
+                Gibbs(:b => PG(10), :x => ESS()),
+                2000;
+                discard_initial=100,
+            )
+            # test_dynamic_bernoulli(chain)
+        end
+
+        @testset "Inserting @addlogprob!" begin
+            # On top of observing x[i], we also add in extra 'observations'
+            @model function dynamic_bernoulli_2()
+                b ~ Bernoulli()
+                x_length = b ? 1 : 2
+                x = Vector{Float64}(undef, x_length)
+                for i in 1:x_length
+                    x[i] ~ Normal(i, 1.0)
+                end
+                if length(x) == 1
+                    # This value is the expectation value of logpdf(Normal(), x) where x ~ Normal().
+                    # See discussion in
+                    # https://github.com/TuringLang/Turing.jl/pull/2629#discussion_r2237323817
+                    @addlogprob!(-1.418849)
+                end
+            end
+            model = dynamic_bernoulli_2()
+            chain = sample(
+                StableRNG(468),
+                model,
+                Gibbs(:b => PG(10), :x => ESS()),
+                2000;
+                discard_initial=100,
+            )
+            test_dynamic_bernoulli(chain)
+        end
+    end
+
     @testset "Demo model" begin
         @testset verbose = true "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
             vns = DynamicPPL.TestUtils.varnames(model)