Merge pull request #111 from jrfiedler/purity_function

Replace leftover util.entropy
JuliaAI · Jan 4, 2020 · a4e88e9 · a4e88e9
2 parents 31bbd90 + 5d18bc2
commit a4e88e9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/src/classification/tree.jl b/src/classification/tree.jl
@@ -192,7 +192,7 @@ module treeclassifier
 
         # no splits honor min_samples_leaf
         @inbounds if (unsplittable
-            || (best_purity / nt + util.entropy(nc, nt) < min_purity_increase))
+            || (best_purity / nt + purity_function(nc, nt) < min_purity_increase))
             node.is_leaf = true
             return
         else

diff --git a/test/classification/digits.jl b/test/classification/digits.jl
@@ -39,6 +39,21 @@ t = DecisionTree.build_tree(
         min_purity_increase)
 @test length(t) == 54
 
+# test that all purity decisions are based on passed-in purity function;
+# if so, this should be same as previous test
+entropy1000(ns, n) = DecisionTree.util.entropy(ns, n) * 1000
+min_samples_leaf    = 3
+min_samples_split   = 5
+min_purity_increase = 0.05 * 1000
+t = DecisionTree.build_tree(
+        Y, X,
+        n_subfeatures, max_depth,
+        min_samples_leaf,
+        min_samples_split,
+        min_purity_increase,
+        loss = entropy1000)
+@test length(t) == 54
+
 n_subfeatures       = 3
 n_trees             = 10
 partial_sampling    = 0.7