Merge pull request #14 from JuliaAI/expose-feature-names

Expose feature names

Author: ablaom

Project.toml

@@ -7,10 +7,12 @@ version = "0.1.4"
 DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 DecisionTree = "0.10"
 MLJModelInterface = "^0.3,^0.4, 1.0"
+Tables = "1.6"
 julia = "1.6"
 
 [extras]

src/MLJDecisionTreeInterface.jl

@@ -3,6 +3,7 @@ module MLJDecisionTreeInterface
 import MLJModelInterface
 using MLJModelInterface.ScientificTypesBase
 import DecisionTree
+import Tables
 
 using Random
 import Random.GLOBAL_RNG
@@ -43,9 +44,16 @@ MMI.@mlj_model mutable struct DecisionTreeClassifier <: MMI.Probabilistic
 end
 
 function MMI.fit(m::DecisionTreeClassifier, verbosity::Int, X, y)
+    schema = Tables.schema(X)
     Xmatrix = MMI.matrix(X)
     yplain  = MMI.int(y)
 
+    if schema === nothing
+        features = [Symbol("x$j") for j in 1:size(Xmatrix, 2)]
+    else
+        features = schema.names |> collect
+    end
+
     classes_seen  = filter(in(unique(y)), MMI.classes(y[1]))
     integers_seen = MMI.int(classes_seen)
 
@@ -61,11 +69,12 @@ function MMI.fit(m::DecisionTreeClassifier, verbosity::Int, X, y)
     end
     verbosity < 2 || DT.print_tree(tree, m.display_depth)
 
-    fitresult = (tree, classes_seen, integers_seen)
+    fitresult = (tree, classes_seen, integers_seen, features)
 
     cache  = nothing
     report = (classes_seen=classes_seen,
-              print_tree=TreePrinter(tree))
+              print_tree=TreePrinter(tree),
+              features=features)
 
     return fitresult, cache, report
 end
@@ -76,7 +85,9 @@ function get_encoding(classes_seen)
 end
 
 MMI.fitted_params(::DecisionTreeClassifier, fitresult) =
-    (tree=fitresult[1], encoding=get_encoding(fitresult[2]))
+    (tree=fitresult[1],
+     encoding=get_encoding(fitresult[2]),
+     features=fitresult[4])
 
 function smooth(scores, smoothing)
     iszero(smoothing) && return scores
@@ -394,6 +405,9 @@ The fields of `fitted_params(mach)` are:
   of tree (obtained by calling `fit!(mach, verbosity=2)` or from
   report - see below)
 
+- `features`: the names of the features encountered in training, in an
+  order consistent with the output of `print_tree` (see below)
+
 
 # Report
 
@@ -405,6 +419,9 @@ The fields of `report(mach)` are:
   tree, with single argument the tree depth; interpretation requires
   internal integer-class encoding (see "Fitted parameters" above).
 
+- `features`: the names of the features encountered in training, in an
+  order consistent with the output of `print_tree` (see below)
+
 
 # Examples
 

test/runtests.jl

@@ -37,11 +37,15 @@ yhat = MLJBase.predict(baretree, fitresult, X);
 yyhat = predict_mode(baretree, fitresult, MLJBase.selectrows(X, 1:3))
 @test MLJBase.classes(yyhat[1]) == MLJBase.classes(y[1])
 
+# check report and fitresult fields:
+@test Set([:classes_seen, :print_tree, :features]) == Set(keys(report))
+@test Set(report.classes_seen) == Set(levels(y))
+@test report.print_tree(2) === nothing # :-(
+@test report.features == [:sepal_length, :sepal_width, :petal_length, :petal_width]
+fp = fitted_params(baretree, fitresult)
+@test Set([:tree, :encoding, :features]) == Set(keys(fp))
+@test fp.features == report.features
 
-# # testing machine interface:
-# tree = machine(baretree, X, y)
-# fit!(tree)
-# yyhat = predict_mode(tree, MLJBase.selectrows(X, 1:3))
 using Random: seed!
 seed!(0)