Include more options

Author: penelopeysm

docs/src/api.md

@@ -505,3 +505,11 @@ There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_va
 DynamicPPL.Experimental.determine_suitable_varinfo
 DynamicPPL.Experimental.is_suitable_varinfo
 ```
+
+### Converting VarInfo to chains
+
+The following function is useful for package developers seeking to extend DynamicPPL:
+
+```@docs
+DynamicPPL.varinfos_to_chains
+```

ext/DynamicPPLMCMCChainsExt.jl

@@ -36,30 +36,51 @@ function chain_sample_to_varname_dict(
     return d
 end
 
+struct ModelEvaluationData{
+    Tval<:AbstractDict,Tpr<:Union{AbstractFloat,Missing},Tlk<:Union{AbstractFloat,Missing}
+}
+    values::Tval
+    logprior::Tpr
+    loglikelihood::Tlk
+end
+
 """
     DynamicPPL.varinfos_to_chains(
         ::Type{MCMCChains.Chains},
-        model::Model,
         varinfos::AbstractArray{<:AbstractVarInfo},
+        model::Union{DynamicPPL.Model,Nothing};
         include_colon_eq::Bool=true,
+        include_log_probs::Bool=true,
     )
 
 Convert an array of `VarInfo`s to an `MCMCChains.Chains` object.
 """
 function DynamicPPL.varinfos_to_chains(
     ::Type{MCMCChains.Chains},
-    model::DynamicPPL.Model,
     varinfos::AbstractMatrix{<:DynamicPPL.AbstractVarInfo},
+    model::Union{DynamicPPL.Model,Nothing};
     include_colon_eq::Bool=true,
+    include_log_probs::Bool=true,
 )
     all_vn_leaves = DynamicPPL.OrderedCollections.OrderedSet{DynamicPPL.VarName}()
-    # Re-evaluate model
-    param_dicts = map(varinfos) do vi
-        # Dict{VarName, Any}
-        full_vals = DynamicPPL.values_as_in_model(model, include_colon_eq, vi)
+    evaluation_data = map(varinfos) do vi
+        # Re-evaluate model if needed
+        new_vi = if model isa DynamicPPL.Model
+            accs = (
+                DynamicPPL.LogPriorAccumulator(),
+                DynamicPPL.LogJacobianAccumulator(),
+                DynamicPPL.LogLikelihoodAccumulator(),
+                DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),
+            )
+            vi = DynamicPPL.setaccs!!(vi, accs)
+            last(DynamicPPL.evaluate!!(model, vi))
+        else
+            vi
+        end
+        full_vals = DynamicPPL.getacc(new_vi, Val(:ValuesAsInModel)).values
         # Separate into individual VarNames.
         vn_leaves_and_vals = if isempty(full_vals)
-            Tuple{VarName,Any}[]
+            Tuple{DynamicPPL.VarName,Any}[]
         else
             iters = map(
                 AbstractPPL.varname_and_value_leaves,
@@ -73,29 +94,58 @@ function DynamicPPL.varinfos_to_chains(
         for vn_leaf in vn_leaves
             push!(all_vn_leaves, vn_leaf)
         end
-        return DynamicPPL.OrderedCollections.OrderedDict(zip(vn_leaves, vals))
+        vn_leaves_dict = DynamicPPL.OrderedCollections.OrderedDict(zip(vn_leaves, vals))
+        logprior = if include_log_probs && DynamicPPL.hasacc(new_vi, Val(:LogPrior))
+            DynamicPPL.getacc(new_vi, Val(:LogPrior)).logp
+        else
+            missing
+        end
+        loglikelihood =
+            if include_log_probs && DynamicPPL.hasacc(new_vi, Val(:LogLikelihood))
+                DynamicPPL.getacc(new_vi, Val(:LogLikelihood)).logp
+            else
+                missing
+            end
+        return ModelEvaluationData(vn_leaves_dict, logprior, loglikelihood)
     end
     vn_leaves = collect(all_vn_leaves)
     vals = [
-        get(param_dicts[i, j], key, missing) for i in eachindex(axes(param_dicts, 1)),
-        key in vn_leaves, j in eachindex(axes(param_dicts, 2))
+        get(evaluation_data[i, j].values, key, missing) for
+        i in eachindex(axes(evaluation_data, 1)), key in vn_leaves,
+        j in eachindex(axes(evaluation_data, 2))
     ]
     symbols = map(Symbol, vn_leaves)
     info = (
         varname_to_symbol=DynamicPPL.OrderedCollections.OrderedDict(
             zip(all_vn_leaves, symbols)
         ),
     )
-    return MCMCChains.Chains(MCMCChains.concretize(vals), symbols; info=info)
+    name_map = NamedTuple()
+    if include_log_probs
+        logpriors = map(e -> e.logprior, evaluation_data)
+        loglikelihoods = map(e -> e.loglikelihood, evaluation_data)
+        logjoints = map(e -> e.logprior + e.loglikelihood, evaluation_data)
+        lps = permutedims(stack([logpriors, loglikelihoods, logjoints]), (1, 3, 2))
+        vals = hcat(vals, lps)
+        lp_symbols = [:logprior, :loglikelihood, :lp]
+        append!(symbols, lp_symbols)
+        name_map = (internals=lp_symbols,)
+    end
+    return MCMCChains.Chains(MCMCChains.concretize(vals), symbols, name_map; info=info)
 end
 function DynamicPPL.varinfos_to_chains(
     ::Type{MCMCChains.Chains},
-    model::DynamicPPL.Model,
     varinfos::AbstractVector{<:DynamicPPL.AbstractVarInfo},
+    model::Union{DynamicPPL.Model,Nothing};
     include_colon_eq::Bool=true,
+    include_log_probs::Bool=true,
 )
     return DynamicPPL.varinfos_to_chains(
-        MCMCChains.Chains, model, hcat(varinfos), include_colon_eq
+        MCMCChains.Chains,
+        hcat(varinfos),
+        model;
+        include_colon_eq=include_colon_eq,
+        include_log_probs=include_log_probs,
     )
 end
 

src/abstract_varinfo.jl

@@ -1186,19 +1186,34 @@ end
 """
     varinfos_to_chains(
         Tout::Type{<:AbstractChains},
-        model::DynamicPPL.Model,
         varinfos::AbstractArray{<:AbstractVarInfo},
-        include_colon_eq::Bool=true
+        model::Union{DynamicPPL.Model,Nothing};
+        include_colon_eq::Bool=true,
+        include_log_probs::Bool=true,
     )
 
 Convert an array of `varinfos` to a chains object of type `Tout`.
 
-The `model` is required in order to account for cases where the varinfo is linked and
-re-evaluation is required. For example, this is the case when the support of a distribution
-depends on other random variables.
+In many cases, it is necessary to re-evaluate the model with each VarInfo to obtain correct
+parameter values as well as log probabilities. This is especially true when using linked
+VarInfos, i.e., when variables have been transformed to unconstrained space, and if this is
+not done, subtle correctness bugs may arise: see, e.g.,
+https://github.com/TuringLang/Turing.jl/issues/2195
+
+There is one case where re-evaluation is not necessary, which is when the VarInfos all
+already contain `DynamicPPL.ValuesAsInModelAccumulator`. This accumulator stores values
+as seen during the model evaluation, so the values can be simply read off. In this case,
+`model` can be set to `nothing`, and no re-evaluation will be performed. However, it is the
+caller's responsibility to ensure that `ValuesAsInModelAccumulator` is indeed
+present.
 
 `include_colon_eq` indicates whether to include variables on the left-hand side of `:=`.
 
+`include_log_probs` indicates whether to include log probabilities (log prior, log
+likelihood, and log joint) in the resulting chains object. By default these are included
+(and also recalculated if a model is provided; this incurs almost no additional cost on top
+of the original re-evaluation).
+
 This function is not implemented here but rather in package extensions for individual chains
 packages.
 """

test/ext/DynamicPPLMCMCChainsExt.jl

@@ -15,59 +15,98 @@ using LinearAlgebra: I
     @test mean(chain_generated) ≈ 0 atol = 0.1
 
     @testset "varinfos_to_chains" begin
-        @model function f()
+        @model function f(z)
             x ~ Normal()
             y := x + 1
-            return z ~ MvNormal(zeros(3), I)
+            return z ~ Normal(y)
         end
 
-        model = f()
+        z = 1.0
+        model = f(z)
 
         @testset "vector" begin
             vis = [VarInfo(model) for _ in 1:50]
-            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, vis)
+            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, vis, model)
             @test c isa MCMCChains.Chains
-            @test size(c) == (50, 5, 1)
+            @test size(c, 1) == 50
+            @test size(c, 3) == 1
+            @test Set(c.name_map.parameters) == Set([:x, :y])
+            @test Set(c.name_map.internals) == Set([:logprior, :loglikelihood, :lp])
+            @test logpdf.(Normal(), c[:x]) ≈ c[:logprior]
+            @test c.info.varname_to_symbol[@varname(x)] == :x
+            @test c.info.varname_to_symbol[@varname(y)] == :y
+        end
+
+        @testset "vector, no reevaluation" begin
+            # check that it throws an error without VAIMAcc
+            vis = [VarInfo(model) for _ in 1:50]
+            @test_throws ErrorException DynamicPPL.varinfos_to_chains(
+                MCMCChains.Chains, vis, nothing
+            )
+            # and that it works with VAIMAcc
+            vi = DynamicPPL.VarInfo(model)
+            vi = DynamicPPL.setaccs!!(vi, (DynamicPPL.ValuesAsInModelAccumulator(true),))
+            vis = [last(DynamicPPL.init!!(model, vi)) for _ in 1:50]
+            c = DynamicPPL.varinfos_to_chains(
+                MCMCChains.Chains, vis, nothing; include_log_probs=false
+            )
+            @test c isa MCMCChains.Chains
+            @test size(c, 1) == 50
+            @test size(c, 3) == 1
+            @test Set(c.name_map.parameters) == Set([:x, :y])
             @test c.info.varname_to_symbol[@varname(x)] == :x
             @test c.info.varname_to_symbol[@varname(y)] == :y
-            @test c.info.varname_to_symbol[@varname(z[1])] == Symbol("z[1]")
-            @test c.info.varname_to_symbol[@varname(z[2])] == Symbol("z[2]")
-            @test c.info.varname_to_symbol[@varname(z[3])] == Symbol("z[3]")
         end
 
         @testset "vector, without include_colon_eq" begin
             vis = [VarInfo(model) for _ in 1:50]
-            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, vis, false)
+            c = DynamicPPL.varinfos_to_chains(
+                MCMCChains.Chains, vis, model; include_colon_eq=false
+            )
             @test c isa MCMCChains.Chains
-            @test size(c) == (50, 4, 1)
+            @test size(c, 1) == 50
+            @test size(c, 3) == 1
+            @test Set(c.name_map.parameters) == Set([:x])
+            @test Set(c.name_map.internals) == Set([:logprior, :loglikelihood, :lp])
+            @test logpdf.(Normal(), c[:x]) ≈ c[:logprior]
             @test c.info.varname_to_symbol[@varname(x)] == :x
-            @test c.info.varname_to_symbol[@varname(z[1])] == Symbol("z[1]")
-            @test c.info.varname_to_symbol[@varname(z[2])] == Symbol("z[2]")
-            @test c.info.varname_to_symbol[@varname(z[3])] == Symbol("z[3]")
+        end
+
+        @testset "vector, without logprobs" begin
+            vis = [VarInfo(model) for _ in 1:50]
+            c = DynamicPPL.varinfos_to_chains(
+                MCMCChains.Chains, vis, model; include_log_probs=false
+            )
+            @test c isa MCMCChains.Chains
+            @test size(c, 1) == 50
+            @test size(c, 3) == 1
+            @test Set(c.name_map.parameters) == Set([:x, :y])
+            @test !haskey(c.name_map, :internals)
         end
 
         @testset "Different VarInfo type" begin
             vis = [SimpleVarInfo(model) for _ in 1:50]
-            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, vis)
+            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, vis, model)
             @test c isa MCMCChains.Chains
-            @test size(c) == (50, 5, 1)
+            @test size(c, 1) == 50
+            @test size(c, 3) == 1
+            @test Set(c.name_map.parameters) == Set([:x, :y])
+            @test Set(c.name_map.internals) == Set([:logprior, :loglikelihood, :lp])
+            @test logpdf.(Normal(), c[:x]) ≈ c[:logprior]
             @test c.info.varname_to_symbol[@varname(x)] == :x
             @test c.info.varname_to_symbol[@varname(y)] == :y
-            @test c.info.varname_to_symbol[@varname(z[1])] == Symbol("z[1]")
-            @test c.info.varname_to_symbol[@varname(z[2])] == Symbol("z[2]")
-            @test c.info.varname_to_symbol[@varname(z[3])] == Symbol("z[3]")
         end
 
         @testset "matrix" begin
             vis = [VarInfo(model) for _ in 1:50, _ in 1:3]
-            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, vis)
+            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, vis, model)
             @test c isa MCMCChains.Chains
             @test size(c) == (50, 5, 3)
+            @test Set(c.name_map.parameters) == Set([:x, :y])
+            @test Set(c.name_map.internals) == Set([:logprior, :loglikelihood, :lp])
+            @test logpdf.(Normal(), c[:x]) ≈ c[:logprior]
             @test c.info.varname_to_symbol[@varname(x)] == :x
             @test c.info.varname_to_symbol[@varname(y)] == :y
-            @test c.info.varname_to_symbol[@varname(z[1])] == Symbol("z[1]")
-            @test c.info.varname_to_symbol[@varname(z[2])] == Symbol("z[2]")
-            @test c.info.varname_to_symbol[@varname(z[3])] == Symbol("z[3]")
         end
     end
 end