Add skeleton for XGBoostExt. Fix utilities.Introducedensembles.

Author: giopaglia

Project.toml

@@ -6,21 +6,19 @@ version = "0.9.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 Lazy = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SoleBase = "4475fa32-7023-44a0-aa70-4813b230e492"
@@ -34,15 +32,18 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
 
 [weakdeps]
 DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
+XGBoost = "009559a3-9522-5dbb-924b-0b6ed2b22bb9"
 MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
+MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MLJDecisionTreeInterface = "c6f25543-311c-4c74-83dc-3ea6d1015661"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 
 [extensions]
 DecisionTreeExt = "DecisionTree"
-MLJExt = "MLJ"
 MLJDecisionTreeInterfaceExt = "MLJDecisionTreeInterface"
+MLJExt = "MLJ"
 MLJModelInterfaceExt = "MLJModelInterface"
+XGBoostExt = "XGBoost"
 
 [compat]
 AbstractTrees = "0.4"
@@ -51,11 +52,16 @@ CSV = "0.10"
 CategoricalArrays = "0.10"
 DataFrames = "1"
 DataStructures = "0.18"
+DecisionTree = "0.12"
 FillArrays = "1"
 FunctionWrappers = "1"
 Graphs = "1.8"
 HTTP = "1.9"
+IterTools = "1"
 Lazy = "0.15.1"
+MLJ = "0.20"
+MLJBase = "1.6"
+MLJDecisionTreeInterface = "0.4"
 MLJModelInterface = "1.8.0"
 PrettyTables = "2.2"
 ProgressMeter = "1"
@@ -65,6 +71,7 @@ Revise = "3"
 SoleBase = "0.12"
 SoleData = "0.15, 0.16"
 SoleLogics = "0.11"
+XGBoost = "2.5"
 StatsBase = "0.30 - 0.34"
 Suppressor = "0.2"
 Tables = "1"
@@ -73,11 +80,15 @@ ZipFile = "0.10"
 julia = "1"
 
 [extras]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 MultiData = "8cc5100c-b3d1-4f82-90cb-0ea93d317aba"
 PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SoleData = "123f1ae1-6307-4526-ab5b-aab3a92a2b8c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Markdown", "InteractiveUtils", "PlutoUI", "MultiData"]
+test = ["Test", "DataFrames", "Random", "MLJ", "MultiData", "Markdown", "InteractiveUtils", "BenchmarkTools", "MLJBase", "XGBoost", "DecisionTree", "MLJDecisionTreeInterface", "SoleData"]

ext/DecisionTreeExt.jl

@@ -1,26 +1,80 @@
 module DecisionTreeExt
 
 using SoleModels
+import SoleModels: solemodel
 
 import DecisionTree as DT
 
-function SoleModels.solemodel(model::DT.Ensemble, args...; kwargs...)
-    return SoleModels.DecisionForest(map(t -> SoleModels.DecisionTree(SoleModels.solemodel(t, args...; kwargs...)), model.trees))
+function SoleModels.solemodel(
+    model::DT.Ensemble,
+    classlabels = nothing,
+    featurenames = nothing,
+    args...;
+    keep_condensed = true,
+    kwargs...
+)
+    if isnothing(classlabels)
+        error("Please, provide classlabels argument, as in solemodel(forest, classlabels; kwargs...). If your forest was trained via MLJ, use `classlabels = (mach).fitresult[2][sortperm((mach).fitresult[3])]`.")
+    end
+    if keep_condensed
+        info = (;
+            apply_preprocess=(y -> UInt32(findfirst(x -> x == y, classlabels))),
+            apply_postprocess=(y -> classlabels[y]),
+        )
+        keep_condensed = !keep_condensed
+        # O = UInt32
+    else
+        info = (;)
+        # O = UInt32
+    end
+    trees = map(t -> SoleModels.solemodel(t, args...; keep_condensed, featurenames, kwargs...), model.trees)
+    # trees = map(t -> let b = SoleModels.solemodel(t, args...; keep_condensed, featurenames, kwargs...); SoleModels.DecisionTree(b, 
+    #     (;
+    #         supporting_predictions=b.info[:supporting_predictions],
+    #         supporting_labels=b.info[:supporting_labels],
+    #     )
+    # ) end, model.trees)
+
+    if !isnothing(featurenames)
+        info = merge(info, (; featurenames=featurenames, ))
+    end
+
+    info = merge(info, (;
+            supporting_predictions=vcat([t.info[:supporting_predictions] for t in trees]...),
+            supporting_labels=vcat([t.info[:supporting_labels] for t in trees]...),
+        )
+    )
+
+    if !isnothing(classlabels)
+        O = eltype(classlabels)
+        # O = eltype(levels(classnames))
+    else
+        O = nothing
+    end
+
+    if isnothing(O)
+        m = DecisionEnsemble(trees, info)
+    else
+        m = DecisionEnsemble{O}(trees, info)
+    end
+    return m
 end
 
-function SoleModels.solemodel(tree::DT.InfoNode, keep_condensed = false; use_featurenames = true, kwargs...)
+function SoleModels.solemodel(tree::DT.InfoNode; keep_condensed = true, featurenames = true, classlabels = tree.info.classlabels, kwargs...)
     # @show fieldnames(typeof(tree))
-    use_featurenames = use_featurenames ? tree.info.featurenames : false
+    featurenames = featurenames == true ? tree.info.featurenames : featurenames
+    
     root, info = begin
         if keep_condensed
-            root = SoleModels.solemodel(tree.node; use_featurenames = use_featurenames, kwargs...)
+            root = SoleModels.solemodel(tree.node; featurenames, kwargs...)
             info = (;
-                apply_preprocess=(y -> UInt32(findfirst(x -> x == y, tree.info.classlabels))),
-                apply_postprocess=(y -> tree.info.classlabels[y]),
+                apply_preprocess=(y -> UInt32(findfirst(x -> x == y, classlabels))),
+                apply_postprocess=(y -> classlabels[y]),
             )
+            keep_condensed = !keep_condensed
             root, info
         else
-            root = SoleModels.solemodel(tree.node; replace_classlabels = tree.info.classlabels, use_featurenames = use_featurenames, kwargs...)
+            root = SoleModels.solemodel(tree.node; replace_classlabels = classlabels, featurenames, kwargs...)
             info = (;)
             root, info
         end
@@ -33,7 +87,19 @@ function SoleModels.solemodel(tree::DT.InfoNode, keep_condensed = false; use_fea
             supporting_labels=root.info[:supporting_labels],
         )
     )
-    return DecisionTree(root, info)
+    
+    # if !isnothing(classlabels)
+    #     O = eltype(classlabels)
+    # else
+    #     O = nothing
+    # end
+
+    # if isnothing(O)
+        dt = DecisionTree(root, info)
+    # else
+    #     dt = DecisionTree{O}(root, info)
+    # end
+    return dt
 end
 
 # function SoleModels.solemodel(tree::DT.Root)
@@ -48,22 +114,24 @@ end
 #     return DecisionTree(root, info)
 # end
 
-function SoleModels.solemodel(tree::DT.Node; replace_classlabels = nothing, use_featurenames = false)
+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 = (use_featurenames != false) ? VariableValue(use_featurenames[tree.featid]) : VariableValue(tree.featid)
+    feature = !isnothing(featurenames) ? VariableValue(featurenames[tree.featid]) : VariableValue(tree.featid)
     cond = ScalarCondition(feature, test_operator, tree.featval)
     antecedent = Atom(cond)
-    lefttree = SoleModels.solemodel(tree.left; replace_classlabels = replace_classlabels, use_featurenames = use_featurenames)
-    righttree = SoleModels.solemodel(tree.right; replace_classlabels = replace_classlabels, use_featurenames = use_featurenames)
+    lefttree = SoleModels.solemodel(tree.left; replace_classlabels, featurenames)
+    righttree = SoleModels.solemodel(tree.right; replace_classlabels, featurenames)
     info = (;
         supporting_predictions = [lefttree.info[:supporting_predictions]..., righttree.info[:supporting_predictions]...],
         supporting_labels = [lefttree.info[:supporting_labels]..., righttree.info[:supporting_labels]...],
     )
     return Branch(antecedent, lefttree, righttree, info)
 end
 
-function SoleModels.solemodel(tree::DT.Leaf; replace_classlabels = nothing, use_featurenames = false)
+function SoleModels.solemodel(tree::DT.Leaf; replace_classlabels = nothing, featurenames = nothing, keep_condensed = false)
+    keep_condensed && error("Cannot keep condensed DecisionTree.Node.")
     # @show fieldnames(typeof(tree))
     prediction = tree.majority
     labels = tree.values

ext/XGBoostExt.jl

@@ -0,0 +1,91 @@
+module XGBoostExt
+
+using XGBoost
+
+import Sole: alphabet, solemodel
+
+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
+
+
+# TODO fix and test. Problem: where are the tree weights? How do I write this in the multi-class case?
+
+# # Convert an XGBoost.Booster to a Sole Ensemble
+# function solemodel(model::XGBoost.Booster; with_stats::Bool = true, kwargs...)
+#     # Extract weights (global scaling factors for trees, if any)
+#     weights = nothing  # XGBoost trees usually don't have individual weights, but modify here if needed.
+
+#     model.params[:objective] == "multi:softprob" || error("Unexpected objective encountered: $(model.params[:objective]).")
+#     isempty(model.feature_names) || error("Unexpected objective encountered: $(model.params[:objective]).")
+
+#     # Convert all trees into Sole models
+#     trees = [solemodel(tree; with_stats, kwargs...) for tree in XGBoost.trees(model; with_stats = with_stats)]
+
+#     # Return a Sole Ensemble
+#     return Sole.Ensemble(trees, weights)
+# end
+
+# # Convert a single XGBoost tree into a Sole tree
+# function solemodel(tree::XGBoost.Node; with_stats::Bool = true, kwargs...)
+#     function _makeleaf(value)
+#         SoleModels.ConstantModel(value, (; supporting_predictions=[value]))
+#     end
+    
+#     # Base case: if it's a leaf node
+#     if length(tree.children) == 0
+#         return _makeleaf(tree.leaf)
+#     end
+
+#     # Recursive case: split node
+#     feature = Sole.VariableValue(tree.split isa String ? Symbol(tree.split) : tree.split)
+#     condition = ScalarCondition(feature, (<), tree.split_condition)
+#     antecedent = Atom(condition)
+    
+#     # Recursively convert left and right branches
+    
+#     if length(tree.children) == 2
+#         left_tree = solemodel(tree.children[1]; with_stats, kwargs...)
+#         right_tree = solemodel(tree.children[2]; with_stats, kwargs...)
+#     else
+#         error("Found $(length(tree.children)) children while 2 were expected: $(tree.children).")
+#     end
+
+#     # Aggregate info (e.g., supporting predictions) from children
+#     info = (;
+#         supporting_predictions=[left_tree.info[:supporting_predictions]..., right_tree.info[:supporting_predictions]...],
+#         this=_makeleaf(tree.leaf),
+#         xgboost_gain = tree.gain,
+#         xgboost_yes = tree.yes,
+#         xgboost_no = tree.no,
+#         xgboost_cover = tree.cover,
+#     )
+
+#     # Create and return a Sole Branch
+#     return Branch(antecedent, left_tree, right_tree, info)
+# end
+# solemodel(model)
+
+
+end

src/SoleModels.jl

@@ -56,8 +56,8 @@ export root
 export nnodes, nleaves
 export height
 
-export DecisionForest
-export trees
+export DecisionEnsemble, models
+export DecisionForest, trees
 
 export MixedModel
 

src/print.jl

@@ -459,7 +459,7 @@ end
 
 function printmodel(
     io::IO,
-    m::DecisionForest;
+    m::DecisionEnsemble;
     header = DEFAULT_HEADER,
     indentation_str = "",
     indentation = default_indentation,
@@ -496,22 +496,22 @@ function printmodel(
     end
 
     ########################################################################################
-    depth == 0 && show_symbols && print(io, "$(MODEL_SYMBOL) Forest of $(ntrees(m)) trees")
+    depth == 0 && show_symbols && print(io, "$(MODEL_SYMBOL) Ensemble{$(outcometype(m))} of $(nmodels(m)) models of type $(modelstype(m))")
     if isnothing(max_depth) || depth < max_depth
         _show_rule_metrics = show_rule_metrics
         println(io, "$(indentation_list_children)")
-        for (i_tree, tree) in enumerate(trees(m))
-            if i_tree < ntrees(m)
-                pipe = indentation_any_first*"[$i_tree/$(ntrees(m))]┐"
+        for (i_model, model) in enumerate(models(m))
+            if i_model < nmodels(m)
+                pipe = indentation_any_first*"[$i_model/$(nmodels(m))]┐"
                 pad_str = indentation_str*indentation_any_space*repeat(indentation_hspace, length(pipe)-length(indentation_any_space)-1-1)
                 ind_str = pad_str*indentation_last_space
             else
-                pipe = indentation_last_first*"[$i_tree/$(ntrees(m))]┐"
+                pipe = indentation_last_first*"[$i_model/$(nmodels(m))]┐"
                 ind_str = indentation_str*indentation_last_space*repeat(indentation_hspace, length(pipe)-length(indentation_last_space)-1-1)*indentation_last_space
             end
             print(io, pipe)
 
-            @_print_submodel io tree ind_str indentation depth max_depth show_subtree_info _show_rule_metrics show_subtree_metrics show_shortforms show_intermediate_finals tree_mode show_symbols syntaxstring_kwargs parenthesize_atoms kwargs
+            @_print_submodel io model ind_str indentation depth max_depth show_subtree_info _show_rule_metrics show_subtree_metrics show_shortforms show_intermediate_finals tree_mode show_symbols syntaxstring_kwargs parenthesize_atoms kwargs
         end
     else
         depth != 0 && print(io, " ")