fix: metrics average over compared pairs and throw on empty input

dataframe-learn/src/DataFrame/Metrics.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
@@ -45,6 +46,7 @@ import Data.Ord (comparing)
 import qualified Data.Text as T
 import qualified Data.Vector.Unboxed as VU
 
+import DataFrame.Errors (DataFrameException (..))
 import DataFrame.Internal.Column (TypedColumn (..), toVector)
 import DataFrame.Internal.DataFrame (DataFrame)
 import DataFrame.Internal.Expression (Expr)
@@ -73,15 +75,20 @@ columnOf df e = case interpret @Double df e of
     Right (TColumn c) -> fromRight VU.empty (toVector @Double @VU.Vector c)
     Left err -> throw err
 
-n2 :: VU.Vector Double -> Double
-n2 = fromIntegral . VU.length
+{- | Compared pairs: 'VU.zipWith' truncates to the shorter vector, so every
+mean below divides by this, never by the length of 'truth' alone.
+-}
+nCompared :: VU.Vector Double -> VU.Vector Double -> Double
+nCompared preds truth = fromIntegral (min (VU.length preds) (VU.length truth))
 
 -- | Mean squared error.
 mse :: Metric
 mse preds truth
-    | VU.null truth = 0
+    | n == 0 = throw (EmptyDataSetException "mse")
     | otherwise =
-        VU.sum (VU.zipWith (\p t -> (p - t) ^ (2 :: Int)) preds truth) / n2 truth
+        VU.sum (VU.zipWith (\p t -> (p - t) ^ (2 :: Int)) preds truth) / n
+  where
+    n = nCompared preds truth
 
 -- | Root mean squared error.
 rmse :: Metric
@@ -90,30 +97,37 @@ rmse preds truth = sqrt (mse preds truth)
 -- | Mean absolute error.
 mae :: Metric
 mae preds truth
-    | VU.null truth = 0
-    | otherwise = VU.sum (VU.zipWith (\p t -> abs (p - t)) preds truth) / n2 truth
+    | n == 0 = throw (EmptyDataSetException "mae")
+    | otherwise = VU.sum (VU.zipWith (\p t -> abs (p - t)) preds truth) / n
+  where
+    n = nCompared preds truth
 
 -- | Coefficient of determination @R²@.
 r2 :: Metric
 r2 preds truth
-    | VU.null truth || ssTot == 0 = 0
+    | n == 0 = throw (EmptyDataSetException "r2")
+    | ssTot == 0 = 0
     | otherwise = 1 - ssRes / ssTot
   where
-    mean = VU.sum truth / n2 truth
-    ssRes = VU.sum (VU.zipWith (\p t -> (t - p) ^ (2 :: Int)) preds truth)
-    ssTot = VU.sum (VU.map (\t -> (t - mean) ^ (2 :: Int)) truth)
+    n = nCompared preds truth
+    truth' = VU.take (min (VU.length preds) (VU.length truth)) truth
+    mean = VU.sum truth' / n
+    ssRes = VU.sum (VU.zipWith (\p t -> (t - p) ^ (2 :: Int)) preds truth')
+    ssTot = VU.sum (VU.map (\t -> (t - mean) ^ (2 :: Int)) truth')
 
 -- | Fraction of exact matches.
 accuracy :: Metric
 accuracy preds truth
-    | VU.null truth = 0
+    | n == 0 = throw (EmptyDataSetException "accuracy")
     | otherwise =
-        fromIntegral (VU.length (VU.filter id (VU.zipWith (==) preds truth))) / n2 truth
+        fromIntegral (VU.length (VU.filter id (VU.zipWith (==) preds truth))) / n
+  where
+    n = nCompared preds truth
 
 -- | Binary log loss; probabilities clamped away from @0@/@1@.
 logLoss :: Metric
 logLoss probs truth
-    | VU.null truth = 0
+    | n == 0 = throw (EmptyDataSetException "logLoss")
     | otherwise =
         negate
             ( VU.sum
@@ -123,8 +137,9 @@ logLoss probs truth
                     truth
                 )
             )
-            / n2 truth
+            / n
   where
+    n = nCompared probs truth
     clampP p = max 1e-15 (min (1 - 1e-15) p)
 
 -- | Averaging strategy for multiclass precision/recall/F1.

dataframe-learn/tests-internal/Learn/EdgeCases.hs

@@ -28,7 +28,7 @@ import DataFrame.LinearModel
 import DataFrame.LinearSolver (sigmoid)
 import DataFrame.PCA
 
-import DataFrame.Internal.Statistics (correlation', variance')
+import DataFrame.Internal.Statistics (correlation', meanSquaredError, variance')
 
 import Test.HUnit
 
@@ -169,6 +169,27 @@ testCorrelationTooFew = TestCase $ do
         Nothing
         (correlation' (VU.fromList [1]) (VU.fromList [2]))
 
+{- meanSquaredError refuses length mismatches and empty inputs rather than
+   averaging over terms it never summed (or indexing out of bounds). -}
+testMeanSquaredErrorGuards :: Test
+testMeanSquaredErrorGuards = TestCase $ do
+    assertEqual
+        "mse of mismatched lengths is Nothing"
+        Nothing
+        (meanSquaredError (VU.fromList [0, 0, 0, 0]) (VU.fromList [2, 2]))
+    assertEqual
+        "mse with the longer prediction does not index out of bounds"
+        Nothing
+        (meanSquaredError (VU.fromList [1]) (VU.fromList [1, 2, 3]))
+    assertEqual
+        "mse of empty inputs is Nothing"
+        Nothing
+        (meanSquaredError VU.empty VU.empty)
+    assertEqual
+        "mse of equal-length inputs is the plain mean"
+        (Just 4.0)
+        (meanSquaredError (VU.fromList [0, 0]) (VU.fromList [2, 2]))
+
 -- ===========================================================================
 -- Category 8: stability inside the model expr layer
 -- ===========================================================================
@@ -425,6 +446,7 @@ tests =
     , testCorrelationPerfect
     , testCorrelationConstantColumnIsNaN
     , testCorrelationTooFew
+    , testMeanSquaredErrorGuards
     , testLogisticProbsExtremeFeatures
     , testOLSOneRow
     , testLogisticSingleClass

dataframe-operations/src-internal/DataFrame/Internal/Statistics.hs

@@ -202,11 +202,14 @@ interQuartileRange' samp =
 {-# INLINE interQuartileRange' #-}
 
 meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double
-meanSquaredError target prediction =
-    let
-        squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction
-     in
-        Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction))
+meanSquaredError target prediction
+    | VU.length target /= VU.length prediction = Nothing
+    | VU.null target = Nothing
+    | otherwise =
+        Just
+            ( VU.sum (VU.zipWith (\t p -> (p - t) ^ (2 :: Int)) target prediction)
+                / fromIntegral (VU.length target)
+            )
 {-# INLINE meanSquaredError #-}
 
 mutualInformationBinned ::

tests/Learn/MetricsTests.hs

@@ -1,8 +1,11 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
 module Learn.MetricsTests (tests) where
 
+import qualified Control.Exception as E
+
 import qualified DataFrame as D
 import qualified DataFrame.Functions as F
 import qualified DataFrame.Internal.Column as DI
@@ -15,7 +18,7 @@
 import DataFrame.Transform
 
 import qualified Data.Vector.Unboxed as VU
 import DataFrame.Model (fit, predict)
 import Test.HUnit
 
 close :: Double -> Double -> Double -> Bool
@@ -43,6 +46,16 @@
     assertBool "rmse" (close 1e-9 (rmse p t) 0.5)
     assertBool "mae" (close 1e-9 (mae p t) 0.25)
     assertBool "r2 in range" (r2 p t <= 1)
+    assertBool
+        "mse averages over compared pairs"
+        (close 1e-9 (mse (VU.fromList [2, 2]) (VU.fromList [0, 0, 0, 0])) 4)
+    assertBool
+        "mae averages over compared pairs"
+        (close 1e-9 (mae (VU.fromList [2, 2]) (VU.fromList [0, 0, 0, 0])) 2)
+    r <- E.try (E.evaluate (mse VU.empty (VU.fromList [5, 5, 5])))
+    case r of
+        Left (_ :: E.SomeException) -> pure ()
+        Right v -> assertFailure ("mse with no pairs returned " ++ show v)
 
 testMulticlassMetrics :: Test
 testMulticlassMetrics = TestCase $ do