diff --git a/Project.toml b/Project.toml index b22c375..9f8942c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "EvoTrees" uuid = "f6006082-12f8-11e9-0c9c-0d5d367ab1e5" authors = ["jeremiedb "] -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 3a0b24c..a3eee50 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 1adde30..7d5c8cb 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 1c77f6e..a2132f8 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 1821934..8a63d70 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