Add VarInfo(::MarginalLogDensity) method

Author: penelopeysm

docs/make.jl

@@ -13,6 +13,9 @@ using DocumenterMermaid
 using MCMCChains
 using MarginalLogDensities: MarginalLogDensities
 
+# Need this to document a method which uses a type inside the extension...
+DPPLMLDExt = Base.get_extension(DynamicPPL, :DynamicPPLMarginalLogDensitiesExt)
+
 # Doctest setup
 DocMeta.setdocmeta!(
     DynamicPPL, :DocTestSetup, :(using DynamicPPL, MCMCChains); recursive=true

docs/src/api.md

@@ -145,6 +145,13 @@ This requires `MarginalLogDensities.jl` to be loaded in your environment.
 marginalize
 ```
 
+A `MarginalLogDensity` object acts as a function which maps non-marginalized parameter values to a marginal log-probability.
+To retrieve a VarInfo object from it, you can use:
+
+```@docs
+VarInfo(::MarginalLogDensities.MarginalLogDensity{<:DPPLMLDExt.LogDensityFunctionWrapper}, ::Union{AbstractVector,Nothing})
+```
+
 ## Models within models
 
 One can include models and call another model inside the model function with `left ~ to_submodel(model)`.

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -6,6 +6,16 @@ using MarginalLogDensities: MarginalLogDensities
 _to_varname(n::Symbol) = VarName{n}()
 _to_varname(n::VarName) = n
 
+# A thin wrapper to adapt a DynamicPPL.LogDensityFunction to the interface expected by
+# MarginalLogDensities. It's helpful to have a struct so that we can dispatch on its type
+# below.
+struct LogDensityFunctionWrapper{L<:DynamicPPL.LogDensityFunction}
+    logdensity::L
+end
+function (lw::LogDensityFunctionWrapper)(x, _)
+    return LogDensityProblems.logdensity(lw.logdensity, x)
+end
+
 """
     marginalize(
         model::DynamicPPL.Model,
@@ -26,7 +36,7 @@ log-density.
 ## Keyword arguments
 
 - `varinfo`: The `varinfo` to use for the model. By default we use a linked `VarInfo`,
-   meaning that the resulting log-density function accepts parameters that have bee_FWDn
+   meaning that the resulting log-density function accepts parameters that have been
    transformed to unconstrained space.
 
 - `getlogprob`: A function which specifies which kind of marginal log-density to compute.
@@ -60,6 +70,26 @@ julia> # The resulting callable computes the marginal log-density of `y`.
 julia> logpdf(Normal(2.0), 1.0)
 -1.4189385332046727
 ```
+
+
+!!! warning
+
+    The default usage of linked VarInfo means that, for example, optimization of the
+    marginal log-density can be performed in unconstrained space. However, care must be
+    taken if the model contains variables where the link transformation depends on a
+    marginalized variable. For example:
+
+    ```julia
+    @model function f()
+        x ~ Normal()
+        y ~ truncated(Normal(); lower=x)
+    end
+    ```
+
+    Here, the support of `y`, and hence the link transformation used, depends on the value
+    of `x`. If we now marginalize over `x`, we obtain a function mapping linked values of
+    `y` to log-probabilities. However, it will not be possible to use DynamicPPL to
+    correctly retrieve _unlinked_ values of `y`.
 """
 function DynamicPPL.marginalize(
     model::DynamicPPL.Model,
@@ -74,15 +104,104 @@ function DynamicPPL.marginalize(
     varindices = reduce(vcat, DynamicPPL.vector_getranges(varinfo, vns))
     # Construct the marginal log-density model.
     f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
-    mdl = MarginalLogDensities.MarginalLogDensity(
-        (x, _) -> LogDensityProblems.logdensity(f, x),
-        varinfo[:],
-        varindices,
-        (),
-        method;
-        kwargs...,
+    mld = MarginalLogDensities.MarginalLogDensity(
+        LogDensityFunctionWrapper(f), varinfo[:], varindices, (), method; kwargs...
+    )
+    return mld
+end
+
+"""
+    VarInfo(
+        mld::MarginalLogDensities.MarginalLogDensity{<:LogDensityFunctionWrapper},
+        unmarginalized_params::Union{AbstractVector,Nothing}=nothing
     )
-    return mdl
+
+Retrieve the `VarInfo` object used in the marginalisation process.
+
+If a Laplace approximation was used for the marginalisation, the values of the marginalized
+parameters are also set to their mode (note that this only happens if the `mld` object has
+been used to compute the marginal log-density at least once, so that the mode has been
+computed).
+
+If a vector of `unmarginalized_params` is specified, the values for the corresponding
+parameters will also be updated in the returned VarInfo. This vector may be obtained e.g. by
+performing an optimization of the marginal log-density.
+
+All other aspects of the VarInfo, such as link status, are preserved from the original
+VarInfo used in the marginalisation.
+
+!!! note
+
+    The other fields of the VarInfo, e.g. accumulated log-probabilities, will not be
+    updated. If you wish to have a fully consistent VarInfo, you should re-evaluate the
+    model with the returned VarInfo (e.g. using `vi = last(DynamicPPL.evaluate!!(model,
+    vi))`).
+
+## Example
+
+```jldoctest
+julia> using DynamicPPL, Distributions, MarginalLogDensities
+
+julia> @model function demo()
+           x ~ Normal()
+           y ~ Beta(2, 2)
+       end
+demo (generic function with 2 methods)
+
+julia> # Note that by default `marginalize` uses a linked VarInfo.
+       mld = marginalize(demo(), [@varname(x)]);
+
+julia> using MarginalLogDensities: Optimization, OptimizationOptimJL
+
+julia> # Find the mode of the marginal log-density of `y`, with an initial point of `y0`.
+       y0 = 2.0; opt_problem = Optimization.OptimizationProblem(mld, [y0])
+OptimizationProblem. In-place: true
+u0: 1-element Vector{Float64}:
+ 2.0
+
+julia> # This tells us the optimal (linked) value of `y` is around 0.
+       opt_solution = Optimization.solve(opt_problem, OptimizationOptimJL.NelderMead())
+retcode: Success
+u: 1-element Vector{Float64}:
+ 4.88281250001733e-5
+
+julia> # Get the VarInfo corresponding to the mode of `y`.
+       vi = VarInfo(mld, opt_solution.u);
+
+julia> # `x` is set to its mode (which for `Normal()` is zero).
+       vi[@varname(x)]
+0.0
+
+julia> # `y` is set to the optimal value we found above.
+       DynamicPPL.getindex_internal(vi, @varname(y))
+1-element Vector{Float64}:
+ 4.88281250001733e-5
+
+julia> # To obtain values in the original constrained space, we can either
+       # use `getindex`:
+       vi[@varname(y)]
+0.5000122070312476
+
+julia> # Or invlink the entire VarInfo object using the model:
+       vi_unlinked = DynamicPPL.invlink(vi, demo()); vi_unlinked[:]
+2-element Vector{Float64}:
+ 0.0
+ 0.5000122070312476
+```
+"""
+function DynamicPPL.VarInfo(
+    mld::MarginalLogDensities.MarginalLogDensity{<:LogDensityFunctionWrapper},
+    unmarginalized_params::Union{AbstractVector,Nothing}=nothing,
+)
+    # Extract the original VarInfo. Its contents will in general be junk.
+    original_vi = mld.logdensity.logdensity.varinfo
+    # `mld.u` will contain the modes for any marginalized parameters
+    full_params = mld.u
+    # We can then set the values for any non-marginalized parameters
+    if unmarginalized_params !== nothing
+        full_params[MarginalLogDensities.ijoint(mld)] = unmarginalized_params
+    end
+    return DynamicPPL.unflatten(original_vi, full_params)
 end
 
 end

test/ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -69,6 +69,38 @@ using ADTypes: AutoForwardDiff
             end
         end
     end
+
+    @testset "retrieving VarInfo from MLD" 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)]; varinfo=vi_unlinked)
+            mx([0.5]) # evaluate at some point to force calculation of Laplace approx
+            vi = VarInfo(mx)
+            @test vi[@varname(x)] ≈ mode(Normal())
+            vi = VarInfo(mx, [0.5]) # this 0.5 is unlinked
+            @test vi[@varname(x)] ≈ mode(Normal())
+            @test vi[@varname(y)] ≈ 0.5
+        end
+
+        @testset "linked VarInfo" begin
+            mx = marginalize(model, [@varname(x)]; varinfo=vi_linked)
+            mx([0.5]) # evaluate at some point to force calculation of Laplace approx
+            vi = VarInfo(mx)
+            @test vi[@varname(x)] ≈ mode(Normal())
+            vi = VarInfo(mx, [0.5]) # this 0.5 is linked
+            binv = Bijectors.inverse(Bijectors.bijector(Beta(2, 2)))
+            @test vi[@varname(x)] ≈ mode(Normal())
+            # when using getindex it always returns unlinked values
+            @test vi[@varname(y)] ≈ binv(0.5)
+        end
+    end
 end
 
 end