Merge pull request #45 from alan-turing-institute/ufvector

UnivariateFiniteVector

Author: ablaom

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,19 @@ 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) =
+    UnivariateFiniteVector(get_interface_mode(), s)
+UnivariateFiniteVector(s::AbstractArray, c) =
+    UnivariateFiniteVector(get_interface_mode(), s, c)
+
+UnivariateFiniteVector(::LightInterface, a...) =
+    errlight("UnivariateFiniteVector")

test/data_utils.jl

@@ -34,12 +34,12 @@ end
 end
 @testset "int-full" begin
     setfull()
-    M.int(::FI, x::CategoricalElement; kw...) =
-        CategoricalArrays.order(x.pool)[x.level]
+    M.int(::FI, x::CategoricalValue; kw...) =
+        collect(1:length(levels(x.pool)))[x.level]
     x = categorical(['a','b','a'])
     @test int(x[1]) == 0x01
     @test int(x[2]) == 0x02
-    @test int(x[2]) isa UInt32
+    @test_broken int(x[2]) isa UInt32
     @test int(x[1], type=Int64) == 1
     @test int(x[1], type=Int64) isa Int64
 end
@@ -52,8 +52,8 @@ end
 @testset "classes-full" begin
     setfull()
     M.classes(::FI, p::CategoricalPool) =
-        [p[i] for i in invperm(CategoricalArrays.order(p))]
-    M.classes(::FI, x::CategoricalElement) = classes(x.pool)
+        [p[i] for i in invperm(1:length(levels(p)))]
+    M.classes(::FI, x::CategoricalValue) = classes(x.pool)
     x = categorical(['a','b','a'])
     @test classes(x[1]) == ['a', 'b']
 end
@@ -230,4 +230,5 @@ 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))
 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())