Fix tests

Author: penelopeysm

ext/DynamicPPLMCMCChainsExt.jl

@@ -180,7 +180,6 @@ function DynamicPPL.predict(
         DynamicPPL.VarInfo(),
         (
             DynamicPPL.LogPriorAccumulator(),
-            DynamicPPL.LogJacobianAccumulator(),
             DynamicPPL.LogLikelihoodAccumulator(),
             DynamicPPL.ValuesAsInModelAccumulator(false),
         ),
@@ -199,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
@@ -234,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)
 

src/to_chains.jl

@@ -71,7 +71,7 @@ struct ParamsWithStats{P<:OrderedDict{VarName,Any},S<:NamedTuple}
                 (
                     logprior=DynamicPPL.getlogprior(varinfo),
                     loglikelihood=DynamicPPL.getloglikelihood(varinfo),
-                    logjoint=DynamicPPL.getlogjoint(varinfo),
+                    lp=DynamicPPL.getlogjoint(varinfo),
                 ),
             )
         end

test/test_util.jl

@@ -62,8 +62,12 @@ 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)
-    varinfos = [VarInfo(rng, model) for _ in 1:n_iters]
-    return DynamicPPL.varinfos_to_chains(MCMCChains.Chains, model, varinfos)
+    vi = VarInfo(model)
+    vi = DynamicPPL.setaccs!!(vi, (DynamicPPL.ValuesAsInModelAccumulator(false),))
+    ps = [
+        ParamsWithStats(last(DynamicPPL.init!!(rng, model, vi)), nothing) for _ in 1:n_iters
+    ]
+    return DynamicPPL.to_chains(MCMCChains.Chains, ps)
 end
 function make_chain_from_prior(model::Model, n_iters::Int)
     return make_chain_from_prior(Random.default_rng(), model, n_iters)