Merge pull request #49 from TuringLang/csp/ntfix

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,

test/runtests.jl

@@ -73,6 +73,20 @@ include("util.jl")
         end
     end
 
+    @testset "MCMCChains" begin
+        spl1 = StaticMH([Normal(0,1), Normal(0, 1)])
+        spl2 = MetropolisHastings((μ = StaticProposal(Normal(0,1)), σ = StaticProposal(Normal(0, 1))))
+
+        chain1 = sample(model, spl1, 10_000; param_names=["μ", "σ"], chain_type=Chains)
+        chain2 = sample(model, spl2, 10_000; chain_type=Chains)
+
+        @test mean(chain1["μ"]) ≈ 0.0 atol=0.1
+        @test mean(chain1["σ"]) ≈ 1.0 atol=0.1
+
+        @test mean(chain2["μ"]) ≈ 0.0 atol=0.1
+        @test mean(chain2["σ"]) ≈ 1.0 atol=0.1
+    end
+
     @testset "Proposal styles" begin
         m1 = DensityModel(x -> logpdf(Normal(x,1), 1.0))
         m2 = DensityModel(x -> logpdf(Normal(x[1], x[2]), 1.0))