Merge pull request #35 from JuliaAI/raw_support

fix issue #34

Author: ablaom

src/types.jl

@@ -451,9 +451,11 @@ function _UnivariateFinite(support,
         issubset(support, _classes) ||
             error("Specified support, $support, not contained in "*
                   "specified pool, $(levels(classes)). ")
-        _support = filter(_classes) do c
-            c in support
-        end
+        idxs = getindex.(
+            Ref(CategoricalArrays.DataAPI.invrefpool(_classes)), 
+            support
+        )
+        _support = _classes[idxs] 
     end
 
     # calls core method:

test/types.jl

@@ -7,6 +7,7 @@ using StableRNGs
 using FillArrays
 using ScientificTypes
 import Random
+import CategoricalDistributions: classes
 
 # coverage of constructor testing is expanded in the other test files
 
@@ -36,6 +37,20 @@ end
 
     UnivariateFinite(supp, probs, pool=missing, augment=true);
 
+    # construction from pool and support does not 
+    # consist of categorical elements (See issue #34)
+    v = categorical(["x", "x", "y", "z", "y", "z", "p"])
+    probs1 = [0.1, 0.2, 0.7]
+    probs2 = [0.1 0.2 0.7; 0.5 0.2 0.3; 0.8 0.1 0.1]
+    unf1 = UnivariateFinite(["y", "x", "z"], probs1, pool=v)
+    unf2 = UnivariateFinite(["y", "x", "z"], probs2, pool=v)
+    @test CategoricalArrays.pool(classes(unf1)) == CategoricalArrays.pool(v)
+    @test CategoricalArrays.pool(classes(unf2)) == CategoricalArrays.pool(v)
+    @test pdf.(unf1, ["y", "x", "z"]) == probs1
+    @test pdf.(unf2, "y") == probs2[:, 1]
+    @test pdf.(unf2, "x") == probs2[:, 2]
+    @test pdf.(unf2, "z") == probs2[:, 3]
+
     # dimension mismatches:
     badprobs = rand(rng, 40, 3)
     @test_throws(CategoricalDistributions.err_dim(supp, badprobs),