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

Project.toml

@@ -1,31 +1,37 @@
 name = "MCMCDiagnosticTools"
 uuid = "be115224-59cd-429b-ad48-344e309966f0"
-version = "0.3.19"
+version = "0.3.20"
 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"
+AliasTables = "66dad0bd-aa9a-41b7-9441-69ab47430ed8"
 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"
+
+[weakdeps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
+MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
+[extensions]
+MCMCDiagnosticToolsMLJBaseExt = ["Distributions", "MLJBase", "MLJModelInterface", "Tables"]
+MCMCDiagnosticToolsStatsBaseExt = "StatsBase"
+
 [compat]
 AbstractFFTs = "0.5, 1"
-DataAPI = "1.12"
-DataStructures = "0.18.13, 0.19"
+AliasTables = "1"
 Distributions = "0.25.57"
 LinearAlgebra = "1.6"
+MLJBase = "0.20.18, 0.21, 1"
 MLJModelInterface = "1.11"
 Random = "1.6"
 SpecialFunctions = "1.8.4, 2"

ext/MCMCDiagnosticToolsMLJBaseExt.jl

@@ -0,0 +1,160 @@
+module MCMCDiagnosticToolsMLJBaseExt
+
+using Distributions: Distributions
+using MCMCDiagnosticTools:
+    MCMCDiagnosticTools, shuffle_split_stratified, split_chain_indices
+using MLJModelInterface: MLJModelInterface as MMI
+using Random: Random
+using Statistics: Statistics
+using Tables: Tables
+
+import MCMCDiagnosticTools:
+    _check_model_supports_continuous_inputs,
+    _check_model_supports_multiclass_predictions,
+    _check_model_supports_multiclass_targets,
+    _rstar
+
+function _rstar(
+    rng::Random.AbstractRNG,
+    classifier,
+    x,
+    y::AbstractVector{Int};
+    subset::Real,
+    split_chains::Int,
+    verbosity::Int,
+)
+    # check the arguments
+    _check_model_supports_continuous_inputs(classifier)
+    _check_model_supports_multiclass_targets(classifier)
+    _check_model_supports_multiclass_predictions(classifier)
+    MMI.nrows(x) != length(y) && throw(DimensionMismatch())
+    0 < subset < 1 || throw(ArgumentError("`subset` must be a number in (0, 1)"))
+
+    # randomly sub-select training and testing set
+    ysplit = split_chain_indices(y, split_chains)
+    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"))
+
+    xtable = _astable(x)
+    ycategorical = MMI.categorical(ysplit)
+
+    # train classifier on training data
+    data = MMI.reformat(classifier, xtable, ycategorical)
+    train_data = MMI.selectrows(classifier, train_ids, data...)
+    fitresult, _ = MMI.fit(classifier, verbosity, train_data...)
+
+    # compute predictions on test data
+    # we exploit that MLJ demands that
+    # reformat(model, args...)[1] = reformat(model, args[1])
+    # (https://alan-turing-institute.github.io/MLJ.jl/dev/adding_models_for_general_use/#Implementing-a-data-front-end)
+    test_data = MMI.selectrows(classifier, test_ids, data[1])
+    predictions = _predict(classifier, fitresult, test_data...)
+
+    # compute statistic
+    ytest = ycategorical[test_ids]
+    result = _rstar(MMI.scitype(predictions), predictions, ytest)
+
+    return result
+end
+
+# check that the model supports the inputs and targets, and has predictions of the desired form
+function _check_model_supports_continuous_inputs(classifier)
+    # ideally we would not allow MMI.Unknown but some models do not implement the traits
+    input_scitype_classifier = MMI.input_scitype(classifier)
+    if input_scitype_classifier !== MMI.Unknown &&
+        !(MMI.Table(MMI.Continuous) <: input_scitype_classifier)
+        throw(
+            ArgumentError(
+                "classifier does not support tables of continuous values as inputs"
+            ),
+        )
+    end
+    return nothing
+end
+function _check_model_supports_multiclass_targets(classifier)
+    target_scitype_classifier = MMI.target_scitype(classifier)
+    if target_scitype_classifier !== MMI.Unknown &&
+        !(AbstractVector{<:MMI.Finite} <: target_scitype_classifier)
+        throw(
+            ArgumentError(
+                "classifier does not support vectors of multi-class labels as targets"
+            ),
+        )
+    end
+    return nothing
+end
+function _check_model_supports_multiclass_predictions(classifier)
+    if !(
+        MMI.predict_scitype(classifier) <: Union{
+            MMI.Unknown,
+            AbstractVector{<:MMI.Finite},
+            AbstractVector{<:MMI.Density{<:MMI.Finite}},
+        }
+    )
+        throw(
+            ArgumentError(
+                "classifier does not support vectors of multi-class labels or their densities as predictions",
+            ),
+        )
+    end
+    return nothing
+end
+
+_astable(x::AbstractVecOrMat) = Tables.table(x)
+_astable(x) = Tables.istable(x) ? x : throw(ArgumentError("Argument is not a valid table"))
+
+# Workaround for https://github.com/JuliaAI/MLJBase.jl/issues/863
+# `MLJModelInterface.predict` sometimes returns predictions and sometimes predictions + additional information
+# TODO: Remove once the upstream issue is fixed
+function _predict(model::MMI.Model, fitresult, x)
+    y = MMI.predict(model, fitresult, x)
+    return if :predict in MMI.reporting_operations(model)
+        first(y)
+    else
+        y
+    end
+end
+
+# R⋆ for deterministic predictions (algorithm 1)
+function _rstar(
+    ::Type{<:AbstractVector{<:MMI.Finite}},
+    predictions::AbstractVector,
+    ytest::AbstractVector,
+)
+    length(predictions) == length(ytest) ||
+        error("numbers of predictions and targets must be equal")
+    mean_accuracy = Statistics.mean(p == y for (p, y) in zip(predictions, ytest))
+    nclasses = length(MMI.classes(ytest))
+    return nclasses * mean_accuracy
+end
+
+# R⋆ for probabilistic predictions (algorithm 2)
+function _rstar(
+    ::Type{<:AbstractVector{<:MMI.Density{<:MMI.Finite}}},
+    predictions::AbstractVector,
+    ytest::AbstractVector,
+)
+    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 scaled_distribution
+end
+
+# unsupported types of predictions and targets
+function _rstar(::Any, predictions, targets)
+    return throw(
+        ArgumentError(
+            "unsupported types of predictions ($(typeof(predictions))) and targets ($(typeof(targets)))",
+        ),
+    )
+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,14 +1,9 @@
 module MCMCDiagnosticTools
 
 using AbstractFFTs: AbstractFFTs
-using DataAPI: DataAPI
-using DataStructures: DataStructures
-using Distributions: Distributions
-using MLJModelInterface: MLJModelInterface as MMI
+using AliasTables: AliasTables
 using SpecialFunctions: SpecialFunctions
-using StatsBase: StatsBase
 using StatsFuns: StatsFuns, sqrt2
-using Tables: Tables
 
 using LinearAlgebra: LinearAlgebra
 using Random: Random
@@ -35,4 +30,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)
@@ -206,33 +206,44 @@ are independent random variables.
 function simulate_DAR1!(X::Matrix{Int}, phi::Float64, sampler)
     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] = rand(rng, sampler)
 
         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)
+                rand(rng, sampler)
             end
         end
     end
 
     return X
 end
 
-function simulate_MC(N::Int, P::Matrix{Float64})
+# Construct a sampler of the next state of a Markov chain for each current state from the
+# rows of the transition matrix `P`. States without transition probabilities (i.e. that
+# were never observed to transition) always transition to the first state.
+function transition_samplers(P::Matrix{Float64})
+    return map(eachrow(P)) do p
+        all(iszero, p) && return AliasTables.AliasTable(Float64.(eachindex(p) .== 1))
+        return AliasTables.AliasTable(p)
+    end
+end
+
+function simulate_MC(N::Int, samplers::AbstractVector{<:AliasTables.AliasTable})
     X = zeros(Int, N)
-    n, m = size(P)
-    X[1] = StatsBase.sample(1:n)
+    rng = Random.default_rng()
+    X[1] = rand(rng, 1:length(samplers))
     for i in 2:N
-        X[i] = StatsBase.wsample(1:n, vec(P[X[i - 1], :]))
+        X[i] = rand(rng, samplers[X[i - 1]])
     end
     return X
 end
@@ -303,18 +314,18 @@ 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))
+                sampler_phat = AliasTables.AliasTable(phat)
                 bstats = zeros(nsim)
                 Y = Matrix{Int}(undef, t, d)
                 for b in 1:nsim
@@ -326,9 +337,10 @@ function diag_all(
                 df0 = Statistics.mean(non_nan_bstats)
                 pval = Statistics.mean(stat <= x for x in non_nan_bstats)
             elseif method == :MCBOOT
+                samplers_mP = transition_samplers(mP)
                 bstats = zeros(Float64, nsim)
                 for b in 1:nsim
-                    Y = reduce(hcat, [simulate_MC(t, mP) for j in 1:d])
+                    Y = reduce(hcat, [simulate_MC(t, samplers_mP) for j in 1:d])
                     s = hangartner_inner(Y, m)[1]
                     bstats[b] = s
                 end
@@ -339,13 +351,14 @@ 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
+                samplers_mP = transition_samplers(mP)
                 bstats = zeros(Float64, nsim)
                 for b in 1:nsim
-                    Y = reduce(hcat, [simulate_MC(t, mP) for j in 1:d])
+                    Y = reduce(hcat, [simulate_MC(t, samplers_mP) for j in 1:d])
                     (s, sd) = bd_inner(Y, m)[1:2]
                     bstats[b] = s / sd
                 end

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