use NamedTuple

Author: sunxd3

ext/DynamicPPLMCMCChainsExt.jl

@@ -147,9 +147,9 @@ function _params_to_array(predictive_samples)
     names_set = DynamicPPL.OrderedCollections.OrderedSet{DynamicPPL.VarName}()
 
     dicts = map(predictive_samples) do t
-        nms_and_vs = t[:values]
-        nms = map(first, nms_and_vs)
-        vs = map(last, nms_and_vs)
+        varname_and_values = t.varname_and_values
+        nms = map(first, varname_and_values)
+        vs = map(last, varname_and_values)
         for nm in nms
             push!(names_set, nm)
         end
@@ -164,15 +164,11 @@ function _params_to_array(predictive_samples)
     return names, vals
 end
 
-function _bundle_predictive_samples(
-    predictive_samples::AbstractArray{
-        <:DynamicPPL.OrderedCollections.OrderedDict{Symbol,Any}
-    },
-)
+function _bundle_predictive_samples(predictive_samples)
     varnames, vals = _params_to_array(predictive_samples)
     varnames_symbol = map(Symbol, varnames)
     extra_params = [:lp]
-    extra_values = reshape([t[:logp] for t in predictive_samples], :, 1)
+    extra_values = reshape([t.logp for t in predictive_samples], :, 1)
     nms = [varnames_symbol; extra_params]
     parray = hcat(vals, extra_values)
     parray = MCMCChains.concretize(parray)

src/model.jl

@@ -1223,15 +1223,13 @@ function predict(
     varinfos::AbstractArray{<:AbstractVarInfo};
     include_all=false,
 )
-    predictive_samples = similar(varinfos, OrderedDict{Symbol,Any})
+    predictive_samples = similar(varinfos, NamedTuple{(:varname_and_values, :logp)})
     for i in eachindex(varinfos)
         model(rng, varinfos[i], SampleFromPrior())
         vals = values_as_in_model(model, varinfos[i])
         iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
-        params = mapreduce(collect, vcat, iters)
-        predictive_samples[i] = OrderedDict(
-            :values => params, :logp => getlogp(varinfos[i])
-        )
+        params = mapreduce(collect, vcat, iters) # returns a vector of tuples (varname, value)
+        predictive_samples[i] = (varname_and_values=params, logp=getlogp(varinfos[i]))
     end
     return predictive_samples
 end