perf improvements to MultinomialSampler

src/multivariate/multinomial.jl

@@ -168,7 +168,26 @@ end
 _rand!(rng::AbstractRNG, d::Multinomial, x::AbstractVector{<:Real}) =
     multinom_rand!(rng, ntrials(d), probs(d), x)
 
-sampler(d::Multinomial) = MultinomialSampler(ntrials(d), probs(d))
+function sampler(d::Multinomial)
+    n = ntrials(d)
+    p = probs(d)
+    k = length(p)
+
+    # the constant λ term should be proportional to the perfomance ratio of
+    # constructing and sampling from a Binomial vs sampling from an AliasTable
+    λ = 10
+    return if n > λ * k
+        MultinomialSamplerBinomial(n, k, p)
+    else
+        alias = AliasTable(p)
+        MultinomialSamplerSequential(
+            n,
+            k,
+            alias,
+            Vector{typeof(alias.at.mask)}(undef, n),
+        )
+    end
+end
 
 
 ## Fit model

src/samplers/multinomial.jl

@@ -39,31 +39,34 @@ function multinom_rand!(rng::AbstractRNG, n::Int, p::AbstractVector{<:Real},
     return x
 end
 
-struct MultinomialSampler{T<:Real} <: Sampleable{Multivariate,Discrete}
+struct MultinomialSamplerBinomial{T<:Real} <: Sampleable{Multivariate,Discrete}
     n::Int
+    k::Int
     prob::Vector{T}
+end
+
+struct MultinomialSamplerSequential{S} <: Sampleable{Multivariate,Discrete}
+    n::Int
+    k::Int
     alias::AliasTable
+    scratch_alias_rng::Vector{S}
 end
 
-function MultinomialSampler(n::Int, prob::Vector{<:Real})
-    return MultinomialSampler(n, prob, AliasTable(prob))
+Base.@deprecate MultinomialSampler(n::Int, prob::Vector{<:Real}) MultinomialSamplerBinomial(n, length(prob), prob) false
+
+function _rand!(rng::AbstractRNG, s::MultinomialSamplerBinomial, x::AbstractVector{<:Real})
+    multinom_rand!(rng, s.n, s.prob, x)
 end
 
-function _rand!(rng::AbstractRNG, s::MultinomialSampler,
-                x::AbstractVector{<:Real})
-    n = s.n
-    k = length(s)
-    if n^2 > k
-        multinom_rand!(rng, n, s.prob, x)
-    else
-        # Use an alias table
-        fill!(x, zero(eltype(x)))
-        a = s.alias
-        for i = 1:n
-            x[rand(rng, a)] += 1
-        end
+function _rand!(rng::AbstractRNG, s::MultinomialSamplerSequential, x::AbstractVector{<:Real})
+    fill!(x, zero(eltype(x)))
+    rand!(rng, s.scratch_alias_rng)
+
+    for r in s.scratch_alias_rng
+        x[AliasTables.sample(r, s.alias.at)] += 1
     end
     return x
 end
 
-length(s::MultinomialSampler) = length(s.prob)
+length(s::MultinomialSamplerBinomial) = s.k
+length(s::MultinomialSamplerSequential) = s.k