Fix extensions

Author: giopaglia

Project.toml

@@ -83,7 +83,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MLJDecisionTreeInterface = "c6f25543-311c-4c74-83dc-3ea6d1015661"
+MLJXGBoostInterface = "54119dfa-1dab-4055-a167-80440f4f7a91"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 
 [targets]
-test = ["Test", "DataFrames", "Random", "MLJ", "MultiData", "Markdown", "InteractiveUtils", "BenchmarkTools", "MLJBase", "XGBoost", "DecisionTree", "MLJDecisionTreeInterface", "SoleData"]
+test = ["Test", "DataFrames", "Random", "MLJ", "MLJXGBoostInterface", "MultiData", "Markdown", "InteractiveUtils", "BenchmarkTools", "MLJBase", "XGBoost", "DecisionTree", "MLJDecisionTreeInterface", "SoleData"]

ext/DecisionTreeExt.jl

@@ -2,9 +2,56 @@ module DecisionTreeExt
 
 using SoleModels
 import SoleModels: solemodel
+import SoleModels: alphabet
 
 import DecisionTree as DT
 
+function get_condition(featid, featval, featurenames)
+    test_operator = (<)
+    # @show fieldnames(typeof(tree))
+    feature = !isnothing(featurenames) ? VariableValue(featurenames[featid]) : VariableValue(featid)
+    return ScalarCondition(feature, test_operator, featval)
+end
+
+function SoleModels.alphabet(
+    model::Union{
+        DT.Ensemble,
+        DT.InfoNode,
+        DT.Node,
+        DT.Leaf,
+    },
+    args...;
+    kwargs...
+)
+
+    function _alphabet!(a::Vector, model::DT.Ensemble, args...; kwargs...)
+        map(t -> _alphabet!(a, t, args...; kwargs...), model.trees)
+        return a
+    end
+
+    function _alphabet!(a::Vector, model::DT.InfoNode, args...; kwargs...)
+        _alphabet!(a, model.left, args...; kwargs...)
+        _alphabet!(a, model.right, args...; kwargs...)
+        return a
+    end
+
+    function _alphabet!(a::Vector, model::DT.Node, args...;
+        featurenames = true,
+        kwargs...
+    )
+        featurenames = featurenames == true ? model.info.featurenames : featurenames
+        push!(a, Atom(get_condition(model.featid, model.featval, featurenames)))
+        return a
+    end
+
+    function _alphabet!(a::Vector, model::DT.Leaf, args...; kwargs...)
+        return a
+    end
+    
+    return SoleData.scalaralphabet(_alphabet!(Atom{ScalarCondition}[], model, args...; kwargs...))
+end
+
+
 function SoleModels.solemodel(
     model::DT.Ensemble,
     classlabels = nothing,
@@ -116,10 +163,7 @@ end
 
 function SoleModels.solemodel(tree::DT.Node; replace_classlabels = nothing, featurenames = nothing, keep_condensed = false)
     keep_condensed && error("Cannot keep condensed DecisionTree.Node.")
-    test_operator = (<)
-    # @show fieldnames(typeof(tree))
-    feature = !isnothing(featurenames) ? VariableValue(featurenames[tree.featid]) : VariableValue(tree.featid)
-    cond = ScalarCondition(feature, test_operator, tree.featval)
+    cond = get_condition(tree.featid, tree.featval, featurenames)
     antecedent = Atom(cond)
     lefttree = SoleModels.solemodel(tree.left; replace_classlabels, featurenames)
     righttree = SoleModels.solemodel(tree.right; replace_classlabels, featurenames)

ext/XGBoostExt.jl

@@ -1,34 +1,35 @@
 module XGBoostExt
 
+using SoleModels
 using XGBoost
 
-import Sole: alphabet, solemodel
-
-# TODO fix and test
-# function alphabet(model::XGBoost.Booster; kwargs...)
-#     function _alphabet!(a::Vector, model::XGBoost.Booster; kwargs...)
-#         return a
-#     end
-#     function _alphabet!(a::Vector, tree::XGBoost.Node; kwargs...)
-#         # Base case: if it's a leaf node
-#         if length(tree.children) == 0
-#             return a
-#         end
-
-#         # Recursive case: split node
-#         feature = Sole.VariableValue(tree.split isa String ? Symbol(tree.split) : tree.split)
-#         condition = ScalarCondition(feature, (<), tree.split_condition)
-#         push!(a, condition)
-#         if length(tree.children) == 2
-#             _alphabet!(a, tree.children[1]; with_stats, kwargs...)
-#             _alphabet!(a, tree.children[2]; with_stats, kwargs...)
-#         else
-#             error("Found $(length(tree.children)) children while 2 were expected: $(tree.children).")
-#         end
-#         return a
-#     end
-#     _alphabet!([], model; kwargs...)
-# end
+import SoleModels: alphabet, solemodel
+
+function alphabet(model::XGBoost.Booster; kwargs...)
+    error("TODO fix and test.")
+    function _alphabet!(a::Vector, model::XGBoost.Booster; kwargs...)
+        return a
+    end
+    function _alphabet!(a::Vector, tree::XGBoost.Node; kwargs...)
+        # Base case: if it's a leaf node
+        if length(tree.children) == 0
+            return a
+        end
+
+        # Recursive case: split node
+        feature = Sole.VariableValue(tree.split isa String ? Symbol(tree.split) : tree.split)
+        condition = ScalarCondition(feature, (<), tree.split_condition)
+        push!(a, condition)
+        if length(tree.children) == 2
+            _alphabet!(a, tree.children[1]; with_stats, kwargs...)
+            _alphabet!(a, tree.children[2]; with_stats, kwargs...)
+        else
+            error("Found $(length(tree.children)) children while 2 were expected: $(tree.children).")
+        end
+        return a
+    end
+    _alphabet!(Atom{ScalarCondition}[], model; kwargs...)
+end
 
 
 # TODO fix and test. Problem: where are the tree weights? How do I write this in the multi-class case?

test/XGBoostExt.jl

@@ -38,7 +38,7 @@ println("Test Accuracy: $acc")
 
 
 
-using Sole
+using SoleModels
 
 @test_nowarn alphabet(fitted_params(mach).fitresult[1])