Add better method to convert Trandition{NamedTuple} to MCMCChains

Author: cpfiffer

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.5.6"
+version = "0.5.7"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/mcmcchains-connect.jl

@@ -2,7 +2,7 @@ import .MCMCChains: Chains
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    ts::Vector{<:Transition},
+    ts::Vector{<:Transition{<:AbstractArray}},
     model::DensityModel,
     sampler::MHSampler,
     state,
@@ -28,6 +28,40 @@ function AbstractMCMC.bundle_samples(
     return Chains(vals, param_names, (internals = [:lp],))
 end
 
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:Transition{<:NamedTuple}},
+    model::DensityModel,
+    sampler::MHSampler,
+    state,
+    chain_type::Type{Chains};
+    param_names=missing,
+    kwargs...
+)
+    # Convert to a Vector{NamedTuple} first
+    nts = AbstractMCMC.bundle_samples(ts, model, sampler, state, Vector{NamedTuple}; param_names=param_names, kwargs...)
+
+    # Get all the keys
+    all_keys = unique(mapreduce(collect∘keys, vcat, nts))
+
+    # Preallocate array
+    # vals = []
+
+    # Push linearized draws onto array
+    trygetproperty(thing, key) = key in keys(thing) ? getproperty(thing, key) : missing
+    vals = map(nt -> [trygetproperty(nt, k) for k in all_keys], nts)
+
+    # Check if we received any parameter names.
+    if ismissing(param_names)
+        param_names = all_keys
+    else
+        # Generate new array to be thread safe.
+        param_names = Symbol.(param_names)
+    end
+
+    # Bundle everything up and return a Chains struct.
+    return Chains(vals, param_names, (internals = [:lp],))
+end
+
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Vector{<:Transition}},
     model::DensityModel,