removing threshold for classification predict method

witgo · Mar 13, 2014 · 632818f · 632818f
1 parent 2116360
commit 632818f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
@@ -1034,8 +1034,12 @@ object DecisionTree extends Serializable with Logging {
   /**
    * Calculates the classifier accuracy.
    */
-  def accuracyScore(model: DecisionTreeModel, data: RDD[LabeledPoint]): Double = {
-    val correctCount = data.filter(y => model.predict(y.features) == y.label).count()
+  def accuracyScore(model: DecisionTreeModel, data: RDD[LabeledPoint],
+                    threshold: Double = 0.5): Double = {
+    def predictedValue(features: Array[Double]) = {
+      if (model.predict(features) < threshold) 0.0 else 1.0
+    }
+    val correctCount = data.filter(y => predictedValue(y.features) == y.label).count()
     val count = data.count()
     logDebug("correct prediction count = " +  correctCount)
     logDebug("data count = " + count)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/model/DecisionTreeModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/model/DecisionTreeModel.scala
@@ -34,14 +34,7 @@ class DecisionTreeModel(val topNode: Node, val algo: Algo) extends Serializable
    * @return Double prediction from the trained model
    */
   def predict(features: Array[Double]): Double = {
-    algo match {
-      case Classification => {
-        if (topNode.predictIfLeaf(features) < 0.5) 0.0 else 1.0
-      }
-      case Regression => {
-        topNode.predictIfLeaf(features)
-      }
-    }
+    topNode.predictIfLeaf(features)
   }
 
   /**