Merge pull request #69 from JuliaAI/dev

For a 0.1.13 release

Author: ablaom

.github/codecov.yml

@@ -0,0 +1,8 @@
+coverage:
+  status:
+    project:
+      default:
+        threshold: 0.5%
+    patch:
+      default:
+        target: 80%

Project.toml

@@ -1,7 +1,7 @@
 name = "CategoricalDistributions"
 uuid = "af321ab8-2d2e-40a6-b165-3d674595d28e"
 authors = ["Anthony D. Blaom <[email protected]>"]
-version = "0.1.12"
+version = "0.1.13"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

src/methods.jl

@@ -1,6 +1,6 @@
 # not for export:
-const UnivariateFiniteUnion =
-    Union{UnivariateFinite, UnivariateFiniteArray}
+const UnivariateFiniteUnion{S,V,R,P} =
+    Union{UnivariateFinite{S,V,R,P}, UnivariateFiniteArray{S,V,R,P}}
 
 """
     classes(d::UnivariateFinite)
@@ -42,18 +42,34 @@ end
 raw_support(d::UnivariateFiniteUnion) = collect(keys(d.prob_given_ref))
 
 """
-    Dist.support(d::UnivariateFinite)
-    Dist.support(d::UnivariateFiniteArray)
+    Distributions.support(d::UnivariateFinite)
+    Distributions.support(d::UnivariateFiniteArray)
 
 Ordered list of classes associated with non-zero probabilities.
 
     v = categorical(["yes", "maybe", "no", "yes"])
     d = UnivariateFinite(v[1:2], [0.3, 0.7])
-    support(d) # CategoricalArray{String,1,UInt32}["maybe", "yes"]
+    Distributions.support(d) # CategoricalArray{String,1,UInt32}["maybe", "yes"]
 
 """
-Dist.support(d::UnivariateFiniteUnion) =
-    map(d.decoder, raw_support(d))
+Distributions.support(d::UnivariateFiniteUnion) = classes(d)[raw_support(d)]
+
+"""
+    fast_support(d::UnivariateFinite)
+
+Same as `Distributions.support(d)` except it returns a vector of `CategoricalValue`s,
+rather than a `CategoricalVector`. It executes faster, about five times faster for a
+three-class `UnivariateFinite` distribution.
+"""
+function fast_support(d::UnivariateFiniteUnion{S,V,R}) where {S,V,R}
+    raw_support = keys(d.prob_given_ref)
+    n = length(raw_support)
+    ret = Vector{CategoricalValue{V,R}}(undef, n)
+    for (i, ref) in enumerate(raw_support)
+        ret[i] = d.decoder(ref)
+    end
+    ret
+end
 
 # TODO: If I manually give a class zero probability, it will appear in
 # support, which is probably confusing. We may need two versions of
@@ -64,8 +80,7 @@ Dist.support(d::UnivariateFiniteUnion) =
 # not exported:
 sample_scitype(d::UnivariateFiniteUnion) = d.scitype
 
-CategoricalArrays.isordered(d::UnivariateFinite) = isordered(classes(d))
-CategoricalArrays.isordered(u::UnivariateFiniteArray) = isordered(classes(u))
+CategoricalArrays.isordered(d::UnivariateFiniteUnion) = isordered(classes(d))
 
 
 ## DISPLAY
@@ -96,8 +111,8 @@ probability pairs.  Returns `false` otherwise.
 
 """
 function Base.isapprox(d1::UnivariateFinite, d2::UnivariateFinite; kwargs...)
-    support1 = Dist.support(d1)
-    support2 = Dist.support(d2)
+    support1 = fast_support(d1)
+    support2 = fast_support(d2)
     for c in support1
         c in support2 || return false
         isapprox(pdf(d1, c), pdf(d2, c); kwargs...) ||
@@ -107,8 +122,8 @@ function Base.isapprox(d1::UnivariateFinite, d2::UnivariateFinite; kwargs...)
 end
 function Base.isapprox(d1::UnivariateFiniteArray,
                        d2::UnivariateFiniteArray; kwargs...)
-    support1 = Dist.support(d1)
-    support2 = Dist.support(d2)
+    support1 = fast_support(d1)
+    support2 = fast_support(d2)
     for c in support1
         c in support2 || return false
         isapprox(pdf.(d1, c), pdf.(d2, c); kwargs...) ||
@@ -206,22 +221,18 @@ function throw_nan_error_if_needed(x)
     end
 end
 
-# mode(v::Vector{UnivariateFinite}) = mode.(v)
-# mode(u::UnivariateFiniteVector{2}) =
-#     [u.support[ifelse(s > 0.5, 2, 1)] for s in u.scores]
-# mode(u::UnivariateFiniteVector{C}) where {C} =
-#     [u.support[findmax(s)[2]] for s in eachrow(u.scores)]
+
+# # HELPERS FOR RAND
 
 """
     _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
-implement random sampling from `d`. We have `C[1] == 0` and `C[end] ==
-1`, assuming the probabilities have been normalized.
+Return the cumulative probability vector `C` for the distribution `d`, using only classes
+in `Distributions.support(d)`, ordered according to the categorical elements used at
+instantiation of `d`. Used only to implement random sampling from `d`. We have `C[1] == 0`
+and `C[end] == 1`, assuming the probabilities have been normalized.
 
 """
 function _cumulative(d::UnivariateFinite{S,V,R,P}) where {S,V,R,P<:Real}
@@ -260,16 +271,54 @@ function _rand(rng, p_cumulative, R)
     return index
 end
 
-Random.eltype(::Type{<:UnivariateFinite{<:Any,V}}) where V = V
+
+# # RAND
+
+Random.eltype(::Type{<:UnivariateFinite{S,V,R}}) where {S,V,R} =
+    CategoricalArrays.CategoricalValue{V,R}
 
 # The Sampler hook into Random's API is discussed in the Julia documentation, in the
 # Standard Library section on Random.
+
+
+## Single samples
+
+Random.Sampler(::AbstractRNG, d::UnivariateFinite, ::Val{1}) = Random.SamplerTrivial(d)
+
+function Base.rand(
+    rng::AbstractRNG,
+    sampler::Random.SamplerTrivial{<:UnivariateFinite{<:Any,<:Any,V,P}},
+    ) where {V, P}
+
+    d = sampler[]
+    u = rand(rng)
+
+    total = zero(P)
+    
+    # For type stability we assign `zero(V)`` as the default ref
+    # This isn't a problem since we know that `rand` is always defined 
+    # as UnivariateFinite objects have non-negative probabilities,
+    # summing up to a non-negative value.
+    rng_key = zero(V)
+    for (ref, prob) in pairs(d.prob_given_ref)
+        total += prob
+        u <= total && begin
+            rng_key = ref
+            break
+        end
+    end
+    return d.decoder(rng_key)
+end
+
+
+## Multiple samples
+
 function Random.Sampler(
     ::AbstractRNG,
     d::UnivariateFinite,
     ::Random.Repetition,
     )
-    data = (_cumulative(d), Dist.support(d))
+    data = (_cumulative(d), fast_support(d))
     Random.SamplerSimple(d, data)
 end
 
@@ -281,6 +330,9 @@ function Base.rand(
     return support[_rand(rng, p_cumulative, R)]
 end
 
+
+## FIT
+
 function Dist.fit(d::Type{<:UnivariateFinite},
                            v::AbstractVector{C}) where C
     C <: CategoricalValue ||

test/methods.jl

@@ -28,6 +28,7 @@ A, S, Q, F = V[1], V[2], V[3], V[4]
     @test classes(d) == classes(s)
     @test levels(d) == levels(s)
     @test support(d) == [f, q, s]
+    @test support(d) == [CategoricalDistributions.fast_support(d)...]
     @test CategoricalDistributions.sample_scitype(d) == OrderedFactor{4}
     # levels!(v, reverse(levels(v)))
     # @test classes(d) == [s, q, f, a]
@@ -54,7 +55,7 @@ A, S, Q, F = V[1], V[2], V[3], V[4]
 
     N = 50
     rng = StableRNG(125)
-    samples = [rand(rng,d) for i in 1:50];
+    samples = [rand(rng, d) for i in 1:N];
     rng = StableRNG(125)
     @test samples == [rand(rng, d) for i in 1:N]
 
@@ -301,15 +302,18 @@ end
 end
 
 @testset "rand signatures" begin
-    d = UnivariateFinite(
-        ["maybe", "no", "yes"],
-        [0.5, 0.4, 0.1];
-        pool=missing,
-    )
+    dict = Dict(s=>0.1, q=>0.2, f=>0.7)
+    d    = UnivariateFinite(dict)
 
-    # smoke test:
     sampler = Random.Sampler(default_rng(), d, Val(1))
-    rand(default_rng(), sampler)
+    @test sampler isa Random.SamplerTrivial
+    sampler = Random.Sampler(default_rng(), d, Val(Inf))
+    @test sampler isa Random.SamplerSimple
+
+    # sampling one at a time, or all at once is the same:
+    rng0 = StableRNG(123)
+    samples = [rand(rng0, d) for i in 1:30]
+    @test samples == [rand(StableRNG(123), d, 30)...]
 
     Random.seed!(123)
     samples = [rand(default_rng(), d) for i in 1:30]