Implement fixes for rstar

Project.toml

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

src/MCMCDiagnosticTools.jl

@@ -22,6 +22,7 @@ export mcse
 export rafterydiag
 export rstar
 
+include("utils.jl")
 include("bfmi.jl")
 include("discretediag.jl")
 include("ess.jl")

src/rstar.jl

@@ -4,7 +4,8 @@
         classifier::MLJModelInterface.Supervised,
         samples,
         chain_indices::AbstractVector{Int};
-        subset::Real=0.8,
+        subset::Real=0.7,
+        split_chains::Int=2,
         verbosity::Int=0,
     )
 
@@ -23,26 +24,25 @@ function rstar(
     classifier::MLJModelInterface.Supervised,
     x,
     y::AbstractVector{Int};
-    subset::Real=0.8,
+    subset::Real=0.7,
+    split_chains::Int=2,
     verbosity::Int=0,
 )
     # checks
     MLJModelInterface.nrows(x) != length(y) && throw(DimensionMismatch())
     0 < subset < 1 || throw(ArgumentError("`subset` must be a number in (0, 1)"))
 
+    ysplit = split_chain_indices(y, split_chains)
+
     # randomly sub-select training and testing set
-    N = length(y)
-    Ntrain = round(Int, N * subset)
-    0 < Ntrain < N ||
+    train_ids, test_ids = shuffle_split_stratified(rng, ysplit, subset)
+    0 < length(train_ids) < length(y) ||
         throw(ArgumentError("training and test data subsets must not be empty"))
-    ids = Random.randperm(rng, N)
-    train_ids = view(ids, 1:Ntrain)
-    test_ids = view(ids, (Ntrain + 1):N)
 
     xtable = _astable(x)
 
     # train classifier on training data
-    ycategorical = MLJModelInterface.categorical(y)
+    ycategorical = MLJModelInterface.categorical(ysplit)
     xtrain = MLJModelInterface.selectrows(xtable, train_ids)
     fitresult, _ = MLJModelInterface.fit(
         classifier, verbosity, xtrain, ycategorical[train_ids]
@@ -79,7 +79,8 @@ end
         rng::Random.AbstractRNG=Random.default_rng(),
         classifier::MLJModelInterface.Supervised,
         samples::AbstractArray{<:Real,3};
-        subset::Real=0.8,
+        subset::Real=0.7,
+        split_chains::Int=2,
         verbosity::Int=0,
     )
 
@@ -91,8 +92,10 @@ This implementation is an adaption of algorithms 1 and 2 described by Lambert an
 
 The `classifier` has to be a supervised classifier of the MLJ framework (see the
 [MLJ documentation](https://alan-turing-institute.github.io/MLJ.jl/dev/list_of_supported_models/#model_list)
-for a list of supported models). It is trained with a `subset` of the samples. The training
-of the classifier can be inspected by adjusting the `verbosity` level.
+for a list of supported models). It is trained with a `subset` of the samples from each
+chain. Each chain is split into `split_chains` separate chains to additionally check for
+within-chain convergence. The training of the classifier can be inspected by adjusting the
+`verbosity` level.
 
 If the classifier is deterministic, i.e., if it predicts a class, the value of the ``R^*``
 statistic is returned (algorithm 1). If the classifier is probabilistic, i.e., if it outputs

src/utils.jl

@@ -0,0 +1,96 @@
+"""
+    unique_indices(x) -> (unique, indices)
+
+Return the results of `unique(collect(x))` along with the a vector of the same length whose
+elements are the indices in `x` at which the corresponding unique element in `unique` is
+found.
+"""
+function unique_indices(x)
+    inds = eachindex(x)
+    T = eltype(inds)
+    ind_map = Dict{eltype(x),Vector{T}}()
+    for i in inds
+        xi = x[i]
+        inds_xi = get!(ind_map, xi) do
+            return T[]
+        end
+        push!(inds_xi, i)
+    end
+    unique = sort!(collect(keys(ind_map)))
+    indices = [ind_map[xi] for xi in unique]
+    return unique, indices
+end
+
+"""
+    split_chain_indices(
+        chain_inds::AbstractVector{Int},
+        split::Int=2,
+    ) -> AbstractVector{Int}
+
+Split each chain in `chain_inds` into `split` chains.
+
+For each chain in `chain_inds`, all entries are assumed to correspond to draws that have
+been ordered by iteration number. The result is a vector of the same length as `chain_inds`
+where each entry is the new index of the chain that the corresponding draw belongs to.
+"""
+function split_chain_indices(c::AbstractVector{<:Int}, split::Int=2)
+    cnew = similar(c)
+    if split == 1
+        copyto!(cnew, c)
+        return cnew
+    end
+    chains, indices = unique_indices(c)
+    chain_ind = 0
+    for (chain, inds) in zip(chains, indices)
+        ndraws_per_split, rem = divrem(length(inds), split)
+        ilast = 0
+        # here we can't use Iterators.partition because it's greedy. e.g. we can't partition
+        # 4 items across 3 partitions because Iterators.partition(1:4, 1) == [[1], [2], [3]]
+        # and Iterators.partition(1:4, 2) == [[1, 2], [3, 4]]. But we would want
+        # [[1, 2], [3], [4]].
+        for j in 1:split
+            chain_ind += 1
+            ndraws_this_split = ndraws_per_split + (j ≤ rem)
+            i = ilast + 1
+            ilast = i + ndraws_this_split - 1
+            @views cnew[inds[i:ilast]] .= chain_ind
+        end
+    end
+    return cnew
+end
+
+"""
+    shuffle_split_stratified(
+        rng::Random.AbstractRNG,
+        group_ids::AbstractVector,
+        frac::Real,
+    ) -> (inds1, inds2)
+
+Randomly split the indices of `group_ids` into two groups, where `frac` indices from each
+group are in `inds1` and the remainder are in `inds2`.
+
+This is used, for example, to split data into training and test data while preserving the
+class balances.
+"""
+function shuffle_split_stratified(
+    rng::Random.AbstractRNG, groups_ids::AbstractVector, frac::Real
+)
+    T = eltype(eachindex(groups_ids))
+    groups, indices = unique_indices(groups_ids)
+    N1_tot = sum(x -> round(Int, length(x) * frac), indices)
+    N2_tot = length(groups_ids) - N1_tot
+    inds1 = Vector{T}(undef, N1_tot)
+    inds2 = Vector{T}(undef, N2_tot)
+    items_in_1 = items_in_2 = 0
+    for (group, inds) in zip(groups, indices)
+        N = length(inds)
+        N1 = round(Int, N * frac)
+        N2 = N - N1
+        ids = Random.randperm(rng, N)
+        @views inds1[(items_in_1 + 1):(items_in_1 + N1)] .= inds[ids[1:N1]]
+        @views inds2[(items_in_2 + 1):(items_in_2 + N2)] .= inds[ids[(N1 + 1):N]]
+        items_in_1 += N1
+        items_in_2 += N2
+    end
+    return inds1, inds2
+end

test/rstar.jl

@@ -30,7 +30,7 @@ const xgboost_deterministic = Pipeline(XGBoostClassifier(); operation=predict_mo
                 @test dist isa LocationScale
                 @test dist.ρ isa PoissonBinomial
                 @test minimum(dist) == 0
-                @test maximum(dist) == 3
+                @test maximum(dist) == 6
             end
             @test mean(dist) ≈ 1 rtol = 0.2
             wrapper === Vector && break
@@ -48,7 +48,7 @@ const xgboost_deterministic = Pipeline(XGBoostClassifier(); operation=predict_mo
                 @test dist isa LocationScale
                 @test dist.ρ isa PoissonBinomial
                 @test minimum(dist) == 0
-                @test maximum(dist) == 4
+                @test maximum(dist) == 8
             end
             @test mean(dist) ≈ 1 rtol = 0.15
 
@@ -58,7 +58,7 @@ const xgboost_deterministic = Pipeline(XGBoostClassifier(); operation=predict_mo
                 100 .* cos.(1:N) 100 .* sin.(1:N)
             ])
             chain_indices = repeat(1:2; inner=N)
-            dist = rstar(classifier, samples, chain_indices)
+            dist = rstar(classifier, samples, chain_indices; split_chains=1)
 
             # Mean of the statistic should be close to 2, i.e., the classifier should be able to
             # learn an almost perfect decision boundary between chains.
@@ -71,6 +71,17 @@ const xgboost_deterministic = Pipeline(XGBoostClassifier(); operation=predict_mo
                 @test maximum(dist) == 2
             end
             @test mean(dist) ≈ 2 rtol = 0.15
+
+            # Compute the R⋆ statistic for identical chains that individually have not mixed.
+            samples = ones(sz)
+            samples[div(N, 2):end, :] .= 2
+            chain_indices = repeat(1:4; outer=div(N, 4))
+            dist = rstar(classifier, samples, chain_indices; split_chains=1)
+            # without split chains cannot distinguish between chains
+            @test mean(dist) ≈ 1 rtol = 0.15
+            dist = rstar(classifier, samples, chain_indices)
+            # with split chains can learn almost perfect decision boundary
+            @test mean(dist) ≈ 2 rtol = 0.15
         end
         wrapper === Vector && continue
 

test/runtests.jl

@@ -10,6 +10,10 @@ using Test
 Random.seed!(1)
 
 @testset "MCMCDiagnosticTools.jl" begin
+    @testset "utils" begin
+        include("utils.jl")
+    end
+
     @testset "Bayesian fraction of missing information" begin
         include("bfmi.jl")
     end

test/utils.jl

@@ -0,0 +1,60 @@
+using MCMCDiagnosticTools
+using Test
+using Random
+
+@testset "unique_indices" begin
+    @testset "indices=$(eachindex(inds))" for inds in [
+        rand(11:14, 100), transpose(rand(11:14, 10, 10))
+    ]
+        unique, indices = @inferred MCMCDiagnosticTools.unique_indices(inds)
+        @test unique isa Vector{Int}
+        if eachindex(inds) isa CartesianIndices{2}
+            @test indices isa Vector{Vector{CartesianIndex{2}}}
+        else
+            @test indices isa Vector{Vector{Int}}
+        end
+        @test issorted(unique)
+        @test issetequal(union(indices...), eachindex(inds))
+        for i in eachindex(unique, indices)
+            @test all(inds[indices[i]] .== unique[i])
+        end
+    end
+end
+
+@testset "split_chain_indices" begin
+    c = [2, 2, 1, 3, 4, 3, 4, 1, 2, 1, 4, 3, 3, 2, 4, 3, 4, 1, 4, 1]
+    @test @inferred(MCMCDiagnosticTools.split_chain_indices(c, 1)) == c
+
+    cnew = @inferred MCMCDiagnosticTools.split_chain_indices(c, 2)
+    unique, indices = MCMCDiagnosticTools.unique_indices(c)
+    uniquenew, indicesnew = MCMCDiagnosticTools.unique_indices(cnew)
+    for (i, inew) in enumerate(1:2:7)
+        @test length(indicesnew[inew]) ≥ length(indicesnew[inew + 1])
+        @test indices[i] == vcat(indicesnew[inew], indicesnew[inew + 1])
+    end
+
+    cnew = MCMCDiagnosticTools.split_chain_indices(c, 3)
+    unique, indices = MCMCDiagnosticTools.unique_indices(c)
+    uniquenew, indicesnew = MCMCDiagnosticTools.unique_indices(cnew)
+    for (i, inew) in enumerate(1:3:11)
+        @test length(indicesnew[inew]) ≥
+            length(indicesnew[inew + 1]) ≥
+            length(indicesnew[inew + 2])
+        @test indices[i] ==
+            vcat(indicesnew[inew], indicesnew[inew + 1], indicesnew[inew + 2])
+    end
+end
+
+@testset "shuffle_split_stratified" begin
+    rng = Random.default_rng()
+    c = rand(1:4, 100)
+    unique, indices = MCMCDiagnosticTools.unique_indices(c)
+    @testset "frac=$frac" for frac in [0.3, 0.5, 0.7]
+        inds1, inds2 = @inferred(MCMCDiagnosticTools.shuffle_split_stratified(rng, c, frac))
+        @test issetequal(vcat(inds1, inds2), eachindex(c))
+        for inds in indices
+            common_inds = intersect(inds1, inds)
+            @test length(common_inds) == round(frac * length(inds))
+        end
+    end
+end