Fix Chains output for scalar parameters

Author: devmotion

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{<:AbstractArray}},
+    ts::Vector{<:AbstractTransition},
     model::DensityModel,
     sampler::MHSampler,
     state,
@@ -63,7 +63,7 @@ function AbstractMCMC.bundle_samples(
 end
 
 function AbstractMCMC.bundle_samples(
-    ts::Vector{<:Vector{<:Transition}},
+    ts::Vector{<:Vector{<:AbstractTransition}},
     model::DensityModel,
     sampler::Ensemble,
     state,

src/structarray-connect.jl

@@ -2,7 +2,7 @@ import .StructArrays: StructArray
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    ts,
+    ts::Vector{<:AbstractTransition},
     model::DensityModel,
     sampler::MHSampler,
     state,

test/runtests.jl

@@ -74,17 +74,34 @@ include("util.jl")
     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)
-
+        # Array of parameters
+        chain1 = sample(
+            model, StaticMH([Normal(0,1), Normal(0, 1)]), 10_000;
+            param_names=["μ", "σ"], chain_type=Chains
+        )
+        @test chain1 isa Chains
         @test mean(chain1["μ"]) ≈ 0.0 atol=0.1
         @test mean(chain1["σ"]) ≈ 1.0 atol=0.1
 
+        # NamedTuple of parameters
+        chain2 = sample(
+            model,
+            MetropolisHastings(
+                (μ = StaticProposal(Normal(0,1)), σ = StaticProposal(Normal(0, 1)))
+            ), 10_000;
+            chain_type=Chains
+        )
+        @test chain2 isa Chains
         @test mean(chain2["μ"]) ≈ 0.0 atol=0.1
         @test mean(chain2["σ"]) ≈ 1.0 atol=0.1
+
+        # Scalar parameter
+        chain3 = sample(
+            DensityModel(x -> loglikelihood(Normal(x, 1), data)),
+            StaticMH(Normal(0, 1)), 10_000; param_names=["μ"], chain_type=Chains
+        )
+        @test chain3 isa Chains
+        @test mean(chain3["μ"]) ≈ 0.0 atol=0.1
     end
 
     @testset "Proposal styles" begin
@@ -194,7 +211,6 @@ include("util.jl")
     end
 
     @testset "MALA" begin
-        
         # Set up the sampler.
         sigma = 1e-1
         spl1 = MALA(x -> MvNormal((sigma^2 / 2) .* x, sigma))
@@ -203,9 +219,8 @@ include("util.jl")
         chain1 = sample(model, spl1, 100000; init_params=ones(2), chain_type=StructArray, param_names=["μ", "σ"])
 
         @test mean(chain1.μ) ≈ 0.0 atol=0.1
-        @test mean(chain1.σ) ≈ 1.0 atol=0.1 
+        @test mean(chain1.σ) ≈ 1.0 atol=0.1
     end
 
     @testset "EMCEE" begin include("emcee.jl") end
-  
 end