.

Author: tlienart

src/MLJModelInterface.jl

@@ -14,7 +14,7 @@ export LightInterface, FullInterface
 # MLJ model hierarchy
 export MLJType, Model, Supervised, Unsupervised,
        Probabilistic, Deterministic, Interval, Static,
-       UnivariateFinite
+       UnivariateFinite, UnivariateFiniteVector
 
 # model constructor + metadata
 export @mlj_model, metadata_pkg, metadata_model

src/data_utils.jl

@@ -292,7 +292,7 @@ _squeeze(v) = first(v)
 
 const UNIVARIATE_FINITE_DOCSTRING =
 """
-        UnivariateFinite(classes, p; pool=nothing, ordered=false)
+    UnivariateFinite(classes, p; pool=nothing, ordered=false)
 
 Construct a discrete univariate distribution whose finite support is
 the elements of the vector `classes`, and whose corresponding
@@ -329,3 +329,17 @@ UnivariateFinite(c::AbstractVector, p; kwargs...) =
 
 UnivariateFinite(::LightInterface, a...; kwargs...) =
     errlight("UnivariateFinite")
+
+const UNIVARIATE_FINITE_VECTOR_DOCSTRING =
+"""
+    UnivariateFiniteVector(scores, classes)
+
+Container for UnivariateFinite elements optimised for efficiency.
+Accessing a single element will construct and return the corresponding
+UnivariateFinite lazily.
+"""
+UnivariateFiniteVector(s::AbstractArray, a...) =
+    UnivariateFiniteVector(get_interface_mode(), s, a...)
+
+UnivariateFiniteVector(::LightInterface, a...) =
+    errlight("UnivariateFiniteVector")

test/data_utils.jl

@@ -34,7 +34,7 @@ end
 end
 @testset "int-full" begin
     setfull()
-    M.int(::FI, x::CategoricalElement; kw...) =
+    M.int(::FI, x::CategoricalValue; kw...) =
         CategoricalArrays.order(x.pool)[x.level]
     x = categorical(['a','b','a'])
     @test int(x[1]) == 0x01
@@ -53,7 +53,7 @@ end
     setfull()
     M.classes(::FI, p::CategoricalPool) =
         [p[i] for i in invperm(CategoricalArrays.order(p))]
-    M.classes(::FI, x::CategoricalElement) = classes(x.pool)
+    M.classes(::FI, x::CategoricalValue) = classes(x.pool)
     x = categorical(['a','b','a'])
     @test classes(x[1]) == ['a', 'b']
 end
@@ -230,4 +230,18 @@ end
     setlight()
     @test_throws M.InterfaceError UnivariateFinite(Dict(2=>3,3=>4))
     @test_throws M.InterfaceError UnivariateFinite(randn(2), randn(2))
+    @test_throws M.InterfaceError UnivariateFiniteVector(randn(2), randn(2))
+
+    setfull()
+    yc = categorical([1,2])
+    c = classes(yc[1])
+    s = rand()
+    u = UnivariateFinite(c, [1-s, s])
+    @test u isa MLJBase.UnivariateFinite
+    @test MLJBase.pdf(u, c[2]) == s
+
+    s = rand(5)
+    u = UnivariateFiniteVector(s)
+    @test u isa MLJBase.UnivariateFiniteVector
+    @test MLJBase.pdf.(u, u.classes[2]) == s
 end

test/runtests.jl

@@ -6,7 +6,6 @@ import MLJBase
 
 const M  = MLJModelInterface
 const FI = M.FullInterface
-const CategoricalElement = Union{CategoricalValue,CategoricalString}
 ScientificTypes.TRAIT_FUNCTION_GIVEN_NAME[:table] = Tables.istable
 
 setlight() = M.set_interface_mode(M.LightInterface())