Implement ParamsWithStats and to_chains functions

penelopeysm · penelopeysm · commit 833cbbf352f5 · 2025-10-23T20:37:59.000+01:00
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,5 +1,9 @@
 # DynamicPPL Changelog
 
+## 0.38.2
+
+Added a new exported struct, `DynamicPPL.ParamsWithStats`, and a corresponding function `DynamicPPL.to_chains`, which automatically converts a collection of `ParamsWithStats` to a given Chains type.
+
 ## 0.38.1
 
 Added `from_linked_vec_transform` and `from_vec_transform` methods for `ProductNamedTupleDistribution`.
diff --git a/Project.toml b/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"
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -505,3 +505,13 @@ There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_va
 DynamicPPL.Experimental.determine_suitable_varinfo
 DynamicPPL.Experimental.is_suitable_varinfo
 ```
+
+### Converting VarInfos to chains
+
+It is a fairly common operation to want to convert a collection of `VarInfo` objects into a chains object for downstream analysis.
+This can be accomplished with the following:
+
+```@docs
+DynamicPPL.ParamsWithStats
+DynamicPPL.to_chains
+```
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
@@ -36,6 +36,76 @@ function chain_sample_to_varname_dict(
     return d
 end
 
+"""
+    DynamicPPL.to_chains(
+        ::Type{MCMCChains.Chains},
+        params_and_stats::AbstractArray{<:ParamsWithStats}
+    )
+
+Convert an array of `DynamicPPL.ParamsWithStats` to an `MCMCChains.Chains` object.
+"""
+function DynamicPPL.to_chains(
+    ::Type{MCMCChains.Chains},
+    params_and_stats::AbstractMatrix{<:DynamicPPL.ParamsWithStats},
+)
+    # Handle parameters
+    all_vn_leaves = DynamicPPL.OrderedCollections.OrderedSet{DynamicPPL.VarName}()
+    split_dicts = map(params_and_stats) do ps
+        # Separate into individual VarNames.
+        vn_leaves_and_vals = if isempty(ps.params)
+            Tuple{DynamicPPL.VarName,Any}[]
+        else
+            iters = map(
+                AbstractPPL.varname_and_value_leaves,
+                keys(ps.params),
+                values(ps.params),
+            )
+            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
+        DynamicPPL.OrderedCollections.OrderedDict(zip(vn_leaves, vals))
+    end
+    vn_leaves = collect(all_vn_leaves)
+    param_vals = [
+        get(split_dicts[i, j], key, missing) for i in eachindex(axes(split_dicts, 1)),
+        key in vn_leaves, j in eachindex(axes(split_dicts, 2))
+    ]
+    param_symbols = map(Symbol, vn_leaves)
+    # Handle statistics
+    stat_keys = DynamicPPL.OrderedCollections.OrderedSet{Symbol}()
+    for ps in params_and_stats
+        for k in keys(ps.stats)
+            push!(stat_keys, k)
+        end
+    end
+    stat_keys = collect(stat_keys)
+    stat_vals = [
+        get(params_and_stats[i, j].stats, key, missing) for
+        i in eachindex(axes(params_and_stats, 1)), key in stat_keys,
+        j in eachindex(axes(params_and_stats, 2))
+    ]
+    # Construct name map and info
+    name_map = (internals=stat_keys,)
+    info = (
+        varname_to_symbol=DynamicPPL.OrderedCollections.OrderedDict(
+            zip(all_vn_leaves, param_symbols)
+        ),
+    )
+    # Concatenate parameter and statistic values
+    vals = cat(param_vals, stat_vals; dims=2)
+    symbols = vcat(param_symbols, stat_keys)
+    return MCMCChains.Chains(MCMCChains.concretize(vals), symbols, name_map; info=info)
+end
+function DynamicPPL.to_chains(
+    ::Type{MCMCChains.Chains}, ps::AbstractVector{<:DynamicPPL.ParamsWithStats}
+)
+    return DynamicPPL.to_chains(MCMCChains.Chains, hcat(ps))
+end
+
 """
     predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 
@@ -110,7 +180,6 @@ function DynamicPPL.predict(
         DynamicPPL.VarInfo(),
         (
             DynamicPPL.LogPriorAccumulator(),
-            DynamicPPL.LogJacobianAccumulator(),
             DynamicPPL.LogLikelihoodAccumulator(),
             DynamicPPL.ValuesAsInModelAccumulator(false),
         ),
@@ -129,23 +198,9 @@ function DynamicPPL.predict(
             varinfo,
             DynamicPPL.InitFromParams(values_dict, DynamicPPL.InitFromPrior()),
         )
-        vals = DynamicPPL.getacc(varinfo, Val(:ValuesAsInModel)).values
-        varname_vals = mapreduce(
-            collect,
-            vcat,
-            map(AbstractPPL.varname_and_value_leaves, keys(vals), values(vals)),
-        )
-
-        return (varname_and_values=varname_vals, logp=DynamicPPL.getlogjoint(varinfo))
+        DynamicPPL.ParamsWithStats(varinfo, nothing)
     end
-
-    chain_result = reduce(
-        MCMCChains.chainscat,
-        [
-            _predictive_samples_to_chains(predictive_samples[:, chain_idx]) for
-            chain_idx in 1:size(predictive_samples, 2)
-        ],
-    )
+    chain_result = DynamicPPL.to_chains(MCMCChains.Chains, predictive_samples)
     parameter_names = if include_all
         MCMCChains.names(chain_result, :parameters)
     else
@@ -164,45 +219,6 @@ function DynamicPPL.predict(
     )
 end
 
-function _predictive_samples_to_arrays(predictive_samples)
-    variable_names_set = DynamicPPL.OrderedCollections.OrderedSet{DynamicPPL.VarName}()
-
-    sample_dicts = map(predictive_samples) do sample
-        varname_value_pairs = sample.varname_and_values
-        varnames = map(first, varname_value_pairs)
-        values = map(last, varname_value_pairs)
-        for varname in varnames
-            push!(variable_names_set, varname)
-        end
-
-        return DynamicPPL.OrderedCollections.OrderedDict(zip(varnames, values))
-    end
-
-    variable_names = collect(variable_names_set)
-    variable_values = [
-        get(sample_dicts[i], key, missing) for i in eachindex(sample_dicts),
-        key in variable_names
-    ]
-
-    return variable_names, variable_values
-end
-
-function _predictive_samples_to_chains(predictive_samples)
-    variable_names, variable_values = _predictive_samples_to_arrays(predictive_samples)
-    variable_names_symbols = map(Symbol, variable_names)
-
-    internal_parameters = [:lp]
-    log_probabilities = reshape([sample.logp for sample in predictive_samples], :, 1)
-
-    parameter_names = [variable_names_symbols; internal_parameters]
-    parameter_values = hcat(variable_values, log_probabilities)
-    parameter_values = MCMCChains.concretize(parameter_values)
-
-    return MCMCChains.Chains(
-        parameter_values, parameter_names, (internals=internal_parameters,)
-    )
-end
-
 """
     returned(model::Model, chain::MCMCChains.Chains)
 
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
@@ -126,6 +126,9 @@ export AbstractVarInfo,
     prefix,
     returned,
     to_submodel,
+    # Chain construction
+    ParamsWithStats,
+    to_chains,
     # Convenience macros
     @addlogprob!,
     value_iterator_from_chain,
@@ -194,6 +197,7 @@ include("model_utils.jl")
 include("extract_priors.jl")
 include("values_as_in_model.jl")
 include("bijector.jl")
+include("to_chains.jl")
 
 include("debug_utils.jl")
 using .DebugUtils
diff --git a/src/to_chains.jl b/src/to_chains.jl
@@ -0,0 +1,116 @@
+"""
+    ParamsWithStats
+
+A struct which contains parameter values extracted from a `VarInfo`, along with any
+statistics associated with the VarInfo. The statistics are provided as a NamedTuple and are
+optional.
+
+    ParamsWithStats(
+        varinfo::AbstractVarInfo,
+        model::Model,
+        stats::NamedTuple=NamedTuple();
+        include_colon_eq::Bool=true,
+        include_log_probs::Bool=true,
+    )
+
+Generate a `ParamsWithStats` by re-evaluating the given `model` with the provided `varinfo`.
+Re-evaluation of the model is often necessary 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.
+
+`include_colon_eq` controls whether variables on the left-hand side of `:=` are included in
+the resulting parameters.
+
+`include_log_probs` controls whether log probabilities (log prior, log likelihood, and log
+joint) are added to the resulting statistics NamedTuple.
+
+    ParamsWithStats(
+        varinfo::AbstractVarInfo,
+        ::Nothing,
+        stats::NamedTuple=NamedTuple();
+        include_log_probs::Bool=true,
+    )
+
+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_log_probs` controls whether log probabilities (log prior, log likelihood, and log
+joint) are added to the resulting statistics NamedTuple.
+"""
+struct ParamsWithStats{P<:OrderedDict{VarName,Any},S<:NamedTuple}
+    params::P
+    stats::S
+
+    function ParamsWithStats(
+        varinfo::AbstractVarInfo,
+        model::DynamicPPL.Model,
+        stats::NamedTuple=NamedTuple();
+        include_colon_eq::Bool=true,
+        include_log_probs::Bool=true,
+    )
+        accs = if include_log_probs
+            (
+                DynamicPPL.LogPriorAccumulator(),
+                DynamicPPL.LogLikelihoodAccumulator(),
+                DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),
+            )
+        else
+            (DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),)
+        end
+        varinfo = DynamicPPL.setaccs!!(varinfo, accs)
+        varinfo = last(DynamicPPL.evaluate!!(model, varinfo))
+        params = DynamicPPL.getacc(varinfo, Val(:ValuesAsInModel)).values
+        if include_log_probs
+            stats = merge(
+                stats,
+                (
+                    logprior=DynamicPPL.getlogprior(varinfo),
+                    loglikelihood=DynamicPPL.getloglikelihood(varinfo),
+                    lp=DynamicPPL.getlogjoint(varinfo),
+                ),
+            )
+        end
+        return new{typeof(params),typeof(stats)}(params, stats)
+    end
+
+    function ParamsWithStats(
+        varinfo::AbstractVarInfo,
+        ::Nothing,
+        stats::NamedTuple=NamedTuple();
+        include_log_probs::Bool=true,
+    )
+        params = DynamicPPL.getacc(varinfo, Val(:ValuesAsInModel)).values
+        if include_log_probs
+            has_prior_acc = DynamicPPL.hasacc(varinfo, Val(:LogPrior))
+            has_likelihood_acc = DynamicPPL.hasacc(varinfo, Val(:LogLikelihood))
+            if has_prior_acc
+                stats = merge(stats, (logprior=DynamicPPL.getlogprior(varinfo),))
+            end
+            if has_likelihood_acc
+                stats = merge(stats, (loglikelihood=DynamicPPL.getloglikelihood(varinfo),))
+            end
+            if has_prior_acc && has_likelihood_acc
+                stats = merge(stats, (logjoint=DynamicPPL.getlogjoint(varinfo),))
+            end
+        end
+        return new{typeof(params),typeof(stats)}(params, stats)
+    end
+end
+
+"""
+    to_chains(
+        Tout::Type{<:AbstractChains},
+        params_and_stats::AbstractArray{<:ParamsWithStats}
+    )
+
+Convert an array of `ParamsWithStats` to a chains object of type `Tout`.
+
+This function is not implemented here but rather in package extensions for individual chains
+packages.
+"""
+function to_chains end
diff --git a/test/ext/DynamicPPLMCMCChainsExt.jl b/test/ext/DynamicPPLMCMCChainsExt.jl
@@ -11,6 +11,43 @@
     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(z)
+            x ~ Normal()
+            y := x + 1
+            return z ~ Normal(y)
+        end
+
+        z = 1.0
+        model = f(z)
+
+        @testset "vector" begin
+            ps = [ParamsWithStats(VarInfo(model), model) for _ in 1:50]
+            c = DynamicPPL.to_chains(MCMCChains.Chains, ps)
+            @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 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 "matrix" begin
+            ps = [ParamsWithStats(VarInfo(model), model) for _ in 1:50, _ in 1:3]
+            c = DynamicPPL.to_chains(MCMCChains.Chains, ps)
+            @test c isa MCMCChains.Chains
+            @test size(c, 1) == 50
+            @test size(c, 3) == 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
+        end
+    end
 end
 
 # test for `predict` is in `test/model.jl`
diff --git a/test/test_util.jl b/test/test_util.jl
diff --git a/test/to_chains.jl b/test/to_chains.jl
