Fix OptimLogDensity invocations

Author: penelopeysm

src/optimisation/Optimisation.jl

@@ -206,7 +206,9 @@ function StatsBase.informationmatrix(
     if linked
         new_vi = DynamicPPL.invlink!!(old_ldf.varinfo, old_ldf.model)
         new_f = OptimLogDensity(
-            old_ldf.model, old_ldf.getlogdensity, new_vi; adtype=old_ldf.adtype
+            DynamicPPL.LogDensityFunction(
+                old_ldf.model, old_ldf.getlogdensity, new_vi; adtype=old_ldf.adtype
+            ),
         )
         m = Accessors.@set m.f = new_f
     end
@@ -221,7 +223,12 @@ function StatsBase.informationmatrix(
         invlinked_ldf = m.f.ldf
         new_vi = DynamicPPL.link!!(invlinked_ldf.varinfo, invlinked_ldf.model)
         new_f = OptimLogDensity(
-            invlinked_ldf.model, old_ldf.getlogdensity, new_vi; adtype=invlinked_ldf.adtype
+            DynamicPPL.LogDensityFunction(
+                invlinked_ldf.model,
+                old_ldf.getlogdensity,
+                new_vi;
+                adtype=invlinked_ldf.adtype,
+            ),
         )
         m = Accessors.@set m.f = new_f
     end

test/optimisation/Optimisation.jl

@@ -44,35 +44,49 @@ using Turing
             m2 = model2() | (x=x,)
             # Doesn't matter if we use getlogjoint or getlogjoint_internal since the
             # VarInfo isn't linked.
-            ld1 = Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogjoint)
-            ld2 = Turing.Optimisation.OptimLogDensity(m2, DynamicPPL.getlogjoint_internal)
+            ld1 = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogjoint)
+            )
+            ld2 = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m2, DynamicPPL.getlogjoint_internal)
+            )
             @test ld1(w) == ld2(w)
         end
 
         @testset "With prefixes" begin
             vn = @varname(inner)
             m1 = prefix(model1(x), vn)
             m2 = prefix((model2() | (x=x,)), vn)
-            ld1 = Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogjoint)
-            ld2 = Turing.Optimisation.OptimLogDensity(m2, DynamicPPL.getlogjoint_internal)
+            ld1 = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogjoint)
+            )
+            ld2 = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m2, DynamicPPL.getlogjoint_internal)
+            )
             @test ld1(w) == ld2(w)
         end
 
         @testset "Joint, prior, and likelihood" begin
             m1 = model1(x)
             a = [0.3]
-            ld_joint = Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogjoint)
-            ld_prior = Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogprior)
+            ld_joint = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogjoint)
+            )
+            ld_prior = Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogprior)
+            )
             ld_likelihood = Turing.Optimisation.OptimLogDensity(
-                m1, DynamicPPL.getloglikelihood
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getloglikelihood)
             )
             @test ld_joint(a) == ld_prior(a) + ld_likelihood(a)
 
             # test that the prior accumulator is calculating the right thing
-            @test Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogprior)([0.3]) ≈
-                -Distributions.logpdf(Uniform(0, 2), 0.3)
-            @test Turing.Optimisation.OptimLogDensity(m1, DynamicPPL.getlogprior)([-0.3]) ≈
-                -Distributions.logpdf(Uniform(0, 2), -0.3)
+            @test Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogprior)
+            )([0.3]) ≈ -Distributions.logpdf(Uniform(0, 2), 0.3)
+            @test Turing.Optimisation.OptimLogDensity(
+                DynamicPPL.LogDensityFunction(m1, DynamicPPL.getlogprior)
+            )([-0.3]) ≈ -Distributions.logpdf(Uniform(0, 2), -0.3)
         end
     end