Addressing reviewers comments mengxr. Added confusion matrix

mllib/src/main/scala/org/apache/spark/mllib/evaluation/MulticlassMetrics.scala

@@ -17,34 +17,52 @@
 
 package org.apache.spark.mllib.evaluation
 
-import org.apache.spark.annotation.Experimental
-import org.apache.spark.rdd.RDD
 import org.apache.spark.Logging
 import org.apache.spark.SparkContext._
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.rdd.RDD
 
 import scala.collection.Map
 
 /**
  * ::Experimental::
  * Evaluator for multiclass classification.
  *
- * @param predictionsAndLabels an RDD of (prediction, label) pairs.
+ * @param predictionAndLabels an RDD of (prediction, label) pairs.
  */
 @Experimental
-class MulticlassMetrics(predictionsAndLabels: RDD[(Double, Double)]) extends Logging {
+class MulticlassMetrics(predictionAndLabels: RDD[(Double, Double)]) extends Logging {
 
-  private lazy val labelCountByClass: Map[Double, Long] = predictionsAndLabels.values.countByValue()
+  private lazy val labelCountByClass: Map[Double, Long] = predictionAndLabels.values.countByValue()
   private lazy val labelCount: Long = labelCountByClass.values.sum
-  private lazy val tpByClass: Map[Double, Int] = predictionsAndLabels
+  private lazy val tpByClass: Map[Double, Int] = predictionAndLabels
     .map { case (prediction, label) =>
-      (label, if (label == prediction) 1 else 0)
-    }.reduceByKey(_ + _)
+    (label, if (label == prediction) 1 else 0)
+  }.reduceByKey(_ + _)
     .collectAsMap()
-  private lazy val fpByClass: Map[Double, Int] = predictionsAndLabels
+  private lazy val fpByClass: Map[Double, Int] = predictionAndLabels
     .map { case (prediction, label) =>
-      (prediction, if (prediction != label) 1 else 0)
-    }.reduceByKey(_ + _)
+    (prediction, if (prediction != label) 1 else 0)
+  }.reduceByKey(_ + _)
     .collectAsMap()
+  private lazy val confusions = predictionAndLabels.map {
+    case (prediction, label) => ((prediction, label), 1)
+  }.reduceByKey(_ + _).collectAsMap()
+
+  /**
+   * Returns confusion matrix:
+   * predicted classes are in columns,
+   * they are ordered by class label ascending,
+   * as in "labels"
+   */
+  lazy val confusionMatrix: Array[Array[Int]] = {
+    val matrix = Array.ofDim[Int](labels.size, labels.size)
+    println(matrix.length, matrix(0).length)
+    for (i <- 0 to labels.size - 1; j <- 0 to labels.size - 1) {
+      matrix(j)(i) = confusions.getOrElse((labels(i), labels(j)), 0)
+    }
+    matrix
+  }
 
   /**
    * Returns true positive rate for a given label (category)
@@ -103,8 +121,8 @@ class MulticlassMetrics(predictionsAndLabels: RDD[(Double, Double)]) extends Log
   /**
    * Returns recall
    * (equals to precision for multiclass classifier
-   *  because sum of all false positives is equal to sum
-   *  of all false negatives)
+   * because sum of all false positives is equal to sum
+   * of all false negatives)
    */
   lazy val recall: Double = precision
 
@@ -114,6 +132,19 @@ class MulticlassMetrics(predictionsAndLabels: RDD[(Double, Double)]) extends Log
    */
   lazy val fMeasure: Double = precision
 
+  /**
+   * Returns weighted true positive rate
+   * (equals to precision, recall and f-measure)
+   */
+  lazy val weightedTruePositiveRate: Double = weightedRecall
+
+  /**
+   * Returns weighted false positive rate
+   */
+  lazy val weightedFalsePositiveRate: Double = labelCountByClass.map { case (category, count) =>
+    falsePositiveRate(category) * count.toDouble / labelCount
+  }.sum
+
   /**
    * Returns weighted averaged recall
    * (equals to precision, recall and f-measure)

mllib/src/test/scala/org/apache/spark/mllib/evaluation/MulticlassMetricsSuite.scala

@@ -17,24 +17,23 @@
 
 package org.apache.spark.mllib.evaluation
 
-import org.scalatest.FunSuite
-
 import org.apache.spark.mllib.util.LocalSparkContext
+import org.scalatest.FunSuite
 
 class MulticlassMetricsSuite extends FunSuite with LocalSparkContext {
   test("Multiclass evaluation metrics") {
     /*
-    * Confusion matrix for 3-class classification with total 9 instances:
-    * |2|1|1| true class0 (4 instances)
-    * |1|3|0| true class1 (4 instances)
-    * |0|0|1| true class2 (1 instance)
-    *
-    */
+     * Confusion matrix for 3-class classification with total 9 instances:
+     * |2|1|1| true class0 (4 instances)
+     * |1|3|0| true class1 (4 instances)
+     * |0|0|1| true class2 (1 instance)
+     */
+    val confusionMatrix = Array(Array(2, 1, 1), Array(1, 3, 0), Array(0, 0, 1))
     val labels = Array(0.0, 1.0, 2.0)
-    val scoreAndLabels = sc.parallelize(
+    val predictionAndLabels = sc.parallelize(
       Seq((0.0, 0.0), (0.0, 1.0), (0.0, 0.0), (1.0, 0.0), (1.0, 1.0),
         (1.0, 1.0), (1.0, 1.0), (2.0, 2.0), (2.0, 0.0)), 2)
-    val metrics = new MulticlassMetrics(scoreAndLabels)
+    val metrics = new MulticlassMetrics(predictionAndLabels)
     val delta = 0.0000001
     val fpRate0 = 1.0 / (9 - 4)
     val fpRate1 = 1.0 / (9 - 4)
@@ -48,6 +47,11 @@ class MulticlassMetricsSuite extends FunSuite with LocalSparkContext {
     val f1measure0 = 2 * precision0 * recall0 / (precision0 + recall0)
     val f1measure1 = 2 * precision1 * recall1 / (precision1 + recall1)
     val f1measure2 = 2 * precision2 * recall2 / (precision2 + recall2)
+    val f2measure0 = (1 + 2 * 2) * precision0 * recall0 / (2 * 2 * precision0 + recall0)
+    val f2measure1 = (1 + 2 * 2) * precision1 * recall1 / (2 * 2 * precision1 + recall1)
+    val f2measure2 = (1 + 2 * 2) * precision2 * recall2 / (2 * 2 * precision2 + recall2)
+
+    assert(metrics.confusionMatrix.deep == confusionMatrix.deep)
     assert(math.abs(metrics.falsePositiveRate(0.0) - fpRate0) < delta)
     assert(math.abs(metrics.falsePositiveRate(1.0) - fpRate1) < delta)
     assert(math.abs(metrics.falsePositiveRate(2.0) - fpRate2) < delta)
@@ -60,17 +64,25 @@ class MulticlassMetricsSuite extends FunSuite with LocalSparkContext {
     assert(math.abs(metrics.fMeasure(0.0) - f1measure0) < delta)
     assert(math.abs(metrics.fMeasure(1.0) - f1measure1) < delta)
     assert(math.abs(metrics.fMeasure(2.0) - f1measure2) < delta)
+    assert(math.abs(metrics.fMeasure(0.0, 2.0) - f2measure0) < delta)
+    assert(math.abs(metrics.fMeasure(1.0, 2.0) - f2measure1) < delta)
+    assert(math.abs(metrics.fMeasure(2.0, 2.0) - f2measure2) < delta)
+
     assert(math.abs(metrics.recall -
       (2.0 + 3.0 + 1.0) / ((2 + 3 + 1) + (1 + 1 + 1))) < delta)
     assert(math.abs(metrics.recall - metrics.precision) < delta)
     assert(math.abs(metrics.recall - metrics.fMeasure) < delta)
     assert(math.abs(metrics.recall - metrics.weightedRecall) < delta)
+    assert(math.abs(metrics.weightedFalsePositiveRate -
+      ((4.0 / 9) * fpRate0 + (4.0 / 9) * fpRate1 + (1.0 / 9) * fpRate2)) < delta)
     assert(math.abs(metrics.weightedPrecision -
       ((4.0 / 9) * precision0 + (4.0 / 9) * precision1 + (1.0 / 9) * precision2)) < delta)
     assert(math.abs(metrics.weightedRecall -
       ((4.0 / 9) * recall0 + (4.0 / 9) * recall1 + (1.0 / 9) * recall2)) < delta)
     assert(math.abs(metrics.weightedFMeasure -
       ((4.0 / 9) * f1measure0 + (4.0 / 9) * f1measure1 + (1.0 / 9) * f1measure2)) < delta)
+    assert(math.abs(metrics.weightedFMeasure(2.0) -
+      ((4.0 / 9) * f2measure0 + (4.0 / 9) * f2measure1 + (1.0 / 9) * f2measure2)) < delta)
     assert(metrics.labels.sameElements(labels))
   }
 }