Allow user to specify VarInfo used for marginalisation

Author: penelopeysm

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -9,6 +9,7 @@ _to_varname(n::VarName) = n
 """
     marginalize(
         model::DynamicPPL.Model,
+        varinfo::DynamicPPL.AbstractVarInfo=VarInfo(model),
         varnames::AbstractVector{<:Union{Symbol,<:VarName}},
         getlogprob=DynamicPPL.getlogjoint,
         method::MarginalLogDensities.AbstractMarginalizer=MarginalLogDensities.LaplaceApprox();
@@ -54,22 +55,19 @@ julia> logpdf(Normal(2.0), 1.0)
 function DynamicPPL.marginalize(
     model::DynamicPPL.Model,
     varnames::AbstractVector{<:Union{Symbol,<:VarName}},
-    getlogprob=DynamicPPL.getlogjoint,
+    varinfo::DynamicPPL.AbstractVarInfo=DynamicPPL.VarInfo(model),
+    getlogprob::Function=DynamicPPL.getlogjoint,
     method::MarginalLogDensities.AbstractMarginalizer=MarginalLogDensities.LaplaceApprox();
     kwargs...,
 )
     # Determine the indices for the variables to marginalise out.
-    varinfo = DynamicPPL.typed_varinfo(model)
     vns = map(_to_varname, varnames)
     varindices = reduce(vcat, DynamicPPL.vector_getranges(varinfo, vns))
     # Construct the marginal log-density model.
-    # Use linked `varinfo` to that we're working in unconstrained space
-    varinfo_linked = DynamicPPL.link(varinfo, model)
-
-    f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo_linked)
+    f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
     mdl = MarginalLogDensities.MarginalLogDensity(
         (x, _) -> LogDensityProblems.logdensity(f, x),
-        varinfo_linked[:],
+        varinfo[:],
         varindices,
         (),
         method;

test/ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -1,25 +1,66 @@
 module MarginalLogDensitiesExtTests
 
+using Bijectors: Bijectors
 using DynamicPPL, Distributions, Test
 using MarginalLogDensities
 using ADTypes: AutoForwardDiff
 
 @testset "MarginalLogDensities" begin
-    # Simple test case.
-    @model function demo()
-        x ~ MvNormal(zeros(2), [1, 1])
-        return y ~ Normal(0, 1)
+    @testset "Basic usage" begin
+        @model function demo()
+            x ~ MvNormal(zeros(2), [1, 1])
+            return y ~ Normal(0, 1)
+        end
+        model = demo()
+        vi = VarInfo(model)
+        # Marginalize out `x`.
+        for vn in [@varname(x), :x]
+            for getlogprob in [DynamicPPL.getlogprior, DynamicPPL.getlogjoint]
+                marginalized = marginalize(
+                    model, [vn], vi, getlogprob; hess_adtype=AutoForwardDiff()
+                )
+                for y in range(-5, 5; length=100)
+                    @test marginalized([y]) ≈ logpdf(Normal(0, 1), y) atol = 1e-5
+                end
+            end
+        end
     end
-    model = demo()
-    # Marginalize out `x`.
 
-    for vn in [@varname(x), :x]
-        for getlogprob in [DynamicPPL.getlogprior, DynamicPPL.getlogjoint]
-            marginalized = marginalize(
-                model, [vn], getlogprob; hess_adtype=AutoForwardDiff()
-            )
-            # Compute the marginal log-density of `y = 0.0`.
-            @test marginalized([0.0]) ≈ logpdf(Normal(0, 1), 0.0) atol = 1e-5
+    @testset "Respects linked status of VarInfo" begin
+        @model function f()
+            x ~ Normal()
+            return y ~ Beta(2, 2)
+        end
+        model = f()
+        vi_unlinked = VarInfo(model)
+        vi_linked = DynamicPPL.link(vi_unlinked, model)
+
+        @testset "unlinked VarInfo" begin
+            mx = marginalize(model, [@varname(x)], vi_unlinked)
+            for x in range(0.01, 0.99; length=10)
+                @test mx([x]) ≈ logpdf(Beta(2, 2), x)
+            end
+            # generally when marginalising Beta it doesn't go to zero
+            my = marginalize(model, [@varname(y)], vi_unlinked)
+            diff = my([0.0]) - logpdf(Normal(), 0.0)
+            for x in range(-5, 5; length=10)
+                @test my([x]) ≈ logpdf(Normal(), x) + diff
+            end
+        end
+
+        @testset "linked VarInfo" begin
+            mx = marginalize(model, [@varname(x)], vi_linked)
+            binv = Bijectors.inverse(Bijectors.bijector(Beta(2, 2)))
+            for y_linked in range(-5, 5; length=10)
+                y_unlinked = binv(y_linked)
+                @test mx([y_linked]) ≈ logpdf(Beta(2, 2), y_unlinked)
+            end
+            # generally when marginalising Beta it doesn't go to zero
+            my = marginalize(model, [@varname(y)], vi_linked)
+            diff = my([0.0]) - logpdf(Normal(), 0.0)
+            for x in range(-5, 5; length=10)
+                @test my([x]) ≈ logpdf(Normal(), x) + diff
+            end
         end
     end
 end