Implement varinfos_to_chains function

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,9 @@
 # DynamicPPL Changelog
 
+## 0.38.2
+
+Added a new exported function, `DynamicPPL.varinfos_to_chains`, which automatically converts a collection of VarInfos to a given Chains type.
+
 ## 0.38.1
 
 Added `from_linked_vec_transform` and `from_vec_transform` methods for `ProductNamedTupleDistribution`.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.38.1"
+version = "0.38.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/DynamicPPLMCMCChainsExt.jl

@@ -36,6 +36,69 @@ function chain_sample_to_varname_dict(
     return d
 end
 
+"""
+    DynamicPPL.varinfos_to_chains(
+        ::Type{MCMCChains.Chains},
+        model::Model,
+        varinfos::AbstractArray{<:AbstractVarInfo},
+        include_colon_eq::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},
+    include_colon_eq::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)
+        # Separate into individual VarNames.
+        vn_leaves_and_vals = if isempty(full_vals)
+            Tuple{VarName,Any}[]
+        else
+            iters = map(
+                AbstractPPL.varname_and_value_leaves,
+                keys(full_vals),
+                values(full_vals),
+            )
+            mapreduce(collect, vcat, iters)
+        end
+        vn_leaves = map(first, vn_leaves_and_vals)
+        vals = map(last, vn_leaves_and_vals)
+        for vn_leaf in vn_leaves
+            push!(all_vn_leaves, vn_leaf)
+        end
+        return DynamicPPL.OrderedCollections.OrderedDict(zip(vn_leaves, vals))
+    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))
+    ]
+    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)
+end
+function DynamicPPL.varinfos_to_chains(
+    ::Type{MCMCChains.Chains},
+    model::DynamicPPL.Model,
+    varinfos::AbstractVector{<:DynamicPPL.AbstractVarInfo},
+    include_colon_eq::Bool=true,
+)
+    return DynamicPPL.varinfos_to_chains(
+        MCMCChains.Chains, model, hcat(varinfos), include_colon_eq
+    )
+end
+
 """
     predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 

src/DynamicPPL.jl

@@ -126,6 +126,7 @@ export AbstractVarInfo,
     prefix,
     returned,
     to_submodel,
+    varinfos_to_chains,
     # Convenience macros
     @addlogprob!,
     value_iterator_from_chain,

src/abstract_varinfo.jl

@@ -1182,3 +1182,24 @@ function linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::Var
     f_to_internal = to_internal_transform(varinfo, vn)
     return f_to_internal ∘ f_from_linked_internal
 end
+
+"""
+    varinfos_to_chains(
+        Tout::Type{<:AbstractChains},
+        model::DynamicPPL.Model,
+        varinfos::AbstractArray{<:AbstractVarInfo},
+        include_colon_eq::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.
+
+`include_colon_eq` indicates whether to include variables on the left-hand side of `:=`.
+
+This function is not implemented here but rather in package extensions for individual chains
+packages.
+"""
+function varinfos_to_chains end

test/ext/DynamicPPLMCMCChainsExt.jl

@@ -1,3 +1,5 @@
+using LinearAlgebra: I
+
 @testset "DynamicPPLMCMCChainsExt" begin
     @model demo() = x ~ Normal()
     model = demo()
@@ -11,6 +13,40 @@
     chain_generated = @test_nowarn returned(model, chain)
     @test size(chain_generated) == (1000, 1)
     @test mean(chain_generated) ≈ 0 atol = 0.1
+
+    @testset "varinfos_to_chains" begin
+        @model function f()
+            x ~ Normal()
+            y ~ Normal(x)
+            return z ~ MvNormal(zeros(3), I)
+        end
+
+        model = f()
+
+        @testset "vector" begin
+            vis = [VarInfo(model) for _ in 1:50]
+            c = DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, vis)
+            @test c isa MCMCChains.Chains
+            @test size(c) == (50, 5, 1)
+            @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)
+            @test c isa MCMCChains.Chains
+            @test size(c) == (50, 5, 3)
+            @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
 
 # test for `predict` is in `test/model.jl`

test/test_util.jl

@@ -62,35 +62,8 @@ Construct an MCMCChains.Chains object by sampling from the prior of `model` for
 `n_iters` iterations.
 """
 function make_chain_from_prior(rng::Random.AbstractRNG, model::Model, n_iters::Int)
-    # Sample from the prior
     varinfos = [VarInfo(rng, model) for _ in 1:n_iters]
-    # Extract all varnames found in any dictionary. Doing it this way guards
-    # against the possibility of having different varnames in different
-    # dictionaries, e.g. for models that have dynamic variables / array sizes
-    varnames = OrderedSet{VarName}()
-    # Convert each varinfo into an OrderedDict of vns => params.
-    # We have to use varname_and_value_leaves so that each parameter is a scalar
-    dicts = map(varinfos) do t
-        vals = DynamicPPL.values_as(t, OrderedDict)
-        iters = map(AbstractPPL.varname_and_value_leaves, keys(vals), values(vals))
-        tuples = mapreduce(collect, vcat, iters)
-        # The following loop is a replacement for:
-        #     push!(varnames, map(first, tuples)...)
-        # which causes a stack overflow if `map(first, tuples)` is too large.
-        # Unfortunately there isn't a union() function for OrderedSet.
-        for vn in map(first, tuples)
-            push!(varnames, vn)
-        end
-        OrderedDict(tuples)
-    end
-    # Convert back to list
-    varnames = collect(varnames)
-    # Construct matrix of values
-    vals = [get(dict, vn, missing) for dict in dicts, vn in varnames]
-    # Construct dict of varnames -> symbol
-    vn_to_sym_dict = Dict(zip(varnames, map(Symbol, varnames)))
-    # Construct and return the Chains object
-    return Chains(vals, varnames; info=(; varname_to_symbol=vn_to_sym_dict))
+    return DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, varinfos)
 end
 function make_chain_from_prior(model::Model, n_iters::Int)
     return make_chain_from_prior(Random.default_rng(), model, n_iters)