add Categorical distribution

Author: rfourquet

src/RandomExtensions.jl

@@ -1,6 +1,7 @@
 module RandomExtensions
 
-export make, Uniform, Normal, Exponential, CloseOpen, OpenClose, OpenOpen, CloseClose, Rand, Bernoulli
+export make, Uniform, Normal, Exponential, CloseOpen, OpenClose, OpenOpen, CloseClose, Rand,
+       Bernoulli, Categorical
 
 # re-exports from Random, which don't overlap with new functionality and not from misc.jl
 export rand!, AbstractRNG, MersenneTwister, RandomDevice

src/distributions.jl

@@ -200,3 +200,34 @@ end
 
 Bernoulli(p::Real=0.5) = Bernoulli(Int, p)
 Bernoulli(::Type{T}, p::Real=0.5) where {T<:Number} = Bernoulli{T}(p)
+
+## Categorical
+
+struct Categorical{T<:Number} <: Distribution{T}
+    cdf::Vector{Float64}
+
+    function Categorical{T}(weigths) where T
+        if !isa(weigths, AbstractArray)
+            # necessary for accumulate
+            # TODO: will not be necessary anymore in Julia 1.5
+            weigths = collect(weigths)
+        end
+        weigths = vec(weigths)
+
+        isempty(weigths) &&
+            throw(ArgumentError("Categorical requires at least one category"))
+
+        s = Float64(sum(weigths))
+        cdf = accumulate(weigths; init=0.0) do x, y
+            x + Float64(y) / s
+        end
+        @assert isapprox(cdf[end], 1.0) # really?
+        cdf[end] = 1.0 # to be sure the algo terminates
+        new{T}(cdf)
+    end
+end
+
+Categorical(weigths) = Categorical{Int}(weigths)
+
+Categorical(n::Number) =
+    Categorical{typeof(n)}(Iterators.repeated(1.0 / Float64(n), Int(n)))

src/sampling.jl

@@ -136,6 +136,22 @@ rand(rng::AbstractRNG, sp::SamplerTag{Bernoulli{T}}) where {T} =
     ifelse(rand(rng, CloseOpen12()) < sp.data, one(T), zero(T))
 
 
+## Categorical
+
+Sampler(RNG::Type{<:AbstractRNG}, c::Categorical, n::Repetition) =
+    SamplerSimple(c, Sampler(RNG, CloseOpen(), n))
+
+# unfortunately requires @inline to avoid allocating
+@inline rand(rng::AbstractRNG, sp::SamplerSimple{Categorical{T}}) where {T} =
+    let c = rand(rng, sp.data)
+        T(findfirst(x -> x >= c, sp[].cdf))
+    end
+
+# NOTE:
+# if length(cdf) is somewhere between 150 and 200, the following gets faster:
+#   T(searchsortedfirst(sp[].cdf, rand(rng, sp.data)))
+
+
 ## random elements from pairs
 
 Sampler(RNG::Type{<:AbstractRNG}, t::Pair, n::Repetition) =

test/runtests.jl

@@ -62,6 +62,24 @@ using Test
         r = rand(Bernoulli(T, 1))
         @test r == 1
     end
+
+    # Categorical
+    n = rand(1:9)
+    @test rand(Categorical(n)) ∈ 1:9
+    @test all(∈(1:9), rand(Categorical(n), 10))
+    @test rand(Categorical(n)) isa Int
+    c = Categorical(Float64(n))
+    @test rand(c) isa Float64
+    @test rand(c) ∈ 1:9
+    c = Categorical([1, 7, 2])
+    # cf. Bernoulli tests
+    @test 620 < count(==(2), rand(c, 1000)) < 780
+    @test rand(c) isa Int
+    @test rand(Categorical{Float64}((1, 2, 3, 4))) isa Float64
+
+    @test_throws ArgumentError Categorical(())
+    @test_throws ArgumentError Categorical([])
+    @test_throws ArgumentError Categorical(x for x in 1:0)
 end
 
 const rInt8 = typemin(Int8):typemax(Int8)