NMI implemented as geval metric

src/GEval/Core.hs

@@ -41,10 +41,13 @@ import GEval.BLEU
 import GEval.Common
 import GEval.ClippEU
 import GEval.PrecisionRecall
+import GEval.ClusteringMetrics
+
+import qualified Data.HashMap.Strict as M
 
 type MetricValue = Double
 
-data Metric = RMSE | MSE | BLEU | Accuracy | ClippEU | FMeasure Double
+data Metric = RMSE | MSE | BLEU | Accuracy | ClippEU | FMeasure Double | NMI
               deriving (Eq)
 
 instance Show Metric where
@@ -54,6 +57,7 @@ instance Show Metric where
   show Accuracy = "Accuracy"
   show ClippEU = "ClippEU"
   show (FMeasure beta) = "F" ++ (show beta)
+  show NMI = "NMI"
 
 instance Read Metric where
   readsPrec _ ('R':'M':'S':'E':theRest) = [(RMSE, theRest)]
@@ -61,6 +65,7 @@ instance Read Metric where
   readsPrec _ ('B':'L':'E':'U':theRest) = [(BLEU, theRest)]
   readsPrec _ ('A':'c':'c':'u':'r':'a':'c':'y':theRest) = [(Accuracy, theRest)]
   readsPrec _ ('C':'l':'i':'p':'p':'E':'U':theRest) = [(ClippEU, theRest)]
+  readsPrec _ ('N':'M':'I':theRest) = [(NMI, theRest)]
   readsPrec p ('F':theRest) = case readsPrec p theRest of
     [(beta, theRest)] -> [(FMeasure beta, theRest)]
     _ -> []
@@ -75,6 +80,7 @@ getMetricOrdering BLEU     = TheHigherTheBetter
 getMetricOrdering Accuracy = TheHigherTheBetter
 getMetricOrdering ClippEU  = TheHigherTheBetter
 getMetricOrdering (FMeasure _) = TheHigherTheBetter
+getMetricOrdering NMI = TheHigherTheBetter
 
 defaultOutDirectory = "."
 defaultTestName = "test-A"
@@ -221,6 +227,8 @@ gevalCore' ClippEU = gevalCore'' parseClippingSpecs parseClippings matchStep cli
     clippeuAgg = CC.foldl countFolder (0, 0, 0)
     finalStep counts = f2MeasureOnCounts counts
 
+gevalCore' NMI = gevalCore'' id id id (CC.foldl updateConfusionMatrix M.empty) normalizedMutualInformationFromConfusionMatrix
+
 data SourceItem a = Got a | Done
 
 gevalCore'' :: (Text -> a) -> (Text -> b) -> ((a, b) -> c) -> (Sink c (ResourceT IO) d) -> (d -> Double) -> String -> String -> IO (MetricValue)

src/GEval/OptionsParser.hs

@@ -74,7 +74,7 @@ metricReader = option auto
                  <> value defaultMetric
                  <> showDefault
                  <> metavar "METRIC"
-                 <> help "Metric to be used - RMSE, MSE, Accuracy, F-measure (specify as F1, F2, F0.25, etc.), BLEU or ClippEU" )
+                 <> help "Metric to be used - RMSE, MSE, Accuracy, F-measure (specify as F1, F2, F0.25, etc.), BLEU, NMI or ClippEU" )
 
 runGEval :: [String] -> IO (Either (ParserResult GEvalOptions) (Maybe MetricValue))
 runGEval args = do

test/Spec.hs

@@ -81,6 +81,9 @@ main = hspec $ do
       normalizedMutualInformation stupidClusteringOneBigCluster `shouldBeAlmost` 0.0
     it "stupid clustering with many small clusters" $ do
       normalizedMutualInformation stupidClusteringManySmallClusters `shouldBeAlmost` 0.61799
+  describe "NMI challenge" $ do
+    it "complex test" $ do
+      runGEvalTest "nmi-complex" `shouldReturnAlmost` 0.36456
   describe "reading options" $ do
     it "can get the metric" $ do
       extractMetric "bleu-complex" `shouldReturn` (Just BLEU)

test/nmi-complex/nmi-complex-solution/test-A/out.tsv

@@ -0,0 +1,17 @@
+3
+1
+1
+1
+1
+1
+1
+2
+2
+2
+2
+2
+2
+3
+3
+3
+3

test/nmi-complex/nmi-complex/config.txt

@@ -0,0 +1 @@
+--metric NMI

test/nmi-complex/nmi-complex/test-A/expected.tsv

@@ -0,0 +1,17 @@
+x
+x
+x
+o
+x
+x
+x
+x
+o
+o
+o
+o
+d
+x
+d
+d
+d