Add DecisionTree.jl models (#4)

* Add DecisionTree.jl models

* Update src/StatsLearnModels.jl

---------

Co-authored-by: Júlio Hoffimann <[email protected]>

Author: eliascarv

Project.toml

@@ -5,6 +5,7 @@ version = "0.2.0"
 
 [deps]
 ColumnSelectors = "9cc86067-7e36-4c61-b350-1ac9833d277f"
+DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
 TableTransforms = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
@@ -16,6 +17,7 @@ StatsLearnModelsMLJModelInterfaceExt = "MLJModelInterface"
 
 [compat]
 ColumnSelectors = "0.1"
+DecisionTree = "0.12"
 MLJModelInterface = "1.9"
 TableTransforms = "1.15"
 Tables = "1.11"

src/StatsLearnModels.jl

@@ -9,9 +9,23 @@ using ColumnSelectors: selector
 using TableTransforms: StatelessFeatureTransform
 import TableTransforms: applyfeat, isrevertible
 
+import DecisionTree as DT
+using DecisionTree: AdaBoostStumpClassifier, DecisionTreeClassifier, RandomForestClassifier
+using DecisionTree: DecisionTreeRegressor, RandomForestRegressor
+
 include("interface.jl")
+include("models/decisiontree.jl")
 include("learn.jl")
 
-export Learn
+export
+  # transform
+  Learn,
+  
+  # models
+  AdaBoostStumpClassifier,
+  DecisionTreeClassifier,
+  RandomForestClassifier,
+  DecisionTreeRegressor,
+  RandomForestRegressor
 
 end

src/models/decisiontree.jl

@@ -0,0 +1,24 @@
+const DTModel = Union{
+  AdaBoostStumpClassifier,
+  DecisionTreeClassifier,
+  RandomForestClassifier,
+  DecisionTreeRegressor,
+  RandomForestRegressor
+}
+
+function fit(model::DTModel, input, output)
+  cols = Tables.columns(output)
+  names = Tables.columnnames(cols)
+  outcol = first(names)
+  y = Tables.getcolumn(cols, outcol)
+  X = Tables.matrix(input)
+  DT.fit!(model, X, y)
+  FittedModel(model, outcol)
+end
+
+function predict(fmodel::FittedModel{<:DTModel}, table)
+  outcol = fmodel.cache
+  X = Tables.matrix(table)
+  ŷ = DT.predict(fmodel.model, X)
+  (; outcol => ŷ) |> Tables.materializer(table)
+end

test/runtests.jl

@@ -1,41 +1,52 @@
 using StatsLearnModels
-using MLJ, MLJDecisionTreeInterface
 using TableTransforms
 using DataFrames
 using Random
 using Test
 
+import MLJ, MLJDecisionTreeInterface
+
 const SLM = StatsLearnModels
 
 @testset "StatsLearnModels.jl" begin
+  iris = DataFrame(MLJ.load_iris())
+  input = iris[:, Not(:target)]
+  output = iris[:, [:target]]
+  train, test = MLJ.partition(1:nrow(input), 0.7, rng=123)
+
   @testset "interface" begin
-    Random.seed!(123)
-    iris = DataFrame(load_iris())
-    input = iris[:, Not(:target)]
-    output = iris[:, [:target]]
-    train, test = partition(1:nrow(input), 0.7, rng=123)
-    Tree = @load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
-    fmodel = SLM.fit(Tree(), input[train, :], output[train, :])
-    pred = SLM.predict(fmodel, input[test, :])
-    accuracy = count(pred.target .== output.target[test]) / length(test)
-    @test accuracy > 0.9
+    @testset "MLJ" begin
+      Random.seed!(123)
+      Tree = MLJ.@load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
+      fmodel = SLM.fit(Tree(), input[train, :], output[train, :])
+      pred = SLM.predict(fmodel, input[test, :])
+      accuracy = count(pred.target .== output.target[test]) / length(test)
+      @test accuracy > 0.9
+    end
+
+    @testset "DecisionTree" begin
+      Random.seed!(123)
+      model = DecisionTreeClassifier()
+      fmodel = SLM.fit(model, input[train, :], output[train, :])
+      pred = SLM.predict(fmodel, input[test, :])
+      accuracy = count(pred.target .== output.target[test]) / length(test)
+      @test accuracy > 0.9
+    end
   end
 
   @testset "Learn" begin
     Random.seed!(123)
-    iris = DataFrame(load_iris())
     outcol = :target
     incols = setdiff(propertynames(iris), [outcol])
-    train, test = partition(1:nrow(iris), 0.7, rng=123)
-    Tree = @load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
-    transform = Learn(iris[train, :], Tree(), incols => outcol)
+    model = DecisionTreeClassifier()
+    transform = Learn(iris[train, :], model, incols => outcol)
     @test !isrevertible(transform)
     pred = transform(iris[test, :])
     accuracy = count(pred.target .== iris.target[test]) / length(test)
     @test accuracy > 0.9
 
     # throws
     # training data is not a table
-    @test_throws ArgumentError Learn(nothing, Tree(), incols => outcol)
+    @test_throws ArgumentError Learn(nothing, model, incols => outcol)
   end
 end