[MLLIB] SPARK-2329 Add multi-label evaluation metrics

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

@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.evaluation
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.SparkContext._
+
+/**
+ * Evaluator for multilabel classification.
+ * @param predictionAndLabels an RDD of (predictions, labels) pairs, both are non-null sets.
+ */
+class MultilabelMetrics(predictionAndLabels: RDD[(Set[Double], Set[Double])]) {
+
+  private lazy val numDocs: Long = predictionAndLabels.count
+
+  private lazy val numLabels: Long = predictionAndLabels.flatMap { case (_, labels) =>
+    labels}.distinct.count
+
+  /**
+   * Returns strict Accuracy
+   * (for equal sets of labels)
+   */
+  lazy val strictAccuracy: Double = predictionAndLabels.filter { case (predictions, labels) =>
+    predictions == labels}.count.toDouble / numDocs
+
+  /**
+   * Returns Accuracy
+   */
+  lazy val accuracy: Double = predictionAndLabels.map { case (predictions, labels) =>
+    labels.intersect(predictions).size.toDouble / labels.union(predictions).size}.sum / numDocs
+
+  /**
+   * Returns Hamming-loss
+   */
+  lazy val hammingLoss: Double = (predictionAndLabels.map { case (predictions, labels) =>
+    labels.diff(predictions).size + predictions.diff(labels).size}.
+    sum).toDouble / (numDocs * numLabels)
+
+  /**
+   * Returns Document-based Precision averaged by the number of documents
+   */
+  lazy val macroPrecisionDoc: Double = (predictionAndLabels.map { case (predictions, labels) =>
+    if (predictions.size > 0) {
+      predictions.intersect(labels).size.toDouble / predictions.size
+    } else 0
+  }.sum) / numDocs
+
+  /**
+   * Returns Document-based Recall averaged by the number of documents
+   */
+  lazy val macroRecallDoc: Double = (predictionAndLabels.map { case (predictions, labels) =>
+    labels.intersect(predictions).size.toDouble / labels.size}.sum) / numDocs
+
+  /**
+   * Returns Document-based F1-measure averaged by the number of documents
+   */
+  lazy val macroF1MeasureDoc: Double = (predictionAndLabels.map { case (predictions, labels) =>
+    2.0 * predictions.intersect(labels).size / (predictions.size + labels.size)}.sum) / numDocs
+
+  /**
+   * Returns micro-averaged document-based Precision
+   * (equals to label-based microPrecision)
+   */
+  lazy val microPrecisionDoc: Double = microPrecisionClass
+
+  /**
+   * Returns micro-averaged document-based Recall
+   * (equals to label-based microRecall)
+   */
+  lazy val microRecallDoc: Double = microRecallClass
+
+  /**
+   * Returns micro-averaged document-based F1-measure
+   * (equals to label-based microF1measure)
+   */
+  lazy val microF1MeasureDoc: Double = microF1MeasureClass
+
+  private lazy val tpPerClass = predictionAndLabels.flatMap { case (predictions, labels) =>
+    predictions.intersect(labels).map(category => (category, 1))}.reduceByKey(_ + _).collectAsMap()
+
+  private lazy val fpPerClass = predictionAndLabels.flatMap { case(predictions, labels) =>
+    predictions.diff(labels).map(category => (category, 1))}.reduceByKey(_ + _).collectAsMap()
+
+  private lazy val fnPerClass = predictionAndLabels.flatMap{ case(predictions, labels) =>
+    labels.diff(predictions).map(category => (category, 1))}.reduceByKey(_ + _).collectAsMap()
+
+  /**
+   * Returns Precision for a given label (category)
+   * @param label the label.
+   */
+  def precisionClass(label: Double) = {
+    val tp = tpPerClass(label)
+    val fp = fpPerClass.getOrElse(label, 0)
+    if (tp + fp == 0) 0 else tp.toDouble / (tp + fp)
+  }
+
+  /**
+   * Returns Recall for a given label (category)
+   * @param label the label.
+   */
+  def recallClass(label: Double) = {
+    val tp = tpPerClass(label)
+    val fn = fnPerClass.getOrElse(label, 0)
+    if (tp + fn == 0) 0 else tp.toDouble / (tp + fn)
+  }
+
+  /**
+   * Returns F1-measure for a given label (category)
+   * @param label the label.
+   */
+  def f1MeasureClass(label: Double) = {
+    val precision = precisionClass(label)
+    val recall = recallClass(label)
+    if((precision + recall) == 0) 0 else 2 * precision * recall / (precision + recall)
+  }
+
+  private lazy val sumTp = tpPerClass.foldLeft(0L){ case (sum, (_, tp)) => sum + tp}
+  private lazy val sumFpClass = fpPerClass.foldLeft(0L){ case (sum, (_, fp)) => sum + fp}
+  private lazy val sumFnClass = fnPerClass.foldLeft(0L){ case (sum, (_, fn)) => sum + fn}
+
+  /**
+   * Returns micro-averaged label-based Precision
+   */
+  lazy val microPrecisionClass = {
+    val sumFp = fpPerClass.foldLeft(0L){ case(sumFp, (_, fp)) => sumFp + fp}
+    sumTp.toDouble / (sumTp + sumFp)
+  }
+
+  /**
+   * Returns micro-averaged label-based Recall
+   */
+  lazy val microRecallClass = {
+    val sumFn = fnPerClass.foldLeft(0.0){ case(sumFn, (_, fn)) => sumFn + fn}
+    sumTp.toDouble / (sumTp + sumFn)
+  }
+
+  /**
+   * Returns micro-averaged label-based F1-measure
+   */
+  lazy val microF1MeasureClass = 2.0 * sumTp / (2 * sumTp + sumFnClass + sumFpClass)
+}

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

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.evaluation
+
+import org.scalatest.FunSuite
+
+import org.apache.spark.mllib.util.LocalSparkContext
+import org.apache.spark.rdd.RDD
+
+class MultilabelMetricsSuite extends FunSuite with LocalSparkContext {
+  test("Multilabel evaluation metrics") {
+    /*
+    * Documents true labels (5x class0, 3x class1, 4x class2):
+    * doc 0 - predict 0, 1 - class 0, 2
+    * doc 1 - predict 0, 2 - class 0, 1
+    * doc 2 - predict none - class 0
+    * doc 3 - predict 2 - class 2
+    * doc 4 - predict 2, 0 - class 2, 0
+    * doc 5 - predict 0, 1, 2 - class 0, 1
+    * doc 6 - predict 1 - class 1, 2
+    *
+    * predicted classes
+    * class 0 - doc 0, 1, 4, 5 (total 4)
+    * class 1 - doc 0, 5, 6 (total 3)
+    * class 2 - doc 1, 3, 4, 5 (total 4)
+    *
+    * true classes
+    * class 0 - doc 0, 1, 2, 4, 5 (total 5)
+    * class 1 - doc 1, 5, 6 (total 3)
+    * class 2 - doc 0, 3, 4, 6 (total 4)
+    *
+    */
+    val scoreAndLabels: RDD[(Set[Double], Set[Double])] = sc.parallelize(
+      Seq((Set(0.0, 1.0), Set(0.0, 2.0)),
+        (Set(0.0, 2.0), Set(0.0, 1.0)),
+        (Set(), Set(0.0)),
+        (Set(2.0), Set(2.0)),
+        (Set(2.0, 0.0), Set(2.0, 0.0)),
+        (Set(0.0, 1.0, 2.0), Set(0.0, 1.0)),
+        (Set(1.0), Set(1.0, 2.0))), 2)
+    val metrics = new MultilabelMetrics(scoreAndLabels)
+    val delta = 0.00001
+    val precision0 = 4.0 / (4 + 0)
+    val precision1 = 2.0 / (2 + 1)
+    val precision2 = 2.0 / (2 + 2)
+    val recall0 = 4.0 / (4 + 1)
+    val recall1 = 2.0 / (2 + 1)
+    val recall2 = 2.0 / (2 + 2)
+    val f1measure0 = 2 * precision0 * recall0 / (precision0 + recall0)
+    val f1measure1 = 2 * precision1 * recall1 / (precision1 + recall1)
+    val f1measure2 = 2 * precision2 * recall2 / (precision2 + recall2)
+    val sumTp = 4 + 2 + 2
+    assert(sumTp == (1 + 1 + 0 + 1 + 2 + 2 + 1))
+    val microPrecisionClass = sumTp.toDouble / (4 + 0 + 2 + 1 + 2 + 2)
+    val microRecallClass = sumTp.toDouble / (4 + 1 + 2 + 1 + 2 + 2)
+    val microF1MeasureClass = 2.0 * sumTp.toDouble /
+      (2 * sumTp.toDouble + (1 + 1 + 2) + (0 + 1 + 2))
+    val macroPrecisionDoc = 1.0 / 7 *
+      (1.0 / 2 + 1.0 / 2 + 0 + 1.0 / 1 + 2.0 / 2 + 2.0 / 3 + 1.0 / 1.0)
+    val macroRecallDoc = 1.0 / 7 *
+      (1.0 / 2 + 1.0 / 2 + 0 / 1 + 1.0 / 1 + 2.0 / 2 + 2.0 / 2 + 1.0 / 2)
+    val macroF1MeasureDoc = (1.0 / 7) *
+      2 * ( 1.0 / (2 + 2) + 1.0 / (2 + 2) + 0 + 1.0 / (1 + 1) +
+        2.0 / (2 + 2) + 2.0 / (3 + 2) + 1.0 / (1 + 2) )
+    val hammingLoss = (1.0 / (7 * 3)) * (2 + 2 + 1 + 0 + 0 + 1 + 1)
+    val strictAccuracy = 2.0 / 7
+    val accuracy = 1.0 / 7 * (1.0 / 3 + 1.0 /3 + 0 + 1.0 / 1 + 2.0 / 2 + 2.0 / 3 + 1.0 / 2)
+    assert(math.abs(metrics.precisionClass(0.0) - precision0) < delta)
+    assert(math.abs(metrics.precisionClass(1.0) - precision1) < delta)
+    assert(math.abs(metrics.precisionClass(2.0) - precision2) < delta)
+    assert(math.abs(metrics.recallClass(0.0) - recall0) < delta)
+    assert(math.abs(metrics.recallClass(1.0) - recall1) < delta)
+    assert(math.abs(metrics.recallClass(2.0) - recall2) < delta)
+    assert(math.abs(metrics.f1MeasureClass(0.0) - f1measure0) < delta)
+    assert(math.abs(metrics.f1MeasureClass(1.0) - f1measure1) < delta)
+    assert(math.abs(metrics.f1MeasureClass(2.0) - f1measure2) < delta)
+    assert(math.abs(metrics.microPrecisionClass - microPrecisionClass) < delta)
+    assert(math.abs(metrics.microRecallClass - microRecallClass) < delta)
+    assert(math.abs(metrics.microF1MeasureClass - microF1MeasureClass) < delta)
+    assert(math.abs(metrics.macroPrecisionDoc - macroPrecisionDoc) < delta)
+    assert(math.abs(metrics.macroRecallDoc - macroRecallDoc) < delta)
+    assert(math.abs(metrics.macroF1MeasureDoc - macroF1MeasureDoc) < delta)
+    assert(math.abs(metrics.hammingLoss - hammingLoss) < delta)
+    assert(math.abs(metrics.strictAccuracy - strictAccuracy) < delta)
+    assert(math.abs(metrics.accuracy - accuracy) < delta)
+  }
+}