Test to verify that custom metrics appear in model insight metrics

core/src/main/scala/com/salesforce/op/evaluators/EvaluationMetrics.scala

@@ -213,10 +213,12 @@ object OpEvaluatorNames extends Enum[OpEvaluatorNames] {
   case object BinScore extends OpEvaluatorNames("binScoreEval", "bin score evaluation metrics")
   case object Multi extends OpEvaluatorNames("multiEval", "multiclass evaluation metrics")
   case object Regression extends OpEvaluatorNames("regEval", "regression evaluation metrics")
-  case object Forecast extends OpEvaluatorNames("regForecast", "regression evaluation metrics")
+  case object Forecast extends OpEvaluatorNames("regForecast", "forecast evaluation metrics")
   case class Custom(name: String, humanName: String) extends OpEvaluatorNames(name, humanName) {
     override def entryName: String = name.toLowerCase
   }
+  def withFriendlyNameInsensitive(name: String): Option[OpEvaluatorNames] =
+    values.collectFirst { case n if n.humanFriendlyName.equalsIgnoreCase(name) => n }
   override def withName(name: String): OpEvaluatorNames = Try(super.withName(name)).getOrElse(Custom(name, name))
   override def withNameInsensitive(name: String): OpEvaluatorNames = super.withNameInsensitiveOption(name)
     .getOrElse(Custom(name, name))

core/src/main/scala/com/salesforce/op/stages/impl/selector/ModelSelectorSummary.scala

@@ -239,26 +239,26 @@ case object ModelSelectorSummary {
 
     ReflectionUtils.classForName(className) match {
       case n if n == classOf[MultiMetrics] =>
-        JsonUtils.fromString[Map[String, Map[String, Any]]](json).map{ d =>
-          val asMetrics = d.flatMap{ case (_, values) => values.map{
-            case (nm: String, mp: Map[String, Any]@unchecked) =>
+        JsonUtils.fromString[Map[String, Map[String, Any]]](json).map { d =>
+          val asMetrics = d.flatMap { case (_, values) =>
+            values.collect { case (nm: String, mp: Map[String, Any]@unchecked) =>
               val valsJson = JsonUtils.toJsonString(mp) // TODO: gross but it works. try to find a better way
-              nm match {
-                case OpEvaluatorNames.Binary.humanFriendlyName =>
-                  nm -> JsonUtils.fromString[BinaryClassificationMetrics](valsJson).get
-                case OpEvaluatorNames.BinScore.humanFriendlyName =>
-                  nm -> JsonUtils.fromString[BinaryClassificationBinMetrics](valsJson).get
-                case OpEvaluatorNames.Multi.humanFriendlyName =>
-                  nm -> JsonUtils.fromString[MultiClassificationMetrics](valsJson).get
-                case OpEvaluatorNames.Regression.humanFriendlyName =>
-                  nm -> JsonUtils.fromString[RegressionMetrics](valsJson).get
-                case _ => nm -> JsonUtils.fromString[SingleMetric](valsJson).get
-              }}
+              nm -> (OpEvaluatorNames.withFriendlyNameInsensitive(nm) match {
+                case Some(OpEvaluatorNames.Binary) => JsonUtils.fromString[BinaryClassificationMetrics](valsJson)
+                case Some(OpEvaluatorNames.BinScore) => JsonUtils.fromString[BinaryClassificationBinMetrics](valsJson)
+                case Some(OpEvaluatorNames.Multi) => JsonUtils.fromString[MultiClassificationMetrics](valsJson)
+                case Some(OpEvaluatorNames.Regression) => JsonUtils.fromString[RegressionMetrics](valsJson)
+                case Some(OpEvaluatorNames.Forecast) => JsonUtils.fromString[ForecastMetrics](valsJson)
+                case _ => // assume a custom metric here, hence trying to parse as single metric value
+                  JsonUtils.fromString[SingleMetric](valsJson)
+              }).get
+            }
           }
           MultiMetrics(asMetrics)
         }.recoverWith { case t: Throwable => error(n, t) }
-      case n => JsonUtils.fromString(json)(ClassTag(n))
-        .recoverWith { case t: Throwable => error(n, t) }
+      case n =>
+        JsonUtils.fromString(json)(ClassTag(n))
+          .recoverWith { case t: Throwable => error(n, t) }
     }
   }
 }

core/src/test/scala/com/salesforce/op/ModelInsightsTest.scala

@@ -30,6 +30,7 @@
 
 package com.salesforce.op
 
+import com.salesforce.op.evaluators._
 import com.salesforce.op.features.types._
 import com.salesforce.op.features.{Feature, FeatureDistributionType, FeatureLike}
 import com.salesforce.op.filters._
@@ -50,7 +51,7 @@ import org.junit.runner.RunWith
 import com.salesforce.op.features.types.Real
 import com.salesforce.op.stages.impl.feature.TextStats
 import com.twitter.algebird.Moments
-import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.{DataFrame, Dataset}
 import org.scalatest.FlatSpec
 import org.scalatest.junit.JUnitRunner
 import org.apache.spark.sql.functions._
@@ -798,4 +799,70 @@ class ModelInsightsTest extends FlatSpec with PassengerSparkFixtureTest with Dou
       }
     }
   }
+
+  it should "return default & custom metrics when having multiple binary classification metrics in model insights" in {
+    val prediction = BinaryClassificationModelSelector
+      .withCrossValidation(seed = 42,
+        trainTestEvaluators = Seq(
+          Evaluators.BinaryClassification.custom(metricName = "second", evaluateFn = _ => 0.0),
+          Evaluators.BinaryClassification.custom(metricName = "third", evaluateFn = _ => 1.0)
+        ),
+        splitter = Option(DataSplitter(seed = 42, reserveTestFraction = 0.1)),
+        modelsAndParameters = models)
+      .setInput(label, checked)
+      .getOutput()
+    val workflow = new OpWorkflow().setResultFeatures(prediction).setParameters(params).setReader(dataReader)
+    val workflowModel = workflow.train()
+    val insights = workflowModel.modelInsights(prediction)
+    val trainEval = insights.selectedModelInfo.get.trainEvaluation
+    trainEval shouldBe a[MultiMetrics]
+    val trainMetric = trainEval.asInstanceOf[MultiMetrics].metrics
+    trainMetric.map { case (metricName, metric) => metricName -> metric.getClass } should contain theSameElementsAs Seq(
+      OpEvaluatorNames.Binary.humanFriendlyName -> classOf[BinaryClassificationMetrics],
+      OpEvaluatorNames.BinScore.humanFriendlyName -> classOf[BinaryClassificationBinMetrics],
+      "second" -> classOf[SingleMetric],
+      "third" -> classOf[SingleMetric]
+    )
+  }
+
+  it should
+    "return default & custom metrics when having multiple multi-class classification metrics in model insights" in {
+    val prediction = MultiClassificationModelSelector
+      .withCrossValidation(seed = 42,
+        trainTestEvaluators = Seq(Evaluators.MultiClassification.custom(metricName = "second", evaluateFn = _ => 0.0)),
+        splitter = Option(DataCutter(seed = 42, reserveTestFraction = 0.1)),
+        modelsAndParameters = models)
+      .setInput(label, checked)
+      .getOutput()
+    val workflow = new OpWorkflow().setResultFeatures(prediction).setParameters(params).setReader(dataReader)
+    val workflowModel = workflow.train()
+    val insights = workflowModel.modelInsights(prediction)
+    val trainEval = insights.selectedModelInfo.get.trainEvaluation
+    trainEval shouldBe a[MultiMetrics]
+    val trainMetric = trainEval.asInstanceOf[MultiMetrics].metrics
+    trainMetric.map { case (metricName, metric) => metricName -> metric.getClass } should contain theSameElementsAs Seq(
+      OpEvaluatorNames.Multi.humanFriendlyName -> classOf[MultiClassificationMetrics],
+      "second" -> classOf[SingleMetric]
+    )
+  }
+
+  it should "return default & custom metrics when having multiple regression metrics in model insights" in {
+    val prediction = RegressionModelSelector
+      .withCrossValidation(seed = 42,
+        trainTestEvaluators = Seq(Evaluators.Regression.custom(metricName = "second", evaluateFn = _ => 0.0)),
+        dataSplitter = Option(DataSplitter(seed = 42, reserveTestFraction = 0.1)),
+        modelsAndParameters = models)
+      .setInput(label, features)
+      .getOutput()
+    val workflow = new OpWorkflow().setResultFeatures(prediction).setParameters(params).setReader(dataReader)
+    val workflowModel = workflow.train()
+    val insights = workflowModel.modelInsights(prediction)
+    val trainEval = insights.selectedModelInfo.get.trainEvaluation
+    trainEval shouldBe a[MultiMetrics]
+    val trainMetric = trainEval.asInstanceOf[MultiMetrics].metrics
+    trainMetric.map { case (metricName, metric) => metricName -> metric.getClass } should contain theSameElementsAs Seq(
+      OpEvaluatorNames.Regression.humanFriendlyName -> classOf[RegressionMetrics],
+      "second" -> classOf[SingleMetric]
+    )
+  }
 }