Merge pull request #32 from JuliaAI/dev

For a 0.1.8 release

Author: ablaom

Project.toml

@@ -1,7 +1,7 @@
 name = "CategoricalDistributions"
 uuid = "af321ab8-2d2e-40a6-b165-3d674595d28e"
 authors = ["Anthony D. Blaom <[email protected]>"]
-version = "0.1.7"
+version = "0.1.8"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
@@ -25,6 +25,7 @@ julia = "1.3"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 
 [targets]
-test = ["Random", "StableRNGs", "Test"]
+test = ["Random", "StableRNGs", "Test", "FillArrays"]

src/methods.jl

@@ -49,7 +49,7 @@ 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", "no"]
+    support(d) # CategoricalArray{String,1,UInt32}["maybe", "yes"]
 
 """
 Dist.support(d::UnivariateFiniteUnion) =

src/types.jl

@@ -242,8 +242,8 @@ _check_augmentable(support, probs) = _check_probs_01(probs) &&
 
 ## AUGMENTING ARRAYS TO MAKE THEM PROBABILITY ARRAYS
 
-_unwrap(A::Array) = A
-_unwrap(A::Vector) = first(A)
+_unwrap(A::AbstractArray) = A
+_unwrap(A::AbstractVector) = first(A)
 
 isbinary(support) = length(support) == 2
 
@@ -276,6 +276,9 @@ function _augment_probs(::Val{true},
     return aug_probs
 end
 
+_array_or_scalar(x::Array) = x
+_array_or_scalar(x::AbstractArray) = copyto!(similar(Array{eltype(x)}, axes(x)), x) 
+_array_or_scalar(x) = x
 
 ## CONSTRUCTORS - FROM DICTIONARY
 
@@ -306,7 +309,7 @@ function UnivariateFinite(
     issubset(_support, parent_classes) ||
         error("Categorical elements are not from the same pool. ")
 
-    pairs = [int(c) => prob_given_class[c]
+    pairs = [int(c) => _array_or_scalar(prob_given_class[c])
                 for c in _support]
 
     probs1 = first(values(prob_given_class))
@@ -341,7 +344,7 @@ function UnivariateFinite(d::AbstractDict{V,<:Prob};
     _classes = classes(pool)
     issubset(raw_support, _classes) ||
         error("Specified support, $raw_support, not contained in "*
-              "specified pool, $(levels(classes)). ")
+              "specified pool, $(levels(_classes)). ")
     support = filter(_classes) do c
         c in raw_support
     end

test/types.jl

@@ -4,6 +4,7 @@ using Test
 using CategoricalDistributions
 using CategoricalArrays
 using StableRNGs
+using FillArrays
 using ScientificTypes
 import Random
 
@@ -55,6 +56,27 @@ end
                                   augment=true,
                                   pool=missing))
 
+    # Test construction from non `Array` `AbstractArray`
+    v = categorical(['x', 'x', 'y', 'x', 'z', 'w'])
+    probs_fillarray = FillArrays.Ones(100, 3)
+    probs_array = ones(100, 3)
+    
+    probs1_fillarray = FillArrays.Fill(0.2, 100, 2)
+    probs1_array = fill(0.2, 100, 2)
+
+    u_from_array = UnivariateFinite(['x', 'y', 'z'], probs_array, pool=v)
+    u_from_fillarray = UnivariateFinite(['x', 'y', 'z'], probs_fillarray, pool=v)
+
+    u1_from_array = UnivariateFinite(
+        ['x', 'y', 'z'], probs1_array, pool=v, augment=true
+    )
+    u1_from_fillarray = UnivariateFinite(
+        ['x', 'y', 'z'], probs1_fillarray, pool=v, augment=true
+    )
+    
+    @test u_from_array.prob_given_ref == u_from_fillarray.prob_given_ref
+    @test u1_from_array.prob_given_ref == u1_from_fillarray.prob_given_ref
+
     # autosupport:
     u = UnivariateFinite(probs, pool=missing, augment=true);
     probs = probs ./ sum(probs)