Give LogDensityFunction the getlogdensity field

Author: mhauru

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -87,7 +87,7 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
         vi = DynamicPPL.link(vi, model)
     end
 
-    f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adbackend)
+    f = DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint, vi, context; adtype=adbackend)
     # The parameters at which we evaluate f.
     θ = vi[:]
 

src/logdensityfunction.jl

@@ -17,7 +17,8 @@ is_supported(::ADTypes.AutoReverseDiff) = true
 """
     LogDensityFunction(
         model::Model,
-        varinfo::AbstractVarInfo=VarInfo(model),
+        getlogdensity::Function=getlogjoint,
+        varinfo::AbstractVarInfo=ldf_default_varinfo(model, getlogdensity),
         context::AbstractContext=DefaultContext();
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing
     )
@@ -28,10 +29,10 @@ A struct which contains a model, along with all the information necessary to:
  - and if `adtype` is provided, calculate the gradient of the log density at
  that point.
 
-At its most basic level, a LogDensityFunction wraps the model together with its
-the type of varinfo to be used, as well as the evaluation context. These must
-be known in order to calculate the log density (using
-[`DynamicPPL.evaluate!!`](@ref)).
+At its most basic level, a LogDensityFunction wraps the model together with
+the type of varinfo to be used, as well as the evaluation context and a function
+to extract the log density from the VarInfo. These must be known in order to
+calculate the log density (using [`DynamicPPL.evaluate!!`](@ref)).
 
 If the `adtype` keyword argument is provided, then this struct will also store
 the adtype along with other information for efficient calculation of the
@@ -73,13 +74,13 @@ julia> LogDensityProblems.dimension(f)
 1
 
 julia> # By default it uses `VarInfo` under the hood, but this is not necessary.
-       f = LogDensityFunction(model, SimpleVarInfo(model));
+       f = LogDensityFunction(model, getlogjoint, SimpleVarInfo(model));
 
 julia> LogDensityProblems.logdensity(f, [0.0])
 -2.3378770664093453
 
-julia> # LogDensityFunction respects the accumulators in VarInfo:
-       f_prior = LogDensityFunction(model, setaccs!!(VarInfo(model), (LogPriorAccumulator(),)));
+julia> # One can also specify evaluating e.g. the log prior only:
+       f_prior = LogDensityFunction(model, getprior);
 
 julia> LogDensityProblems.logdensity(f_prior, [0.0]) == logpdf(Normal(), 0.0)
 true
@@ -94,11 +95,13 @@ julia> LogDensityProblems.logdensity_and_gradient(f, [0.0])
 ```
 """
 struct LogDensityFunction{
-    M<:Model,V<:AbstractVarInfo,C<:AbstractContext,AD<:Union{Nothing,ADTypes.AbstractADType}
+    M<:Model,F<:Function,V<:AbstractVarInfo,C<:AbstractContext,AD<:Union{Nothing,ADTypes.AbstractADType}
 }
     "model used for evaluation"
     model::M
-    "varinfo used for evaluation"
+    "function to be called on `varinfo` to extract the log density. By default `getlogjoint`."
+    getlogdensity::F
+    "varinfo used for evaluation. If not specified, generated with `ldf_default_varinfo`."
     varinfo::V
     "context used for evaluation; if `nothing`, `leafcontext(model.context)` will be used when applicable"
     context::C
@@ -109,7 +112,8 @@ struct LogDensityFunction{
 
     function LogDensityFunction(
         model::Model,
-        varinfo::AbstractVarInfo=VarInfo(model),
+        getlogdensity::Function=getlogjoint,
+        varinfo::AbstractVarInfo=ldf_default_varinfo(model, getlogdensity),
         context::AbstractContext=leafcontext(model.context);
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing,
     )
@@ -125,21 +129,22 @@ struct LogDensityFunction{
             x = map(identity, varinfo[:])
             if use_closure(adtype)
                 prep = DI.prepare_gradient(
-                    x -> logdensity_at(x, model, varinfo, context), adtype, x
+                    x -> logdensity_at(x, model, getlogdensity, varinfo, context), adtype, x
                 )
             else
                 prep = DI.prepare_gradient(
                     logdensity_at,
                     adtype,
                     x,
                     DI.Constant(model),
+                    DI.Constant(getlogdensity),
                     DI.Constant(varinfo),
                     DI.Constant(context),
                 )
             end
         end
-        return new{typeof(model),typeof(varinfo),typeof(context),typeof(adtype)}(
-            model, varinfo, context, adtype, prep
+        return new{typeof(model),typeof(getlogdensity),typeof(varinfo),typeof(context),typeof(adtype)}(
+            model, getlogdensity, varinfo, context, adtype, prep
         )
     end
 end
@@ -164,64 +169,80 @@ function LogDensityFunction(
     end
 end
 
+"""
+    ldf_default_varinfo(model::Model, getlogdensity::Function)
+
+Create the default AbstractVarInfo that should be used for evaluating the log density.
+
+Only the accumulators necesessary for `getlogdensity` will be used.
+"""
+function ldf_default_varinfo(::Model, getlogdensity::Function)
+        msg = """
+        LogDensityFunction does not know what sort of VarInfo should be used when \
+        `getlogdensity` is $getlogdensity. Please specify a VarInfo explicitly.
+        """
+        error(msg)
+    end
+
+ldf_default_varinfo(model::Model, ::typeof(getlogjoint)) = VarInfo(model)
+
+function ldf_default_varinfo(model::Model, ::typeof(getlogprior))
+    return setaccs!!(VarInfo(model), (LogPriorAccumulator(),))
+end
+
+function ldf_default_varinfo(model::Model, ::typeof(getloglikelihood))
+    return setaccs!!(VarInfo(model), (LogLikelihoodAccumulator(),))
+end
+
 """
     logdensity_at(
         x::AbstractVector,
         model::Model,
+        getlogdensity::Function,
         varinfo::AbstractVarInfo,
         context::AbstractContext
     )
 
 Evaluate the log density of the given `model` at the given parameter values `x`,
 using the given `varinfo` and `context`. Note that the `varinfo` argument is provided
 only for its structure, in the sense that the parameters from the vector `x` are inserted
-into it, and its own parameters are discarded. It does, however, determine whether the log
-prior, likelihood, or joint is returned, based on which accumulators are set in it.
+into it, and its own parameters are discarded. `getlogdensity` is the function that extracts
+the log density from the evaluated varinfo.
 """
 function logdensity_at(
-    x::AbstractVector, model::Model, varinfo::AbstractVarInfo, context::AbstractContext
+    x::AbstractVector, model::Model, getlogdensity::Function, varinfo::AbstractVarInfo, context::AbstractContext
 )
     varinfo_new = unflatten(varinfo, x)
     varinfo_eval = last(evaluate!!(model, varinfo_new, context))
-    has_prior = hasacc(varinfo_eval, Val(:LogPrior))
-    has_likelihood = hasacc(varinfo_eval, Val(:LogLikelihood))
-    if has_prior && has_likelihood
-        return getlogjoint(varinfo_eval)
-    elseif has_prior
-        return getlogprior(varinfo_eval)
-    elseif has_likelihood
-        return getloglikelihood(varinfo_eval)
-    else
-        error("LogDensityFunction: varinfo tracks neither log prior nor log likelihood")
-    end
+    return getlogdensity(varinfo_eval)
 end
 
 ### LogDensityProblems interface
 
 function LogDensityProblems.capabilities(
-    ::Type{<:LogDensityFunction{M,V,C,Nothing}}
-) where {M,V,C}
+    ::Type{<:LogDensityFunction{M,F,V,C,Nothing}}
+) where {M,F,V,C}
     return LogDensityProblems.LogDensityOrder{0}()
 end
 function LogDensityProblems.capabilities(
-    ::Type{<:LogDensityFunction{M,V,C,AD}}
-) where {M,V,C,AD<:ADTypes.AbstractADType}
+    ::Type{<:LogDensityFunction{M,F,V,C,AD}}
+) where {M,F,V,C,AD<:ADTypes.AbstractADType}
     return LogDensityProblems.LogDensityOrder{1}()
 end
 function LogDensityProblems.logdensity(f::LogDensityFunction, x::AbstractVector)
-    return logdensity_at(x, f.model, f.varinfo, f.context)
+    return logdensity_at(x, f.model, f.getlogdensity, f.varinfo, f.context)
 end
 function LogDensityProblems.logdensity_and_gradient(
-    f::LogDensityFunction{M,V,C,AD}, x::AbstractVector
-) where {M,V,C,AD<:ADTypes.AbstractADType}
+    f::LogDensityFunction{M,F,V,C,AD}, x::AbstractVector
+) where {M,F,V,C,AD<:ADTypes.AbstractADType}
     f.prep === nothing &&
         error("Gradient preparation not available; this should not happen")
     x = map(identity, x)  # Concretise type
     # Make branching statically inferrable, i.e. type-stable (even if the two
     # branches happen to return different types)
     return if use_closure(f.adtype)
         DI.value_and_gradient(
-            x -> logdensity_at(x, f.model, f.varinfo, f.context), f.prep, f.adtype, x
+            x -> logdensity_at(x, f.model, f.getlogdensity, f.varinfo, f.context), f.prep, f.adtype, x
         )
     else
         DI.value_and_gradient(
@@ -230,6 +251,7 @@ function LogDensityProblems.logdensity_and_gradient(
             f.adtype,
             x,
             DI.Constant(f.model),
+            DI.Constant(f.getlogdensity),
             DI.Constant(f.varinfo),
             DI.Constant(f.context),
         )
@@ -304,7 +326,7 @@ getmodel(f::DynamicPPL.LogDensityFunction) = f.model
 Set the `DynamicPPL.Model` in the given log-density function `f` to `model`.
 """
 function setmodel(f::DynamicPPL.LogDensityFunction, model::DynamicPPL.Model)
-    return LogDensityFunction(model, f.varinfo, f.context; adtype=f.adtype)
+    return LogDensityFunction(model, f.getlogdensity, f.varinfo, f.context; adtype=f.adtype)
 end
 
 """

test/ad.jl

@@ -24,10 +24,10 @@ using DynamicPPL: LogDensityFunction
 
             @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
                 linked_varinfo = DynamicPPL.link(varinfo, m)
-                f = LogDensityFunction(m, linked_varinfo)
+                f = LogDensityFunction(m, getlogjoint, linked_varinfo)
                 x = DynamicPPL.getparams(f)
                 # Calculate reference logp + gradient of logp using ForwardDiff
-                ref_ldf = LogDensityFunction(m, linked_varinfo; adtype=ref_adtype)
+                ref_ldf = LogDensityFunction(m, getlogjoint, linked_varinfo; adtype=ref_adtype)
                 ref_logp, ref_grad = LogDensityProblems.logdensity_and_gradient(ref_ldf, x)
 
                 @testset "$adtype" for adtype in test_adtypes
@@ -106,7 +106,7 @@ using DynamicPPL: LogDensityFunction
         spl = Sampler(MyEmptyAlg())
         vi = VarInfo(model)
         ldf = LogDensityFunction(
-            model, vi, SamplingContext(spl); adtype=AutoReverseDiff(; compile=true)
+            model, getlogjoint, vi, SamplingContext(spl); adtype=AutoReverseDiff(; compile=true)
         )
         @test LogDensityProblems.logdensity_and_gradient(ldf, vi[:]) isa Any
     end

test/logdensityfunction.jl

@@ -15,8 +15,17 @@ end
         vns = DynamicPPL.TestUtils.varnames(model)
         varinfos = DynamicPPL.TestUtils.setup_varinfos(model, example_values, vns)
 
+        vi = first(varinfos)
+        theta = vi[:]
+        ldf_joint = DynamicPPL.LogDensityFunction(model)
+        @test LogDensityProblems.logdensity(ldf_joint, theta) ≈ logjoint(model, vi)
+        ldf_prior = DynamicPPL.LogDensityFunction(model, getlogprior)
+        @test LogDensityProblems.logdensity(ldf_prior, theta) ≈ logprior(model, vi)
+        ldf_likelihood = DynamicPPL.LogDensityFunction(model, getloglikelihood)
+        @test LogDensityProblems.logdensity(ldf_likelihood, theta) ≈ loglikelihood(model, vi)
+
         @testset "$(varinfo)" for varinfo in varinfos
-            logdensity = DynamicPPL.LogDensityFunction(model, varinfo)
+            logdensity = DynamicPPL.LogDensityFunction(model, getlogjoint, varinfo)
             θ = varinfo[:]
             @test LogDensityProblems.logdensity(logdensity, θ) ≈ logjoint(model, varinfo)
             @test LogDensityProblems.dimension(logdensity) == length(θ)

test/test_util.jl

@@ -14,7 +14,7 @@ function test_model_ad(model, logp_manual)
     x = vi[:]
 
     # Log probabilities using the model.
-    ℓ = DynamicPPL.LogDensityFunction(model, vi)
+    ℓ = DynamicPPL.LogDensityFunction(model, getlogjoint, vi)
     logp_model = Base.Fix1(LogDensityProblems.logdensity, ℓ)
 
     # Check that both functions return the same values.