Moved LogDensityFunction from Turing to DPPL (#447)

Related to TuringLang/Turing.jl#1936

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.21.4"
+version = "0.21.5"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -12,6 +12,7 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -28,6 +29,7 @@ ChainRulesCore = "0.9.7, 0.10, 1"
 ConstructionBase = "1"
 Distributions = "0.23.8, 0.24, 0.25"
 DocStringExtensions = "0.8, 0.9"
+LogDensityProblems = "2"
 MacroTools = "0.5.6"
 OrderedCollections = "1"
 Setfield = "0.7.1, 0.8, 1"

docs/Project.toml

@@ -3,6 +3,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -12,6 +13,7 @@ DataStructures = "0.18"
 Distributions = "0.25"
 Documenter = "0.27"
 FillArrays = "0.13"
+LogDensityProblems = "2"
 MLUtils = "0.3, 0.4"
 Setfield = "0.7.1, 0.8, 1"
 StableRNGs = "1"

docs/src/api.md

@@ -58,6 +58,14 @@ loglikelihood
 logjoint
 ```
 
+### LogDensityProblems.jl interface
+
+The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) interface is also supported by simply wrapping a [`Model`](@ref) in a `DynamicPPL.LogDensityFunction`:
+
+```@docs
+DynamicPPL.LogDensityFunction
+```
+
 ## Condition and decondition
 
 A [`Model`](@ref) can be conditioned on a set of observations with [`AbstractPPL.condition`](@ref) or its alias [`|`](@ref).
@@ -133,6 +141,9 @@ Finally, the following methods can also be of use:
 ```@docs
 DynamicPPL.TestUtils.varnames
 DynamicPPL.TestUtils.posterior_mean
+DynamicPPL.TestUtils.setup_varinfos
+DynamicPPL.TestUtils.update_values!!
+DynamicPPL.TestUtils.test_values
 ```
 
 ## Advanced

src/DynamicPPL.jl

@@ -13,6 +13,7 @@ using MacroTools: MacroTools
 using ConstructionBase: ConstructionBase
 using Setfield: Setfield
 using ZygoteRules: ZygoteRules
+using LogDensityProblems: LogDensityProblems
 
 using DocStringExtensions
 
@@ -163,5 +164,6 @@ include("loglikelihoods.jl")
 include("submodel_macro.jl")
 include("test_utils.jl")
 include("transforming.jl")
+include("logdensityfunction.jl")
 
 end # module

src/abstract_varinfo.jl

@@ -519,14 +519,25 @@ end
 
 # Utilities
 """
-    unflatten(vi::AbstractVarInfo[, spl::AbstractSampler], x::AbstractVector)
+    unflatten(vi::AbstractVarInfo[, context::AbstractContext], x::AbstractVector)
 
 Return a new instance of `vi` with the values of `x` assigned to the variables.
 
-If `spl` is provided, `x` is assumed to be realizations only for variables related 
-to `spl`.
+If `context` is provided, `x` is assumed to be realizations only for variables not
+filtered out by `context`.
 """
-function unflatten end
+function unflatten(varinfo::AbstractVarInfo, context::AbstractContext, θ)
+    if hassampler(context)
+        unflatten(getsampler(context), varinfo, context, θ)
+    else
+        DynamicPPL.unflatten(varinfo, θ)
+    end
+end
+
+# TODO: deprecate this once `sampler` is no longer the main way of filtering out variables.
+function unflatten(sampler::AbstractSampler, varinfo::AbstractVarInfo, ::AbstractContext, θ)
+    return unflatten(varinfo, sampler, θ)
+end
 
 """
     tonamedtuple(vi::AbstractVarInfo)

src/contexts.jl

@@ -163,6 +163,28 @@ function setchildcontext(parent::SamplingContext, child)
     return SamplingContext(parent.rng, parent.sampler, child)
 end
 
+"""
+    hassampler(context)
+
+Return `true` if `context` has a sampler.
+"""
+hassampler(::SamplingContext) = true
+hassampler(context::AbstractContext) = hassampler(NodeTrait(context), context)
+hassampler(::IsLeaf, context::AbstractContext) = false
+hassampler(::IsParent, context::AbstractContext) = hassampler(childcontext(context))
+
+"""
+    getsampler(context)
+
+Return the sampler of the context `context`.
+
+This will traverse the context tree until it reaches the first [`SamplingContext`](@ref),
+at which point it will return the sampler of that context.
+"""
+getsampler(context::SamplingContext) = context.sampler
+getsampler(context::AbstractContext) = getsampler(NodeTrait(context), context)
+getsampler(::IsParent, context::AbstractContext) = getsampler(childcontext(context))
+
 """
     struct DefaultContext <: AbstractContext end
 

src/logdensityfunction.jl

@@ -0,0 +1,94 @@
+"""
+    LogDensityFunction
+
+A callable representing a log density function of a `model`.
+
+# Fields
+$(FIELDS)
+
+# Examples
+```jldoctest
+julia> using Distributions
+
+julia> using DynamicPPL: LogDensityFunction
+
+julia> @model function demo(x)
+           m ~ Normal()
+           x ~ Normal(m, 1)
+       end
+demo (generic function with 2 methods)
+
+julia> model = demo(1.0);
+
+julia> f = LogDensityFunction(model);
+
+julia> # It implements the interface of LogDensityProblems.jl.
+       using LogDensityProblems
+
+julia> LogDensityProblems.logdensity(f, [0.0])
+-2.3378770664093453
+
+julia> LogDensityProblems.dimension(f)
+1
+
+julia> # By default it uses `VarInfo` under the hood, but this is not necessary.
+       f = LogDensityFunction(model, SimpleVarInfo(model));
+
+julia> LogDensityProblems.logdensity(f, [0.0])
+-2.3378770664093453
+```
+"""
+struct LogDensityFunction{V,M,C}
+    "varinfo used for evaluation"
+    varinfo::V
+    "model used for evaluation"
+    model::M
+    "context used for evaluation"
+    context::C
+end
+
+# TODO: Deprecate.
+function LogDensityFunction(
+    varinfo::AbstractVarInfo,
+    model::Model,
+    sampler::AbstractSampler,
+    context::AbstractContext,
+)
+    return LogDensityFunction(varinfo, model, SamplingContext(sampler, context))
+end
+
+function LogDensityFunction(
+    model::Model,
+    varinfo::AbstractVarInfo=VarInfo(model),
+    context::AbstractContext=DefaultContext(),
+)
+    return LogDensityFunction(varinfo, model, context)
+end
+
+# HACK: heavy usage of `AbstractSampler` for, well, _everything_, is being phased out. In the mean time
+# we need to define these annoying methods to ensure that we stay compatible with everything.
+getsampler(f::LogDensityFunction) = getsampler(f.context)
+hassampler(f::LogDensityFunction) = hassampler(f.context)
+
+_get_indexer(ctx::AbstractContext) = _get_indexer(NodeTrait(ctx), ctx)
+_get_indexer(ctx::SamplingContext) = ctx.sampler
+_get_indexer(::IsParent, ctx::AbstractContext) = _get_indexer(childcontext(ctx))
+_get_indexer(::IsLeaf, ctx::AbstractContext) = Colon()
+
+"""
+    getparams(f::LogDensityFunction)
+
+Return the parameters of the wrapped varinfo as a vector.
+"""
+getparams(f::LogDensityFunction) = f.varinfo[_get_indexer(f.context)]
+
+# LogDensityProblems interface
+function LogDensityProblems.logdensity(f::LogDensityFunction, θ::AbstractVector)
+    vi_new = unflatten(f.varinfo, f.context, θ)
+    return getlogp(last(evaluate!!(f.model, vi_new, f.context)))
+end
+function LogDensityProblems.capabilities(::Type{<:LogDensityFunction})
+    return LogDensityProblems.LogDensityOrder{0}()
+end
+# TODO: should we instead implement and call on `length(f.varinfo)` (at least in the cases where no sampler is involved)?
+LogDensityProblems.dimension(f::LogDensityFunction) = length(getparams(f))

src/simple_varinfo.jl

@@ -248,7 +248,7 @@ function SimpleVarInfo{T}(
     return SimpleVarInfo(values, convert(T, getlogp(vi)))
 end
 
-unflatten(svi::SimpleVarInfo, spl, x::AbstractVector) = unflatten(svi, x)
+unflatten(svi::SimpleVarInfo, spl::AbstractSampler, x::AbstractVector) = unflatten(svi, x)
 function unflatten(svi::SimpleVarInfo, x::AbstractVector)
     return Setfield.@set svi.values = unflatten(svi.values, x)
 end

src/test_utils.jl

@@ -31,6 +31,50 @@ function varname_leaves(vn::VarName, val::AbstractArray)
     )
 end
 
+"""
+    update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)
+
+Return instance similar to `vi` but with `vns` set to values from `vals`.
+"""
+function update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)
+    for vn in vns
+        vi = DynamicPPL.setindex!!(vi, get(vals, vn), vn)
+    end
+    return vi
+end
+
+"""
+    test_values(vi::AbstractVarInfo, vals::NamedTuple, vns)
+
+Test that `vi[vn]` corresponds to the correct value in `vals` for every `vn` in `vns`.
+"""
+function test_values(vi::AbstractVarInfo, vals::NamedTuple, vns; isequal=isequal, kwargs...)
+    for vn in vns
+        @test isequal(vi[vn], get(vals, vn); kwargs...)
+    end
+end
+
+"""
+    setup_varinfos(model::Model, example_values::NamedTuple, varnames)
+
+Return a tuple of instances for different implementations of `AbstractVarInfo` with
+each `vi`, supposedly, satisfying `vi[vn] == get(example_values, vn)` for `vn` in `varnames`.
+"""
+function setup_varinfos(model::Model, example_values::NamedTuple, varnames)
+    # <:VarInfo
+    vi_untyped = VarInfo()
+    model(vi_untyped)
+    vi_typed = DynamicPPL.TypedVarInfo(vi_untyped)
+    # <:SimpleVarInfo
+    svi_typed = SimpleVarInfo(example_values)
+    svi_untyped = SimpleVarInfo(OrderedDict())
+
+    return map((vi_untyped, vi_typed, svi_typed, svi_untyped)) do vi
+        # Set them all to the same values.
+        update_values!!(vi, example_values, varnames)
+    end
+end
+
 """
     logprior_true(model, args...)
 

src/varinfo.jl

@@ -140,7 +140,7 @@ VarInfo(model::Model, args...) = VarInfo(Random.GLOBAL_RNG, model, args...)
 unflatten(vi::VarInfo, x::AbstractVector) = unflatten(vi, SampleFromPrior(), x)
 
 # TODO: deprecate.
-unflatten(vi::VarInfo, spl, x::AbstractVector) = VarInfo(vi, spl, x)
+unflatten(vi::VarInfo, spl::AbstractSampler, x::AbstractVector) = VarInfo(vi, spl, x)
 
 # without AbstractSampler
 function VarInfo(rng::Random.AbstractRNG, model::Model, context::AbstractContext)