Add from_chains as well

Author: penelopeysm

docs/src/api.md

@@ -515,3 +515,11 @@ This can be accomplished with the following:
 DynamicPPL.ParamsWithStats
 DynamicPPL.to_chains
 ```
+
+Furthermore, one can convert chains back into a collection of parameter dictionaries and/or stats with:
+
+```@docs
+DynamicPPL.from_chains
+```
+
+This is useful if you want to use the result of a chain in further model evaluations.

ext/DynamicPPLMCMCChainsExt.jl

@@ -106,6 +106,43 @@ function DynamicPPL.to_chains(
     return DynamicPPL.to_chains(MCMCChains.Chains, hcat(ps))
 end
 
+function DynamicPPL.from_chains(
+    ::Type{T}, chain::MCMCChains.Chains
+) where {T<:AbstractDict{<:DynamicPPL.VarName}}
+    idxs = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+    matrix = map(idxs) do (sample_idx, chain_idx)
+        d = T()
+        for vn in DynamicPPL.varnames(chain)
+            d[vn] = DynamicPPL.getindex_varname(chain, sample_idx, vn, chain_idx)
+        end
+        d
+    end
+    return matrix
+end
+function DynamicPPL.from_chains(::Type{NamedTuple}, chain::MCMCChains.Chains)
+    idxs = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+    matrix = map(idxs) do (sample_idx, chain_idx)
+        get(chain[sample_idx, :, chain_idx], keys(chain); flatten=true)
+    end
+    return matrix
+end
+function DynamicPPL.from_chains(
+    ::Type{DynamicPPL.ParamsWithStats}, chain::MCMCChains.Chains
+)
+    idxs = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+    internals_chain = MCMCChains.get_sections(chain, :internals)
+    params = DynamicPPL.from_chains(
+        DynamicPPL.OrderedCollections.OrderedDict{DynamicPPL.VarName,eltype(chain.value)},
+        chain,
+    )
+    stats = DynamicPPL.from_chains(NamedTuple, internals_chain)
+    return map(idxs) do (sample_idx, chain_idx)
+        DynamicPPL.ParamsWithStats(
+            params[sample_idx, chain_idx], stats[sample_idx, chain_idx]
+        )
+    end
+end
+
 """
     predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 

src/to_chains.jl

@@ -42,10 +42,16 @@ 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}
+struct ParamsWithStats{P<:OrderedDict{<:VarName,<:Any},S<:NamedTuple}
     params::P
     stats::S
 
+    function ParamsWithStats(
+        params::P, stats::S
+    ) where {P<:OrderedDict{<:VarName,<:Any},S<:NamedTuple}
+        return new{P,S}(params, stats)
+    end
+
     function ParamsWithStats(
         varinfo::AbstractVarInfo,
         model::DynamicPPL.Model,
@@ -113,11 +119,40 @@ maybe_to_typed_varinfo(vi::AbstractVarInfo) = vi
     to_chains(
         Tout::Type{<:AbstractChains},
         params_and_stats::AbstractArray{<:ParamsWithStats}
-    )
+    )::Tout
 
 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
+
+"""
+    from_chains(
+        ::Type{Tout},
+        chain::AbstractChains
+    )::AbstractMatrix{<:Tout}
+
+Convert a chains object to an array of size (niters * nchains) with element type `Tout`.
+
+Note that even if `chain` contains only a single chain, this is still returned as a matrix,
+not a vector.
+
+This function is not implemented here but rather in package extensions for individual chains
+packages.
+
+Common implementations include:
+
+ - `Tout = ParamsWithStats`: obtain both parameters and statistics
+ - `Tout <: AbstractDict{<:VarName}`: obtain the parameters only (since stats are not stored
+   as `VarName`s
+ - `Tout = NamedTuple`: obtain both parameters and statistics as a NamedTuple
+
+!!! warning
+    Note that `Tout = NamedTuple` potentially causes a loss of information especially when
+    used with `MCMCChains.Chains`, since variable names are not preserved. This may lead to
+    bugs if the NamedTuple is later used for other purposes, such as evaluating a model. To
+    avoid this, you should always use something like `Tout = OrderedDict{VarName,Any}`.
+"""
+function from_chains end

test/ext/DynamicPPLMCMCChainsExt.jl

@@ -12,7 +12,7 @@
     @test size(chain_generated) == (1000, 1)
     @test mean(chain_generated) ≈ 0 atol = 0.1
 
-    @testset "varinfos_to_chains" begin
+    @testset "to_chains" begin
         @model function f(z)
             x ~ Normal()
             y := x + 1
@@ -48,6 +48,58 @@
             @test c.info.varname_to_symbol[@varname(y)] == :y
         end
     end
+
+    @testset "from_chains" begin
+        @model function f(z)
+            x ~ Normal()
+            y := x + 1
+            return z ~ Normal(y)
+        end
+
+        z = 1.0
+        model = f(z)
+        ps = [ParamsWithStats(VarInfo(model), model) for _ in 1:50]
+        c = DynamicPPL.to_chains(MCMCChains.Chains, ps)
+
+        @testset "OrderedDict" begin
+            arr_dicts = DynamicPPL.from_chains(OrderedDict{VarName,Any}, c)
+            @test size(arr_dicts) == (50, 1)
+            for i in 1:50
+                dict = arr_dicts[i, 1]
+                @test dict isa OrderedDict{VarName,Any}
+                p = ps[i].params
+                @test dict[@varname(x)] == p[@varname(x)]
+                @test dict[@varname(y)] == p[@varname(y)]
+                @test length(dict) == 2
+            end
+        end
+
+        @testset "NamedTuple" begin
+            arr_nts = DynamicPPL.from_chains(NamedTuple, c)
+            @test size(arr_nts) == (50, 1)
+            for i in 1:50
+                nt = arr_nts[i, 1]
+                @test length(nt) == 5
+                p = ps[i]
+                @test nt.x == p.params[@varname(x)]
+                @test nt.y == p.params[@varname(y)]
+                @test nt.lp == p.stats.lp
+                @test nt.logprior == p.stats.logprior
+                @test nt.loglikelihood == p.stats.loglikelihood
+            end
+        end
+
+        @testset "ParamsWithStats" begin
+            arr_pss = DynamicPPL.from_chains(ParamsWithStats, c)
+            @test size(arr_pss) == (50, 1)
+            for i in 1:50
+                new_p = arr_pss[i, 1]
+                p = ps[i]
+                @test new_p.params == p.params
+                @test new_p.stats == p.stats
+            end
+        end
+    end
 end
 
 # test for `predict` is in `test/model.jl`