feat: Remove Distributions, DataStructures, and MLJModelInterface/Tables dependencies

Project.toml

@@ -1,29 +1,31 @@
 name = "MCMCDiagnosticTools"
 uuid = "be115224-59cd-429b-ad48-344e309966f0"
-version = "0.3.19"
+version = "0.4.0"
 authors = ["David Widmann", "Seth Axen"]
 
 [workspace]
 projects = ["docs", "test"]
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
-DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
+[weakdeps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+
+[extensions]
+MCMCDiagnosticToolsDistributionsExt = "Distributions"
+MCMCDiagnosticToolsStatsBaseExt = "StatsBase"
+
 [compat]
 AbstractFFTs = "0.5, 1"
-DataAPI = "1.12"
-DataStructures = "0.18.13, 0.19"
 Distributions = "0.25.57"
 LinearAlgebra = "1.6"
 MLJModelInterface = "1.11"

ext/MCMCDiagnosticToolsDistributionsExt.jl

@@ -0,0 +1,16 @@
+module MCMCDiagnosticToolsDistributionsExt
+
+using MCMCDiagnosticTools: MCMCDiagnosticTools
+using Distributions: Distributions
+
+function MCMCDiagnosticTools._rstar_distribution(predictions, ytest, nclasses::Int)
+    # create Poisson binomial distribution with support `0:length(predictions)`
+    distribution = Distributions.PoissonBinomial(map(Distributions.pdf, predictions, ytest))
+
+    # scale distribution to support in `[0, nclasses]`
+    scaled_distribution = (nclasses//length(predictions)) * distribution
+
+    return scaled_distribution
+end
+
+end

ext/MCMCDiagnosticToolsStatsBaseExt.jl

@@ -0,0 +1,12 @@
+module MCMCDiagnosticToolsStatsBaseExt
+
+using MCMCDiagnosticTools: MCMCDiagnosticTools
+using Statistics: Statistics
+using StatsBase: StatsBase
+
+function MCMCDiagnosticTools._expectand_proxy(::typeof(StatsBase.mad), x)
+    x_folded = MCMCDiagnosticTools._fold_around_median(x)
+    return MCMCDiagnosticTools._expectand_proxy(Statistics.median, x_folded)
+end
+
+end

src/MCMCDiagnosticTools.jl

@@ -1,12 +1,8 @@
 module MCMCDiagnosticTools
 
 using AbstractFFTs: AbstractFFTs
-using DataAPI: DataAPI
-using DataStructures: DataStructures
-using Distributions: Distributions
 using MLJModelInterface: MLJModelInterface as MMI
 using SpecialFunctions: SpecialFunctions
-using StatsBase: StatsBase
 using StatsFuns: StatsFuns, sqrt2
 using Tables: Tables
 
@@ -35,4 +31,11 @@ include("heideldiag.jl")
 include("mcse.jl")
 include("rafterydiag.jl")
 include("rstar.jl")
+
+function __init__()
+    if isdefined(Base.Experimental, :register_error_hint)
+        Base.Experimental.register_error_hint(_rstar_error_hint, MethodError)
+    end
+end
+
 end

src/discretediag.jl

@@ -71,7 +71,7 @@ function weiss(X::AbstractMatrix)
     stat = (n / ca) * sum(chi_stat)
     pval = NaN
     if ((m_tot - 1) * (d - 1)) >= 1
-        pval = Distributions.ccdf(Distributions.Chisq((m_tot - 1) * (d - 1)), stat)
+        pval = StatsFuns.chisqccdf((m_tot - 1) * (d - 1), stat)
     end
 
     return (stat, m_tot, pval, ca)
@@ -185,7 +185,7 @@ function bd_inner(Y::AbstractMatrix, m::Int)
 end
 
 @doc raw"""
-    simulate_DAR1!(X::Matrix{Int}, phi::Float64, sampler)
+    simulate_DAR1!(X::Matrix{Int}, phi::Float64, prob::AbstractVector{<:Real})
 
 Simulate a DAR(1) model independently in each column of `X`.
 
@@ -196,30 +196,31 @@ X_t = \alpha_t X_{t-1} + (1 - \alpha_t) \epsilon_{t-1},
 where
 ```math
 \begin{aligned}
-X_1 \sim \text{sampler}, \\
+X_1 \sim \operatorname{Categorical}(\mathrm{prob}), \\
 \alpha_t \sim \operatorname{Bernoulli}(\phi), \\
-\epsilon_{t-1} \sim \text{sampler},
+\epsilon_{t-1} \sim \operatorname{Categorical}(\mathrm{prob}),
 \end{aligned}
 ```
 are independent random variables.
 """
-function simulate_DAR1!(X::Matrix{Int}, phi::Float64, sampler)
+function simulate_DAR1!(X::Matrix{Int}, phi::Float64, prob::AbstractVector{<:Real})
     n = size(X, 1)
     n > 0 || error("output matrix must be non-empty")
+    rng = Random.default_rng()
 
     # for all simulations
     @inbounds for j in axes(X, 2)
         # sample first value from categorical distribution with probabilities `prob`
-        X[1, j] = rand(sampler)
+        X[1, j] = _wsample(rng, prob)
 
         for t in 2:n
             # compute next value
-            X[t, j] = if rand() <= phi
+            X[t, j] = if rand(rng) <= phi
                 # copy previous value with probability `phi`
                 X[t - 1, j]
             else
                 # sample value with probability `1-phi`
-                rand(sampler)
+                _wsample(rng, prob)
             end
         end
     end
@@ -230,9 +231,10 @@ end
 function simulate_MC(N::Int, P::Matrix{Float64})
     X = zeros(Int, N)
     n, m = size(P)
-    X[1] = StatsBase.sample(1:n)
+    rng = Random.default_rng()
+    X[1] = rand(rng, 1:n)
     for i in 2:N
-        X[i] = StatsBase.wsample(1:n, vec(P[X[i - 1], :]))
+        X[i] = _wsample(rng, vec(P[X[i - 1], :]))
     end
     return X
 end
@@ -303,22 +305,21 @@ function diag_all(
                 stat = t * sum(chi_stat)
                 df0 = (m - 1) * (d - 1)
                 if m > 1 && !isnan(stat)
-                    pval = Distributions.ccdf(Distributions.Chisq(df0), stat)
+                    pval = StatsFuns.chisqccdf(df0, stat)
                 end
             elseif method == :weiss
                 stat = (t / ca) * sum(chi_stat)
                 df0 = (m - 1) * (d - 1)
                 pval = NaN
                 if m > 1 && !isnan(stat)
-                    pval = Distributions.ccdf(Distributions.Chisq(df0), stat)
+                    pval = StatsFuns.chisqccdf(df0, stat)
                 end
             elseif method == :DARBOOT
                 stat = t * sum(chi_stat)
-                sampler_phat = Distributions.sampler(Distributions.Categorical(phat))
                 bstats = zeros(nsim)
                 Y = Matrix{Int}(undef, t, d)
                 for b in 1:nsim
-                    simulate_DAR1!(Y, phia, sampler_phat)
+                    simulate_DAR1!(Y, phia, phat)
                     s = hangartner_inner(Y, m)[1]
                     @inbounds bstats[b] = s
                 end
@@ -339,7 +340,7 @@ function diag_all(
                 stat = hot_stat
                 df0 = df
                 if df > 0 && !isnan(hot_stat)
-                    pval = Distributions.ccdf(Distributions.Chisq(df), hot_stat)
+                    pval = StatsFuns.chisqccdf(df, hot_stat)
                 end
             elseif method == :billingsleyBOOT
                 stat = hot_stat

src/ess_rhat.jl

@@ -640,10 +640,6 @@ end
 function _expectand_proxy(::typeof(Statistics.std), x)
     return (x .- Statistics.mean(x; dims=_sample_dims(x))) .^ 2
 end
-function _expectand_proxy(::typeof(StatsBase.mad), x)
-    x_folded = _fold_around_median(x)
-    return _expectand_proxy(Statistics.median, x_folded)
-end
 function _expectand_proxy(f::Base.Fix2{typeof(Statistics.quantile),<:Real}, x)
     y = similar(x)
     # currently quantile does not support a dims keyword argument

src/gelmandiag.jl

@@ -44,7 +44,7 @@ function _gelmandiag(psi::AbstractArray{<:Real,3}; alpha::Real=0.05)
         estimates[i] = sqrt(correction * (rfixed + rrandom))
 
         if !isnan(rrandom)
-            rrandom *= Distributions.quantile(Distributions.FDist(B_df, W_df[i]), q)
+            rrandom *= StatsFuns.fdistinvcdf(B_df, W_df[i], q)
         end
         upperlimits[i] = sqrt(correction * (rfixed + rrandom))
     end

src/rafterydiag.jl

@@ -35,7 +35,7 @@ function rafterydiag(
         kthin = -1
         burnin = total = NaN
     else
-        dichot = Int[(x .<= StatsBase.quantile(x, q))...]
+        dichot = Int[(x .<= Statistics.quantile(x, q))...]
         kthin = 0
         bic = 1.0
         local test, ntest
@@ -44,7 +44,7 @@ function rafterydiag(
             test = dichot[1:kthin:nx]
             ntest = length(test)
             temp = test[1:(ntest - 2)] + 2 * test[2:(ntest - 1)] + 4 * test[3:ntest]
-            trantest = reshape(StatsBase.counts(temp, 0:7), 2, 2, 2)
+            trantest = reshape(_counts(temp, 0:7), 2, 2, 2)
             g2 = 0.0
             for i1 in 1:2, i2 in 1:2, i3 in 1:2
                 tt = trantest[i1, i2, i3]
@@ -58,7 +58,7 @@ function rafterydiag(
             bic = g2 - 2.0 * log(ntest - 2.0)
         end
 
-        tranfinal = StatsBase.counts(test[1:(ntest - 1)] + 2 * test[2:ntest], 0:3)
+        tranfinal = _counts(test[1:(ntest - 1)] + 2 * test[2:ntest], 0:3)
         alpha = tranfinal[3] / (tranfinal[1] + tranfinal[3])
         beta = tranfinal[2] / (tranfinal[2] + tranfinal[4])
         kthin *= step(range)

src/rstar.jl

@@ -147,7 +147,8 @@ within-chain convergence. The training of the classifier can be inspected by adj
 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
 probabilities of classes, the scaled Poisson-binomial distribution of the ``R^*`` statistic
-is returned (algorithm 2).
+is returned (algorithm 2). Probabilistic classifiers require
+[Distributions.jl](https://github.com/JuliaStats/Distributions.jl) to be loaded.
 
 !!! note
     The correctness of the statistic depends on the convergence of the `classifier` used
@@ -164,6 +165,7 @@ julia> samples = fill(4.0, 100, 3, 2);
 One can compute the distribution of the ``R^*`` statistic (algorithm 2) with a
 probabilistic classifier.
 For instance, we can use a gradient-boosted trees model with `nrounds = 100` sequentially stacked trees and learning rate `eta = 0.05`:
+Because MLJBase already loads Distributions, we don't need to do so ourselves.
 
 ```jldoctest rstar
 julia> model = EvoTreeClassifier(; nrounds=100, eta=0.05);
@@ -254,14 +256,22 @@ function _rstar(
     length(predictions) == length(ytest) ||
         error("numbers of predictions and targets must be equal")
 
-    # create Poisson binomial distribution with support `0:length(predictions)`
-    distribution = Distributions.PoissonBinomial(map(Distributions.pdf, predictions, ytest))
-
-    # scale distribution to support in `[0, nclasses]`
     nclasses = length(MMI.classes(ytest))
-    scaled_distribution = (nclasses//length(predictions)) * distribution
+    return _rstar_distribution(predictions, ytest, nclasses)
+end
+
+# Implemented in MCMCDiagnosticToolsDistributionsExt
+function _rstar_distribution end
 
-    return scaled_distribution
+function _rstar_error_hint(io, exc, argtypes, kwargs)
+    if exc.f === _rstar_distribution
+        print(
+            io,
+            "\n`rstar` with a probabilistic classifier requires Distributions to be loaded, " *
+            "e.g. with `using Distributions`.",
+        )
+    end
+    return nothing
 end
 
 # unsupported types of predictions and targets

src/utils.jl

@@ -50,16 +50,16 @@ found.
 function unique_indices(x)
     inds = eachindex(x)
     T = eltype(inds)
-    ind_map = DataStructures.SortedDict{eltype(x),Vector{T}}()
+    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 = collect(keys(ind_map))
-    indices = collect(values(ind_map))
+    unique = sort!(collect(keys(ind_map)))
+    indices = map(Base.Fix1(getindex, ind_map), unique)
     return unique, indices
 end
 
@@ -177,12 +177,67 @@ function _rank_normalize!(values, x)
         fill!(values, missing)
         return values
     end
-    rank = StatsBase.tiedrank(x)
-    _normal_quantiles_from_ranks!(values, rank)
+    _tiedrank!(values, x)
+    _normal_quantiles_from_ranks!(values, values)
     map!(StatsFuns.norminvcdf, values, values)
     return values
 end
 
+# rank `x` in ascending order, assigning tied values the mean of their ranks.
+# equivalent to `StatsBase.tiedrank(x)`
+function _tiedrank!(ranks::AbstractArray, x::AbstractArray)
+    _tiedrank!(vec(ranks), vec(x))
+    return ranks
+end
+function _tiedrank!(ranks::AbstractVector, x::AbstractVector)
+    perm = sortperm(x)
+    n = lastindex(perm)
+    start = firstindex(perm)
+    offset = 1 - start
+    while start ≤ n
+        xstart = x[perm[start]]
+        # find the last index `stop` of the run of values tied with xstart
+        stop = start
+        while stop < n && x[perm[stop + 1]] == xstart
+            stop += 1
+        end
+        rank = (start + stop) / 2 + offset
+        @views ranks[perm[start:stop]] .= rank
+        start = stop + 1
+    end
+    return ranks
+end
+
+# sample an index `i` of `weights` with probability proportional to `weights[i]`.
+# for `Float64` weights, equivalent to `StatsBase.wsample(rng, eachindex(weights), weights)`
+function _wsample(rng::Random.AbstractRNG, weights::AbstractVector{T}) where {T<:Real}
+    t = rand(rng, float(T)) * sum(weights)
+    i = firstindex(weights)
+    ilast = lastindex(weights)
+    cw = weights[i]
+    while true
+        cw ≥ t && return i
+        # may happen with floating point weights due to numerical inaccuracies.
+        # never return an index with zero weight.
+        i == ilast && return something(findlast(!iszero, weights), i)
+        i += 1
+        cw += weights[i]
+    end
+end
+
+# count the number of occurrences of each value of `levels` in `x`.
+# equivalent to `StatsBase.counts(x, levels)`
+function _counts(x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer})
+    counts = zeros(Int, length(levels))
+    offset = 1 - first(levels)
+    for xi in x
+        if xi in levels
+            counts[xi + offset] += 1
+        end
+    end
+    return counts
+end
+
 # transform the ranks to quantiles of a standard normal distribution applying the
 # "α-β correction" recommended in Eq 6.10.3 of
 # Blom. Statistical Estimates and Transformed Beta-Variables. Wiley; New York, 1958

test/Project.toml

@@ -26,7 +26,7 @@ MCMCDiagnosticTools = {path = ".."}
 Aqua = "0.8.14"
 Distributions = "0.25.57"
 DynamicHMC = "3.2"
-EvoTrees = "0.16.2, 0.18, 0.19"
+EvoTrees = "0.16.3, 0.18, 0.19"
 FFTW = "1.3"
 LogDensityProblems = "0.12, 1, 2"
 LogExpFunctions = "0.3.26, 1"