Support MCMCDiagnosticTools 0.3 and MCMCChains 6 (#79)

* Support MCMCDiagnosticTools 0.3 and MCMCChains 6

* Apply suggestions from code review

Co-authored-by: Seth Axen <[email protected]>

* Remove Turing test dependency

* Update Project.toml

---------

Co-authored-by: Seth Axen <[email protected]>

Project.toml

@@ -1,7 +1,7 @@
 name = "ParetoSmooth"
 uuid = "a68b5a21-f429-434e-8bfa-46b447300aac"
 authors = ["Carlos Parada <[email protected]>"]
-version = "0.7.5"
+version = "0.7.6"
 
 [deps]
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
@@ -20,8 +20,8 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 AxisKeys = "0.1.18, 0.2"
 DynamicPPL = "0.21, 0.22"
 LogExpFunctions = "0.3"
-MCMCChains = "5"
-MCMCDiagnosticTools = "0.1.0"
+MCMCChains = "6"
+MCMCDiagnosticTools = "0.3.2"
 NamedDims = "0.2.35, 1"
 PrettyTables = "2.1,2.2"
 Requires = "1.1.3"
@@ -42,10 +42,9 @@ RData = "df47a6cb-8c03-5eed-afd8-b6050d6c41da"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
-Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [targets]
-test = ["CSV", "DataFrames", "Distributions", "DynamicPPL", "MCMCChains", "RData", "StatsFuns", "Test", "TestSetExtensions", "Turing"]
+test = ["CSV", "DataFrames", "Distributions", "DynamicPPL", "MCMCChains", "RData", "StatsFuns", "Test", "TestSetExtensions"]
 
 [weakdeps]
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"

src/ESS.jl

@@ -1,33 +1,39 @@
-using MCMCDiagnosticTools
+import MCMCDiagnosticTools
 
 export relative_eff, psis_ess, sup_ess
 
 """
     relative_eff(
-        sample::AbstractArray{<:Real, 3}; 
-        method=MCMCDiagnosticTools.FFTESSMethod()
+        sample::AbstractArray{<:Real, 3};
+        source::Union{AbstractString, Symbol} = "default",
+        maxlag::Int = typemax(Int),
+        kwargs..., 
     )
 
-Calculate the relative efficiency of an MCMC chain, i.e. the effective sample size divided
+Calculate the relative efficiency of an MCMC chain, i.e., the effective sample size divided
 by the nominal sample size.
 
+If `lowercase(String(source))` is `"default"` or `"mcmc"`, the relative effective sample size is computed with `MCMCDiagnosticTools.ess`, using keyword arguments `kind = :basic`, `maxlag = maxlag`, and the remaining keyword arguments `kwargs...`.
+Otherwise a vector of ones for each chain is returned.
+
 # Arguments 
 
-  - `sample::AbstractArray{<:Real, 3}`: An array of log-likelihood values.
+  - `sample::AbstractArray{<:Real, 3}`: An array of log-likelihood values of the shape `(parameters, draws, chains)`.
 """
 function relative_eff(
     sample::AbstractArray{<:Real,3}; 
-    source::Union{AbstractString, Symbol}="default", maxlag=size(sample, 2), kwargs...
+    source::Union{AbstractString, Symbol}="default",
+    maxlag=typemax(Int),
+    kwargs...,
 )
-    if lowercase(String(source)) ∉ ["mcmc", "default"]
-        return ones(size(sample, 1))
+    if lowercase(String(source)) ∉ ("mcmc", "default")
+        # Avoid type instability by computing the return type of `ess`
+        T = promote_type(eltype(sample), typeof(zero(eltype(sample)) / 1))
+        res = similar(sample, T, (axes(sample, 3),))
+        return fill!(res, 1)
     end
-
-    dims = size(sample)
-    post_sample_size = dims[2] * dims[3]
-    ess_sample = permutedims(sample, [2, 1, 3])
-    ess, = MCMCDiagnosticTools.ess_rhat(ess_sample; maxlag=maxlag, kwargs...)
-    return r_eff = ess / post_sample_size
+    ess_sample = PermutedDimsArray(sample, (2, 3, 1))
+    return MCMCDiagnosticTools.ess(ess_sample; maxlag, kwargs..., kind=:basic, relative=true)
 end