Change dimension ordering

Project.toml

@@ -1,7 +1,7 @@
 name = "MCMCDiagnosticTools"
 uuid = "be115224-59cd-429b-ad48-344e309966f0"
 authors = ["David Widmann"]
-version = "0.1.5"
+version = "0.2.0"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/discretediag.jl

@@ -423,9 +423,11 @@ function discretediag_sub(
 end
 
 """
-    discretediag(chains::AbstractArray{<:Real,3}; frac=0.3, method=:weiss, nsim=1_000)
+    discretediag(samples::AbstractArray{<:Real,3}; frac=0.3, method=:weiss, nsim=1_000)
 
-Compute discrete diagnostic where `method` can be one of `:weiss`, `:hangartner`,
+Compute discrete diagnostic on `samples` with shape `(parameters, draws, chains)`.
+
+`method` can be one of `:weiss`, `:hangartner`,
 `:DARBOOT`, `:MCBOOT`, `:billinsgley`, and `:billingsleyBOOT`.
 
 # References
@@ -441,9 +443,9 @@ function discretediag(
     )
     0 < frac < 1 || throw(ArgumentError("`frac` must be in (0,1)"))
 
-    num_iters = size(chains, 1)
+    num_iters = size(chains, 2)
     between_chain_vals, within_chain_vals, _, _ = discretediag_sub(
-        chains, frac, method, nsim, num_iters, num_iters
+        permutedims(chains, (2, 1, 3)), frac, method, nsim, num_iters, num_iters
     )
 
     return between_chain_vals, within_chain_vals

src/ess.jl

@@ -201,7 +201,7 @@ end
     )
 
 Estimate the effective sample size and the potential scale reduction of the `samples` of
-shape (draws, parameters, chains) with the `method` and a maximum lag of `maxlag`.
+shape `(parameters, draws, chains)` with the `method` and a maximum lag of `maxlag`.
 
 See also: [`ESSMethod`](@ref), [`FFTESSMethod`](@ref), [`BDAESSMethod`](@ref)
 """
@@ -211,8 +211,8 @@ function ess_rhat(
     maxlag::Int=250,
 )
     # compute size of matrices (each chain is split!)
-    niter = size(chains, 1) ÷ 2
-    nparams = size(chains, 2)
+    niter = size(chains, 2) ÷ 2
+    nparams = size(chains, 1)
     nchains = 2 * size(chains, 3)
     ntotal = niter * nchains
 
@@ -238,7 +238,7 @@ function ess_rhat(
     rhat = Vector{T}(undef, nparams)
 
     # for each parameter
-    for (i, chains_slice) in enumerate((view(chains, :, i, :) for i in axes(chains, 2)))
+    for (i, chains_slice) in enumerate((selectdim(chains, 1, i) for i in axes(chains, 1)))
         # check that no values are missing
         if any(x -> x === missing, chains_slice)
             rhat[i] = missing

src/gelmandiag.jl

@@ -1,14 +1,14 @@
 function _gelmandiag(psi::AbstractArray{<:Real,3}; alpha::Real=0.05)
-    niters, nparams, nchains = size(psi)
+    nparams, niters, nchains = size(psi)
     nchains > 1 || error("Gelman diagnostic requires at least 2 chains")
 
     rfixed = (niters - 1) / niters
     rrandomscale = (nchains + 1) / (nchains * niters)
 
-    S2 = map(Statistics.cov, (view(psi, :, :, i) for i in axes(psi, 3)))
+    S2 = map(x -> Statistics.cov(x; dims=2), (view(psi, :, :, i) for i in axes(psi, 3)))
     W = Statistics.mean(S2)
 
-    psibar = dropdims(Statistics.mean(psi; dims=1); dims=1)'
+    psibar = dropdims(Statistics.mean(psi; dims=2); dims=2)'
     B = niters .* Statistics.cov(psibar)
 
     w = LinearAlgebra.diag(W)
@@ -52,9 +52,10 @@ function _gelmandiag(psi::AbstractArray{<:Real,3}; alpha::Real=0.05)
 end
 
 """
-    gelmandiag(chains::AbstractArray{<:Real,3}; alpha::Real=0.95)
+    gelmandiag(samples::AbstractArray{<:Real,3}; alpha::Real=0.95)
 
-Compute the Gelman, Rubin and Brooks diagnostics [^Gelman1992] [^Brooks1998].  Values of the
+Compute the Gelman, Rubin and Brooks diagnostics [^Gelman1992] [^Brooks1998] on `samples`
+with shape `(parameters, draws, chains)`.  Values of the
 diagnostic’s potential scale reduction factor (PSRF) that are close to one suggest
 convergence.  As a rule-of-thumb, convergence is rejected if the 97.5 percentile of a PSRF
 is greater than 1.2.
@@ -70,12 +71,13 @@ function gelmandiag(chains::AbstractArray{<:Real,3}; kwargs...)
 end
 
 """
-    gelmandiag_multivariate(chains::AbstractArray{<:Real,3}; alpha::Real=0.05)
+    gelmandiag_multivariate(samples::AbstractArray{<:Real,3}; alpha::Real=0.05)
 
-Compute the multivariate Gelman, Rubin and Brooks diagnostics.
+Compute the multivariate Gelman, Rubin and Brooks diagnostics on `samples` with shape
+`(parameters, draws, chains)`.
 """
 function gelmandiag_multivariate(chains::AbstractArray{<:Real,3}; kwargs...)
-    niters, nparams, nchains = size(chains)
+    nparams, niters, nchains = size(chains)
     if nparams < 2
         error(
             "computation of the multivariate potential scale reduction factor requires ",

src/mcse.jl

@@ -1,7 +1,9 @@
 """
-    mcse(x::AbstractVector{<:Real}; method::Symbol=:imse, kwargs...)
+    mcse(samples::AbstractVector{<:Real}; method::Symbol=:imse, kwargs...)
+    mcse(samples::AbstractArray{<:Real,3}; method::Symbol=:imse, kwargs...)
 
-Compute the Monte Carlo standard error (MCSE) of samples `x`.
+Compute the Monte Carlo standard error (MCSE) of `samples` of shape `(draws,)` or
+`(parameters, draws, chains)`
 The optional argument `method` describes how the errors are estimated. Possible options are:
 
 - `:bm` for batch means [^Glynn1991]
@@ -23,6 +25,9 @@ function mcse(x::AbstractVector{<:Real}; method::Symbol=:imse, kwargs...)
         throw(ArgumentError("unsupported MCSE method $method"))
     end
 end
+function mcse(x::AbstractArray{<:Real,3}; kwargs...)
+    return dropdims(mapslices(xi -> mcse(vec(xi); kwargs...), x; dims=(2, 3)); dims=(2, 3))
+end
 
 function mcse_bm(x::AbstractVector{<:Real}; size::Int=floor(Int, sqrt(length(x))))
     n = length(x)

src/rstar.jl

@@ -2,14 +2,16 @@
     rstar(
         rng=Random.GLOBAL_RNG,
         classifier,
-        samples::AbstractMatrix,
-        chain_indices::AbstractVector{Int};
+        samples::AbstractArray,
+        [chain_indices::AbstractVector{Int}];
         subset::Real=0.8,
         verbosity::Int=0,
     )
 
-Compute the ``R^*`` convergence statistic of the `samples` with shape (draws, parameters)
-and corresponding chains `chain_indices` with the `classifier`.
+Compute the ``R^*`` convergence statistic of the `samples` with the `classifier`.
+
+Either `samples` has shape `(parameters, draws, chains)`, or `samples` has shape
+`(parameters, draws)` and `chain_indices` must be provided.
 
 This implementation is an adaption of algorithms 1 and 2 described by Lambert and Vehtari.
 
@@ -32,16 +34,14 @@ is returned (algorithm 2).
 ```jldoctest rstar; setup = :(using Random; Random.seed!(100))
 julia> using MLJBase, MLJXGBoostInterface, Statistics
 
-julia> samples = fill(4.0, 300, 2);
-
-julia> chain_indices = repeat(1:3; outer=100);
+julia> samples = fill(4.0, 2, 100, 3);
 ```
 
 One can compute the distribution of the ``R^*`` statistic (algorithm 2) with the
 probabilistic classifier.
 
 ```jldoctest rstar
-julia> distribution = rstar(XGBoostClassifier(), samples, chain_indices);
+julia> distribution = rstar(XGBoostClassifier(), samples);
 
 julia> isapprox(mean(distribution), 1; atol=0.1)
 true
@@ -54,7 +54,7 @@ predicting the mode. In MLJ this corresponds to a pipeline of models.
 ```jldoctest rstar
 julia> xgboost_deterministic = Pipeline(XGBoostClassifier(); operation=predict_mode);
 
-julia> value = rstar(xgboost_deterministic, samples, chain_indices);
+julia> value = rstar(xgboost_deterministic, samples);
 
 julia> isapprox(value, 1; atol=0.2)
 true
@@ -73,7 +73,7 @@ function rstar(
     verbosity::Int=0,
 )
     # checks
-    size(x, 1) != length(y) && throw(DimensionMismatch())
+    size(x, 2) != length(y) && throw(DimensionMismatch())
     0 < subset < 1 || throw(ArgumentError("`subset` must be a number in (0, 1)"))
 
     # randomly sub-select training and testing set
@@ -88,11 +88,11 @@ function rstar(
     # train classifier on training data
     ycategorical = MLJModelInterface.categorical(y)
     fitresult, _ = MLJModelInterface.fit(
-        classifier, verbosity, Tables.table(x[train_ids, :]), ycategorical[train_ids]
+        classifier, verbosity, Tables.table(x[:, train_ids]'), ycategorical[train_ids]
     )
 
     # compute predictions on test data
-    xtest = Tables.table(x[test_ids, :])
+    xtest = Tables.table(x[:, test_ids]')
     predictions = _predict(classifier, fitresult, xtest)
 
     # compute statistic
@@ -114,6 +114,17 @@ function _predict(model::MLJModelInterface.Model, fitresult, x)
     end
 end
 
+function rstar(
+    rng::Random.AbstractRNG,
+    classifier::MLJModelInterface.Supervised,
+    x::AbstractArray{<:Any,3};
+    kwargs...
+)
+    samples = reshape(x, size(x, 1), :)
+    chain_inds = repeat(axes(x, 3); inner=size(x, 2))
+    return rstar(rng, classifier, samples, chain_inds; kwargs...)
+end
+
 function rstar(
     classif::MLJModelInterface.Supervised,
     x::AbstractMatrix,
@@ -123,6 +134,14 @@ function rstar(
     return rstar(Random.GLOBAL_RNG, classif, x, y; kwargs...)
 end
 
+function rstar(
+    classif::MLJModelInterface.Supervised,
+    x::AbstractArray{<:Any,3};
+    kwargs...,
+)
+    return rstar(Random.GLOBAL_RNG, classif, x; kwargs...)
+end
+
 # R⋆ for deterministic predictions (algorithm 1)
 function _rstar(predictions::AbstractVector{T}, ytest::AbstractVector{T}) where {T}
     length(predictions) == length(ytest) ||

test/discretediag.jl

@@ -1,7 +1,7 @@
 @testset "discretediag.jl" begin
     nparams = 4
     nchains = 2
-    samples = rand(-100:100, 100, nparams, nchains)
+    samples = rand(-100:100, nparams, 100, nchains)
 
     @testset "results" begin
         for method in

test/ess.jl

@@ -31,7 +31,7 @@
     end
 
     @testset "ESS and R̂ (IID samples)" begin
-        rawx = randn(10_000, 40, 10)
+        rawx = randn(40, 10_000, 10)
 
         # Repeat tests with different scales
         for scale in (1, 50, 100)
@@ -58,7 +58,7 @@
     end
 
     @testset "ESS and R̂ (identical samples)" begin
-        x = ones(10_000, 40, 10)
+        x = ones(40, 10_000, 10)
 
         ess_standard, rhat_standard = ess_rhat(x)
         ess_standard2, rhat_standard2 = ess_rhat(x; method=ESSMethod())
@@ -75,15 +75,15 @@
     end
 
     @testset "ESS and R̂ (single sample)" begin # check that issue #137 is fixed
-        x = rand(1, 5, 3)
+        x = rand(5, 1, 3)
 
         for method in (ESSMethod(), FFTESSMethod(), BDAESSMethod())
             # analyze array
             ess_array, rhat_array = ess_rhat(x; method=method)
 
-            @test length(ess_array) == size(x, 2)
+            @test length(ess_array) == size(x, 1)
             @test all(ismissing, ess_array) # since min(maxlag, niter - 1) = 0
-            @test length(rhat_array) == size(x, 2)
+            @test length(rhat_array) == size(x, 1)
             @test all(ismissing, rhat_array)
         end
     end

test/gelmandiag.jl

@@ -1,7 +1,7 @@
 @testset "gelmandiag.jl" begin
     nparams = 4
     nchains = 2
-    samples = randn(100, nparams, nchains)
+    samples = randn(nparams, 100, nchains)
 
     @testset "results" begin
         result = @inferred(gelmandiag(samples))
@@ -24,6 +24,6 @@
 
     @testset "exceptions" begin
         @test_throws ErrorException gelmandiag(samples[:, :, 1:1])
-        @test_throws ErrorException gelmandiag_multivariate(samples[:, 1:1, :])
+        @test_throws ErrorException gelmandiag_multivariate(samples[1:1, :, :])
     end
 end

test/mcse.jl

@@ -1,7 +1,6 @@
 @testset "mcse.jl" begin
-    samples = randn(100)
-
-    @testset "results" begin
+    @testset "results 1d" begin
+        samples = randn(100)
         result = @inferred(mcse(samples))
         @test result isa Float64
         @test result > 0
@@ -13,13 +12,32 @@
         end
     end
 
+    @testset  "results 3d" begin
+        nparams = 2
+        nchains = 4
+        samples = randn(nparams, 100, nchains)
+        result = mcse(samples)  # mapslices is not type-inferrable
+        @test result isa Vector{Float64}
+        @test length(result) == nparams
+        @test all(r -> r > 0, result)
+
+        for method in (:imse, :ipse, :bm)
+            result = mcse(samples)  # mapslices is not type-inferrable
+            @test result isa Vector{Float64}
+            @test length(result) == nparams
+            @test all(r -> r > 0, result)
+        end
+    end
+
     @testset "warning" begin
+        samples = randn(100)
         for size in (51, 75, 100, 153)
             @test_logs (:warn,) mcse(samples; method=:bm, size=size)
         end
     end
 
     @testset "exception" begin
+        samples = randn(100)
         @test_throws ArgumentError mcse(samples; method=:somemethod)
     end
 end

test/rstar/Project.toml

@@ -4,6 +4,7 @@ MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MLJLIBSVMInterface = "61c7150f-6c77-4bb1-949c-13197eac2a52"
 MLJXGBoostInterface = "54119dfa-1dab-4055-a167-80440f4f7a91"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]