From 4caa1269e1a663830887e248e980dc63494dfe3e Mon Sep 17 00:00:00 2001
From: Jeremie Desgagne-Bouchard <jeremiedb@users.noreply.github.com>
Date: Sun, 31 Mar 2024 16:03:04 -0400
Subject: [PATCH] Issue/267 (#268)

* fix #267

* fix #267
---
 Project.toml                |  2 +-
 ext/EvoTreesCUDAExt/init.jl |  3 +++
 src/MLJ.jl                  |  2 +-
 src/init.jl                 |  3 +++
 test/MLJ.jl                 | 12 ++++++++++++
 5 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index b22c3757..9f8942c7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "EvoTrees"
 uuid = "f6006082-12f8-11e9-0c9c-0d5d367ab1e5"
 authors = ["jeremiedb <jeremie.db@evovest.com>"]
-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
index 3a0b24c9..a3eee50b 100644
--- 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
index 1adde30c..7d5c8cb2 100644
--- 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
index 1c77f6ee..a2132f81 100644
--- 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
index 18219342..8a63d709 100644
--- a/test/MLJ.jl
+++ b/test/MLJ.jl
@@ -397,3 +397,15 @@ end
     fit!(mach)
     predict(mach, X)
 end
+
+##################################################
+### issue #267: ordered target
+##################################################
+@testset "MLJ - supported ordered factor predictions" begin
+    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)
+end