Merge pull request #34 from TuringLang/csp/parnames

Fix parameter names when Transitions hold NamedTuples

Author: cpfiffer

Project.toml

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

src/AdvancedMH.jl

@@ -76,6 +76,25 @@ function AbstractMCMC.bundle_samples(
     return nts
 end
 
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:Transition{<:NamedTuple}},
+    model::DensityModel,
+    sampler::MHSampler,
+    state,
+    chain_type::Type{Vector{NamedTuple}};
+    param_names=missing,
+    kwargs...
+)
+    # If the element type of ts is NamedTuples, just use the names in the
+    # struct.
+
+    # Extract NamedTuples
+    nts = map(x -> merge(x.params, (lp=x.lp,)), ts)
+
+    # Return em'
+    return nts
+end
+
 function __init__()
     @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("mcmcchains-connect.jl")
     @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("structarray-connect.jl")

test/runtests.jl

@@ -84,6 +84,11 @@ using Test
         c2 = sample(m2, MetropolisHastings(p2), 100; chain_type=Vector{NamedTuple})
         c3 = sample(m3, MetropolisHastings(p3), 100; chain_type=Vector{NamedTuple})
         c4 = sample(m4, MetropolisHastings(p4), 100; chain_type=Vector{NamedTuple})
+
+        @test keys(c1[1]) == (:param_1, :lp)
+        @test keys(c2[1]) == (:param_1, :param_2, :lp)
+        @test keys(c3[1]) == (:a, :b, :lp)
+        @test keys(c4[1]) == (:param_1, :lp)
     end
 
     @testset "Initial parameters" begin