Merge pull request #66 from JuliaAI/rand-fix

Fix implementation of `rand`

Author: ablaom

src/methods.jl

@@ -215,6 +215,8 @@ end
 """
     _cumulative(d::UnivariateFinite)
 
+**Private method.**
+
 Return the cumulative probability vector `C` for the distribution `d`,
 using only classes in the support of `d`, ordered according to the
 categorical elements used at instantiation of `d`. Used only to
@@ -238,6 +240,8 @@ end
 """
 _rand(rng, p_cumulative, R)
 
+**Private method.**
+
 Randomly sample the distribution with discrete support `R(1):R(n)`
 which has cumulative probability vector `p_cumulative` (see
 [`_cummulative`](@ref)).
@@ -256,26 +260,27 @@ function _rand(rng, p_cumulative, R)
     return index
 end
 
-function Base.rand(rng::AbstractRNG,
-                   d::UnivariateFinite{<:Any,<:Any,R}) where R
-    p_cumulative = _cumulative(d)
-    return Dist.support(d)[_rand(rng, p_cumulative, R)]
+Random.eltype(::Type{<:UnivariateFinite{<:Any,V}}) where V = V
+
+# The Sampler hook into Random's API is discussed in the Julia documentation, in the
+# Standard Library section on Random.
+function Random.Sampler(
+    ::AbstractRNG,
+    d::UnivariateFinite,
+    ::Random.Repetition,
+    )
+    data = (_cumulative(d), Dist.support(d))
+    Random.SamplerSimple(d, data)
 end
 
-function Base.rand(rng::AbstractRNG,
-                   d::UnivariateFinite{<:Any,<:Any,R},
-                   dim1::Int, moredims::Int...) where R # ref type
-    p_cumulative = _cumulative(d)
-    A = Array{R}(undef, dim1, moredims...)
-    for i in eachindex(A)
-        @inbounds A[i] = _rand(rng, p_cumulative, R)
-    end
-    support = Dist.support(d)
-    return broadcast(i -> support[i], A)
+function Base.rand(
+    rng::AbstractRNG,
+    sampler::Random.SamplerSimple{<:UnivariateFinite{<:Any,<:Any,R}},
+    ) where R
+    p_cumulative, support = sampler.data
+    return support[_rand(rng, p_cumulative, R)]
 end
 
-rng(d::UnivariateFinite, args...) = rng(Random.GLOBAL_RNG, d, args...)
-
 function Dist.fit(d::Type{<:UnivariateFinite},
                            v::AbstractVector{C}) where C
     C <: CategoricalValue ||

test/methods.jl

@@ -8,6 +8,7 @@ using StableRNGs
 import Random
 rng = StableRNG(123)
 using ScientificTypes
+import Random.default_rng
 
 import CategoricalDistributions: classes, ERR_NAN_FOUND
 
@@ -127,7 +128,7 @@ end
 @testset "broadcasting pdf over single UnivariateFinite object" begin
     d = UnivariateFinite(["a", "b"], [0.1, 0.9], pool=missing);
     @test pdf.(d, ["a", "b"]) == [0.1, 0.9]
-end 
+end
 
 @testset "constructor arguments not categorical values" begin
     @test_throws ArgumentError UnivariateFinite(Dict('f'=>0.7, 'q'=>0.2))
@@ -299,6 +300,28 @@ end
     @test displays_okay([5 + 3im, 4 - 7im])
 end
 
+@testset "rand signatures" begin
+    d = UnivariateFinite(
+        ["maybe", "no", "yes"],
+        [0.5, 0.4, 0.1];
+        pool=missing,
+    )
+
+    # smoke test:
+    sampler = Random.Sampler(default_rng(), d, Val(1))
+    rand(default_rng(), sampler)
+
+    Random.seed!(123)
+    samples = [rand(default_rng(), d) for i in 1:30]
+    Random.seed!(123)
+    @test [rand(d) for i in 1:30] == samples
+
+    Random.seed!(123)
+    samples = rand(Random.default_rng(), d, 3, 5)
+    Random.seed!(123)
+    @test samples == rand(d, 3, 5)
+end
+
 end # module
 
 true