fix #267

Evovest · Mar 31, 2024 · 6e92369 · 6e92369
1 parent b4dfe6f
commit 6e92369
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "EvoTrees"
 uuid = "f6006082-12f8-11e9-0c9c-0d5d367ab1e5"
 authors = ["jeremiedb <[email protected]>"]
-version = "0.16.6"
+version = "0.16.7"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"

diff --git a/ext/EvoTreesCUDAExt/init.jl b/ext/EvoTreesCUDAExt/init.jl
@@ -7,6 +7,7 @@ function EvoTrees.init_core(params::EvoTrees.EvoTypes{L}, ::Type{<:EvoTrees.GPU}
     T = Float32
 
     target_levels = nothing
+    target_isordered = false
     if L == EvoTrees.Logistic
         @assert eltype(y_train) <: Real && minimum(y_train) >= 0 && maximum(y_train) <= 1
         K = 1
@@ -22,6 +23,7 @@ function EvoTrees.init_core(params::EvoTrees.EvoTypes{L}, ::Type{<:EvoTrees.GPU}
     elseif L == EvoTrees.MLogLoss
         if eltype(y_train) <: EvoTrees.CategoricalValue
             target_levels = EvoTrees.CategoricalArrays.levels(y_train)
+            target_isordered = isordered(y_train)
             y = UInt32.(EvoTrees.CategoricalArrays.levelcode.(y_train))
         elseif eltype(y_train) <: Integer || eltype(y_train) <: Bool || eltype(y_train) <: String || eltype(y_train) <: Char
             target_levels = sort(unique(y_train))
@@ -89,6 +91,7 @@ function EvoTrees.init_core(params::EvoTrees.EvoTypes{L}, ::Type{<:EvoTrees.GPU}
     info = Dict(
         :fnames => fnames,
         :target_levels => target_levels,
+        :target_isordered => target_isordered,
         :edges => edges,
         :featbins => featbins,
         :feattypes => feattypes,

diff --git a/src/MLJ.jl b/src/MLJ.jl
@@ -48,7 +48,7 @@ end
 
 function predict(::EvoTreeClassifier, fitresult, A)
   pred = predict(fitresult, A)
-  return MMI.UnivariateFinite(fitresult.info[:target_levels], pred, pool=missing)
+  return MMI.UnivariateFinite(fitresult.info[:target_levels], pred, pool=missing, ordered=fitresult.info[:target_isordered])
 end
 
 function predict(::EvoTreeCount, fitresult, A)

diff --git a/src/init.jl b/src/init.jl
@@ -7,6 +7,7 @@ function init_core(params::EvoTypes{L}, ::Type{CPU}, data, fnames, y_train, w, o
     T = Float32
 
     target_levels = nothing
+    target_isordered = false
     if L == Logistic
         @assert eltype(y_train) <: Real && minimum(y_train) >= 0 && maximum(y_train) <= 1
         K = 1
@@ -22,6 +23,7 @@ function init_core(params::EvoTypes{L}, ::Type{CPU}, data, fnames, y_train, w, o
     elseif L == MLogLoss
         if eltype(y_train) <: CategoricalValue
             target_levels = CategoricalArrays.levels(y_train)
+            target_isordered = isordered(y_train)
             y = UInt32.(CategoricalArrays.levelcode.(y_train))
         elseif eltype(y_train) <: Integer || eltype(y_train) <: Bool || eltype(y_train) <: String || eltype(y_train) <: Char
             target_levels = sort(unique(y_train))
@@ -87,6 +89,7 @@ function init_core(params::EvoTypes{L}, ::Type{CPU}, data, fnames, y_train, w, o
     info = Dict(
         :fnames => fnames,
         :target_levels => target_levels,
+        :target_isordered => target_isordered,
         :edges => edges,
         :featbins => featbins,
         :feattypes => feattypes,

diff --git a/test/MLJ.jl b/test/MLJ.jl
@@ -397,3 +397,21 @@ end
     fit!(mach)
     predict(mach, X)
 end
+
+##################################################
+### issue #267: ordered target
+##################################################
+using CategoricalArrays
+y = categorical(collect("cbbba"), levels=['b', 'a', 'c'], ordered=true)
+lvls = levels(y)
+eltype(y) <: CategoricalValue
+isordered(y)
+
+using MLJBase, EvoTrees
+# using StatisticalMeasures
+X = (; x=rand(10))
+y = coerce(rand("ab", 10), OrderedFactor)
+model = EvoTreeClassifier()
+mach = machine(model, X, y) |> fit!
+yhat = predict(mach, X)
+@assert isordered(yhat)