keep parameter order (#466)

* keep parameter order * Test order of parameters for get(c; section) --------- Co-authored-by: Markus Hauru <mhauru@turing.ac.uk>
TuringLang · Dec 6, 2024 · 69bbd98 · 69bbd98
1 parent aa6161b
commit 69bbd98
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/src/chains.jl b/src/chains.jl
@@ -240,7 +240,7 @@ julia> get(chn, :param_1; flatten=true)
 ```
 """
 function Base.get(c::Chains, vs::Vector{Symbol}; flatten=false)
-    pairs = Dict()
+    pairs = OrderedCollections.OrderedDict()
     for v in vs
         syms = namesingroup(c, v)
         len = length(syms)
@@ -289,7 +289,7 @@ function Base.get(
     section::Union{Symbol,AbstractVector{Symbol}},
     flatten = false
 )
-    names = Set(Symbol[])
+    names = OrderedCollections.OrderedSet(Symbol[])
     regex = r"[^\[]*"
     _section = section isa Symbol ? (section,) : section
     for v in _section

diff --git a/src/utils.jl b/src/utils.jl
@@ -56,7 +56,7 @@ function cummean(x::AbstractVector)
     return y
 end
 
-function _dict2namedtuple(d::Dict)
+function _dict2namedtuple(d::AbstractDict)
     t_keys = ntuple(x -> Symbol(collect(keys(d))[x]), length(keys(d)))
     t_vals = ntuple(x -> collect(values(d))[x], length(values(d)))
     return NamedTuple{t_keys}(t_vals)

diff --git a/test/sections_tests.jl b/test/sections_tests.jl
@@ -1,9 +1,12 @@
 using  MCMCChains, Test
 
 # https://github.com/TuringLang/AdvancedMH.jl/pull/63
+# https://github.com/TuringLang/MCMCChains.jl/issues/443
 @testset "order of parameters" begin
-    chains = Chains(rand(2, 10, 4), vcat(["μ[$i]" for i in 1:9], :lp), (internals=[:lp],))
-    @test names(chains, :parameters) == [Symbol("μ[$i]") for i in 1:9]
+    params = vcat(["μ[$i]" for i in 1:9], :p1, :p2, :p3)
+    chains = Chains(rand(2, length(params)+1, 4), vcat(params, :lp), (internals=[:lp],))
+    @test names(chains, :parameters) == map(Symbol, params)
+    @test collect(keys(get(chains; section=:parameters))) == [:μ, :p1, :p2, :p3]
 end
 
 @testset "describe sections" begin