From 5dedc3026b80bc0e09332f504912d8cc1245e40f Mon Sep 17 00:00:00 2001
From: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
Date: Mon, 18 Nov 2024 17:38:42 +0000
Subject: [PATCH] Patch un-imported `Accessors` in `MCMCChainsExt` (#239)

---
 Project.toml                  |  2 +-
 ext/JuliaBUGSMCMCChainsExt.jl |  1 +
 test/ext/mcmchains.jl         | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index f9783e7da..6ef056daf 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "JuliaBUGS"
 uuid = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
diff --git a/ext/JuliaBUGSMCMCChainsExt.jl b/ext/JuliaBUGSMCMCChainsExt.jl
index b17d575ba..eed7d72bd 100644
--- a/ext/JuliaBUGSMCMCChainsExt.jl
+++ b/ext/JuliaBUGSMCMCChainsExt.jl
@@ -6,6 +6,7 @@ using JuliaBUGS.AbstractPPL
 using JuliaBUGS.BUGSPrimitives
 using JuliaBUGS.LogDensityProblems
 using JuliaBUGS.LogDensityProblemsAD
+using JuliaBUGS: Accessors
 using AbstractMCMC
 using MCMCChains: Chains
 
diff --git a/test/ext/mcmchains.jl b/test/ext/mcmchains.jl
index fa544bd4a..0b86efdad 100644
--- a/test/ext/mcmchains.jl
+++ b/test/ext/mcmchains.jl
@@ -91,4 +91,36 @@
     @test means[:beta].nt.mean[1] ≈ 2.1 atol = 0.2
     @test means[:sigma].nt.mean[1] ≈ 0.9 atol = 0.2
     @test means[:gen_quant].nt.mean[1] ≈ 4.2 atol = 0.2
+
+    # test for more complicated varnames
+    model_def = @bugs begin
+        A[1, 1:3] ~ Dirichlet(ones(3))
+        A[2, 1:3] ~ Dirichlet(ones(3))
+        A[3, 1:3] ~ Dirichlet(ones(3))
+
+        mu[1:3] ~ MvNormal(zeros(3), 10 * Diagonal(ones(3)))
+        sigma[1] ~ InverseGamma(2, 3)
+        sigma[2] ~ InverseGamma(2, 3)
+        sigma[3] ~ InverseGamma(2, 3)
+    end
+    model = compile(model_def, (;))
+    ad_model = ADgradient(:ReverseDiff, model; compile=Val(true))
+    hmc_chain = AbstractMCMC.sample(ad_model, NUTS(0.8), 10; chain_type=Chains)
+    @test hmc_chain.name_map[:parameters] == [
+        Symbol("sigma[3]"),
+        Symbol("sigma[2]"),
+        Symbol("sigma[1]"),
+        Symbol("mu[1:3][1]"),
+        Symbol("mu[1:3][2]"),
+        Symbol("mu[1:3][3]"),
+        Symbol("A[3, 1:3][1]"),
+        Symbol("A[3, 1:3][2]"),
+        Symbol("A[3, 1:3][3]"),
+        Symbol("A[2, 1:3][1]"),
+        Symbol("A[2, 1:3][2]"),
+        Symbol("A[2, 1:3][3]"),
+        Symbol("A[1, 1:3][1]"),
+        Symbol("A[1, 1:3][2]"),
+        Symbol("A[1, 1:3][3]"),
+    ]
 end