expose feature names in fitresult and report

Author: ablaom

src/MLJDecisionTreeInterface.jl

@@ -3,6 +3,7 @@ module MLJDecisionTreeInterface
 import MLJModelInterface
 using MLJModelInterface.ScientificTypesBase
 import DecisionTree
+import Tables
 
 using Random
 import Random.GLOBAL_RNG
@@ -50,7 +51,7 @@ function MMI.fit(m::DecisionTreeClassifier, verbosity::Int, X, y)
     if schema === nothing
         features = [Symbol("x$j") for j in 1:size(Xmatrix, 2)]
     else
-        features = schema.names
+        features = schema.names |> collect
     end
 
     classes_seen  = filter(in(unique(y)), MMI.classes(y[1]))
@@ -84,7 +85,9 @@ function get_encoding(classes_seen)
 end
 
 MMI.fitted_params(::DecisionTreeClassifier, fitresult) =
-    (tree=fitresult[1], encoding=get_encoding(fitresult[2]), features=features)
+    (tree=fitresult[1],
+     encoding=get_encoding(fitresult[2]),
+     features=fitresult[4])
 
 function smooth(scores, smoothing)
     iszero(smoothing) && return scores
@@ -402,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
 
@@ -413,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)