added regression support

Signed-off-by: Manish Amde <[email protected]>

mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTreeRunner.scala

@@ -16,6 +16,7 @@
  */
 package org.apache.spark.mllib.tree
 
+import org.apache.spark.SparkContext._
 import org.apache.spark.{Logging, SparkContext}
 import org.apache.spark.mllib.tree.impurity.{Gini,Entropy,Variance}
 import org.apache.spark.rdd.RDD
@@ -95,6 +96,9 @@ object DecisionTreeRunner extends Logging {
     val accuracy = accuracyScore(model, testData)
     logDebug("accuracy = " + accuracy)
 
+    val mse = meanSquaredError(model,testData)
+    logDebug("mean square error = " + mse)
+
     sc.stop()
   }
 
@@ -126,6 +130,14 @@ object DecisionTreeRunner extends Logging {
     correctCount.toDouble / count
   }
 
+  //TODO: Make these generic MLTable metrics
+  def meanSquaredError(tree : DecisionTreeModel, data : RDD[LabeledPoint]) : Double = {
+    val meanSumOfSquares = data.map(y => (tree.predict(y.features) - y.label)*(tree.predict(y.features) - y.label)).mean()
+    println("meanSumOfSquares = " + meanSumOfSquares)
+    meanSumOfSquares
+  }
+
+
 
 
 }

mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Entropy.scala

@@ -16,6 +16,8 @@
  */
 package org.apache.spark.mllib.tree.impurity
 
+import javax.naming.OperationNotSupportedException
+
 object Entropy extends Impurity {
 
    def log2(x: Double) = scala.math.log(x) / scala.math.log(2)
@@ -31,4 +33,6 @@ object Entropy extends Impurity {
      }
    }
 
- }
+  def calculate(count: Double, sum: Double, sumSquares: Double): Double =
+    throw new OperationNotSupportedException("Entropy.calculate")
+}

mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Gini.scala

@@ -16,6 +16,8 @@
  */
 package org.apache.spark.mllib.tree.impurity
 
+import javax.naming.OperationNotSupportedException
+
 object Gini extends Impurity {
 
   def calculate(c0 : Double, c1 : Double): Double = {
@@ -29,4 +31,5 @@ object Gini extends Impurity {
     }
   }
 
- }
+  def calculate(count: Double, sum: Double, sumSquares: Double): Double = throw new OperationNotSupportedException("Gini.calculate")
+}

mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Impurity.scala

@@ -20,4 +20,6 @@ trait Impurity extends Serializable {
 
   def calculate(c0 : Double, c1 : Double): Double
 
+  def calculate(count : Double, sum : Double, sumSquares : Double) : Double
+
 }

mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Variance.scala

@@ -17,7 +17,14 @@
 package org.apache.spark.mllib.tree.impurity
 
 import javax.naming.OperationNotSupportedException
+import org.apache.spark.Logging
 
-object Variance extends Impurity {
+object Variance extends Impurity with Logging {
    def calculate(c0: Double, c1: Double): Double = throw new OperationNotSupportedException("Variance.calculate")
- }
+
+  def calculate(count: Double, sum: Double, sumSquares: Double): Double = {
+    val squaredLoss = sumSquares - (sum*sum)/count
+    squaredLoss/count
+  }
+
+}

mllib/src/test/scala/org/apache/spark/mllib/tree/DecisionTreeSuite.scala

@@ -49,7 +49,7 @@ class DecisionTreeSuite extends FunSuite with BeforeAndAfterAll {
     val arr = DecisionTreeSuite.generateOrderedLabeledPointsWithLabel1()
     assert(arr.length == 1000)
     val rdd = sc.parallelize(arr)
-    val strategy = new Strategy(Regression,Gini,3,100)
+    val strategy = new Strategy(Classification,Gini,3,100)
     val (splits, bins) = DecisionTree.find_splits_bins(rdd,strategy)
     assert(splits.length==2)
     assert(bins.length==2)
@@ -62,7 +62,7 @@ class DecisionTreeSuite extends FunSuite with BeforeAndAfterAll {
     val arr = DecisionTreeSuite.generateOrderedLabeledPointsWithLabel0()
     assert(arr.length == 1000)
     val rdd = sc.parallelize(arr)
-    val strategy = new Strategy(Regression,Gini,3,100)
+    val strategy = new Strategy(Classification,Gini,3,100)
     val (splits, bins) = DecisionTree.find_splits_bins(rdd,strategy)
     assert(splits.length==2)
     assert(splits(0).length==99)
@@ -88,7 +88,7 @@ class DecisionTreeSuite extends FunSuite with BeforeAndAfterAll {
     val arr = DecisionTreeSuite.generateOrderedLabeledPointsWithLabel1()
     assert(arr.length == 1000)
     val rdd = sc.parallelize(arr)
-    val strategy = new Strategy(Regression,Gini,3,100)
+    val strategy = new Strategy(Classification,Gini,3,100)
     val (splits, bins) = DecisionTree.find_splits_bins(rdd,strategy)
     assert(splits.length==2)
     assert(splits(0).length==99)
@@ -114,7 +114,7 @@ class DecisionTreeSuite extends FunSuite with BeforeAndAfterAll {
     val arr = DecisionTreeSuite.generateOrderedLabeledPointsWithLabel0()
     assert(arr.length == 1000)
     val rdd = sc.parallelize(arr)
-    val strategy = new Strategy(Regression,Entropy,3,100)
+    val strategy = new Strategy(Classification,Entropy,3,100)
     val (splits, bins) = DecisionTree.find_splits_bins(rdd,strategy)
     assert(splits.length==2)
     assert(splits(0).length==99)
@@ -139,7 +139,7 @@ class DecisionTreeSuite extends FunSuite with BeforeAndAfterAll {
     val arr = DecisionTreeSuite.generateOrderedLabeledPointsWithLabel1()
     assert(arr.length == 1000)
     val rdd = sc.parallelize(arr)
-    val strategy = new Strategy(Regression,Entropy,3,100)
+    val strategy = new Strategy(Classification,Entropy,3,100)
     val (splits, bins) = DecisionTree.find_splits_bins(rdd,strategy)
     assert(splits.length==2)
     assert(splits(0).length==99)