Remove parsing of Stan-style indices in MCMCChains (#347)

* Remove support for separating indices with dots

* Remove Stan stats key map

* Increment minor version number

* Add test that dots are not split

* Bump ArviZ compat for test and docs environments

* Fix for unicode character variable names

* Remove unused function

* Update ext/ArviZMCMCChainsExt.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix installation and path of CmdStan for docs build

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

.github/workflows/documenter.yml

@@ -8,7 +8,7 @@ on:
     - cron: "0 0 * * *"
 
 env:
-  CMDSTAN_VERSION: "2.33.1"
+  CMDSTAN_VERSION: "2.34.1"
   CMDSTAN_PATH: "${{ GITHUB.WORKSPACE }}/.cmdstan/"
   JULIA_NUM_THREADS: 2
 
@@ -50,6 +50,6 @@ jobs:
         with:
           install-package: false
         env:
-          JULIA_CMDSTAN_HOME: ${{ env.CMDSTAN_PATH }}/cmdstan-${{ env.CMDSTAN_VERSION }}/ # required by CmdStan.jl
+          CMDSTAN: ${{ env.CMDSTAN_PATH }}/cmdstan-${{ env.CMDSTAN_VERSION }}/ # required by StanSample.jl
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}

Project.toml

@@ -1,7 +1,7 @@
 name = "ArviZ"
 uuid = "131c737c-5715-5e2e-ad31-c244f01c1dc7"
 authors = ["Seth Axen <[email protected]>"]
-version = "0.10.5"
+version = "0.11.0"
 
 [deps]
 InferenceObjects = "b5cf5a8d-e756-4ee3-b014-01d49d192c00"

docs/Project.toml

@@ -27,7 +27,7 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
 AlgebraOfGraphics = "0.6.9"
-ArviZ = ">=0.10"
+ArviZ = "0.11.0"
 ArviZExampleData = "0.1.5"
 ArviZPythonPlots = "0.1"
 CairoMakie = "0.8.9, 0.9, 0.10, 0.11"

ext/ArviZMCMCChainsExt.jl

@@ -8,35 +8,26 @@ else
     using ..MCMCChains: MCMCChains
 end
 
-const turing_key_map = Dict(
+const stats_key_map = Dict(
     :hamiltonian_energy => :energy,
     :hamiltonian_energy_error => :energy_error,
     :is_adapt => :tune,
     :max_hamiltonian_energy_error => :max_energy_error,
     :nom_step_size => :step_size_nom,
     :numerical_error => :diverging,
 )
-const stan_key_map = Dict(
-    :accept_stat__ => :acceptance_rate,
-    :divergent__ => :diverging,
-    :energy__ => :energy,
-    :lp__ => :lp,
-    :n_leapfrog__ => :n_steps,
-    :stepsize__ => :step_size,
-    :treedepth__ => :tree_depth,
-)
-const stats_key_map = merge(turing_key_map, stan_key_map)
-
-headtail(x) = x[1], x[2:end]
 
 function split_locname(name::AbstractString)
-    name = replace(name, r"[\[,]" => '.')
-    name = replace(name, ']' => "")
-    name, loc = headtail(split(name, '.'))
-    isempty(loc) && return name, ()
-    loc = tryparse.(Int, loc)
-    Nothing <: eltype(loc) && return name, ()
-    return name, tuple(loc...)
+    endswith(name, "]") || return name, ()
+    basename, index = rsplit(name[1:(end - 1)], "["; limit=2)
+    isempty(index) && return name, ()
+    try
+        loc = parse.(Int, split(index, ','))
+        return basename, tuple(loc...)
+    catch e
+        e isa ArgumentError && return name, ()
+        rethrow(e)
+    end
 end
 function split_locname(name::Symbol)
     subname, loc = split_locname(string(name))

test/ext/ArviZMCMCChainsExt/Project.toml

@@ -7,7 +7,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ArviZ = "0.10.3"
+ArviZ = "0.11.0"
 DimensionalData = "0.24, 0.25, 0.26"
 MCMCChains = "6"
 OrderedCollections = "1"

test/ext/ArviZMCMCChainsExt/runtests.jl

@@ -73,7 +73,7 @@ end
     end
 
     @testset "multivariate" begin
-        var_names = Symbol.(["a[1][1]", "a.2.2", "a[2,1]", "a[1, 2]"])
+        var_names = Symbol.(["a[1, 1]", "a[2,2]", "a[2,1]", "a[1, 2]"])
         coords = (ai=1:2, aj=["aj1", "aj2"])
         dims = (a=[:ai, :aj],)
         nchains, ndraws = 4, 20
@@ -91,6 +91,26 @@ end
         @test arr[:, :, 1, 2] == chns.value[:, ET(var_names[4]), :]
     end
 
+    @testset "dots are not split" begin
+        var_names = Symbol.(["a.b[1]", "a.b[2]", "a.θ[2]", "a.θ[1]"])
+        coords = (dim=["x", "y"],)
+        dims = (var"a.b"=[:dim], var"a.θ"=[:dim])
+        nchains, ndraws = 4, 20
+        chns = makechains(var_names, ndraws, nchains)
+
+        # String or Symbol, which depends on MCMCChains version
+        ET = Base.promote_typeof(chns.name_map.parameters...)
+
+        idata = from_mcmcchains(chns; coords, dims)
+        test_chains_data(chns, idata, :posterior, Symbol.(["a.b", "a.θ"]); coords, dims)
+        arr = idata.posterior.var"a.b"
+        @test arr[:, :, 1] == chns.value[:, ET(var_names[1]), :]
+        @test arr[:, :, 2] == chns.value[:, ET(var_names[2]), :]
+        arr = idata.posterior.var"a.θ"
+        @test arr[:, :, 2] == chns.value[:, ET(var_names[3]), :]
+        @test arr[:, :, 1] == chns.value[:, ET(var_names[4]), :]
+    end
+
     @testset "specify eltypes" begin
         # https://github.com/arviz-devs/ArviZ.jl/issues/125
         nchains, ndraws = 4, 20

test/ext/ArviZSampleChainsDynamicHMCExt/Project.toml

@@ -7,7 +7,7 @@ SampleChainsDynamicHMC = "6d9fd711-e8b2-4778-9c70-c1dfb499d4c4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ArviZ = "0.10.3"
+ArviZ = "0.11.0"
 DimensionalData = "0.24, 0.25, 0.26"
 Random = "1.6"
 SampleChains = "0.5"

test/ext/ArviZSampleChainsExt/Project.toml

@@ -6,7 +6,7 @@ SampleChains = "754583d1-7fc4-4dab-93b5-5eaca5c9622e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ArviZ = "0.10.3"
+ArviZ = "0.11.0"
 DimensionalData = "0.24, 0.25, 0.26"
 Random = "1.6"
 SampleChains = "0.5"