Metrics.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | Evaluation metrics for fitted models.
module DataFrame.Metrics (
    -- * Metric type + evaluation
    Metric,
    evaluate,
    predictColumn,
    columnOf,

    -- * Regression metrics
    mse,
    rmse,
    mae,
    r2,

    -- * Classification metrics
    accuracy,
    logLoss,
    Average (..),
    precision,
    recall,
    f1,
    rocAuc,

    -- * Per-class helpers (for reports)
    classCounts,
    precOf,
    recOf,
    f1Of,
) where

import Control.Exception (throw)
import Data.Either (fromRight)
import Data.List (nub, sort)
import Data.Ord (comparing)
import qualified Data.Text as T
import qualified Data.Vector.Algorithms.Intro as VA
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)
import DataFrame.Internal.Interpreter (interpret)
import DataFrame.Operations.Transformations (derive)

-- | A metric maps predictions and ground truth to a scalar score.
type Metric = VU.Vector Double -> VU.Vector Double -> Double

-- | Evaluate a model's prediction expression against a truth column.
evaluate :: Metric -> Expr Double -> Expr Double -> DataFrame -> Double
evaluate metric predExpr truthExpr df =
    metric (columnOf df predExpr) (columnOf df truthExpr)

-- | Add a model's prediction expression to a frame as a named column.
predictColumn :: T.Text -> Expr Double -> DataFrame -> DataFrame
predictColumn = derive

-- | Interpret an expression to a @Double@ vector over a frame.
columnOf :: DataFrame -> Expr Double -> VU.Vector Double
columnOf df e = case interpret @Double df e of
    Right (TColumn c) -> fromRight VU.empty (toVector @Double @VU.Vector c)
    Left err -> throw err

-- | Number of compared pairs, the shorter of the two vectors' lengths.
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
    | n == 0 = throw (EmptyDataSetException "mse")
    | otherwise =
        VU.sum (VU.zipWith (\p t -> (p - t) ^ (2 :: Int)) preds truth) / n
  where
    n = nCompared preds truth

-- | Root mean squared error.
rmse :: Metric
rmse preds truth = sqrt (mse preds truth)

-- | Mean absolute error.
mae :: Metric
mae preds 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
    | n == 0 = throw (EmptyDataSetException "r2")
    | ssTot == 0 = 0
    | otherwise = 1 - ssRes / ssTot
  where
    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
    | n == 0 = throw (EmptyDataSetException "accuracy")
    | otherwise =
        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
    | n == 0 = throw (EmptyDataSetException "logLoss")
    | otherwise =
        negate
            ( VU.sum
                ( VU.zipWith
                    (\p y -> let q = clampP p in y * log q + (1 - y) * log (1 - q))
                    probs
                    truth
                )
            )
            / n
  where
    n = nCompared probs truth
    clampP p = max 1e-15 (min (1 - 1e-15) p)

-- | Averaging strategy for multiclass precision/recall/F1.
data Average
    = -- | one class is positive; the rest negative
      Binary Double
    | -- | unweighted mean over classes
      Macro
    | -- | pool per-class counts (equals accuracy for single-label)
      Micro
    | -- | support-weighted mean over classes
      Weighted
    deriving (Eq, Show)

-- | Per-class @(tp, fp, fn, support)@ over the class set of @truth ∪ preds@.
classCounts ::
    VU.Vector Double -> VU.Vector Double -> [(Double, (Int, Int, Int, Int))]
classCounts preds truth =
    [(c, countsFor c) | c <- classes]
  where
    classes = sort (nub (VU.toList truth ++ VU.toList preds))
    countsFor c =
        VU.foldl'
            ( \(tp, fp, fn, sup) (p, y) ->
                ( if p == c && y == c then tp + 1 else tp
                , if p == c && y /= c then fp + 1 else fp
                , if p /= c && y == c then fn + 1 else fn
                , if y == c then sup + 1 else sup
                )
            )
            (0, 0, 0, 0)
            (VU.zip preds truth)

safeDiv :: Int -> Int -> Double
safeDiv a b = if b == 0 then 0 else fromIntegral a / fromIntegral b

precOf, recOf :: (Int, Int, Int, Int) -> Double
precOf (tp, fp, _, _) = safeDiv tp (tp + fp)
recOf (tp, _, fn, _) = safeDiv tp (tp + fn)

f1Of :: (Int, Int, Int, Int) -> Double
f1Of cs =
    let p = precOf cs; r = recOf cs in if p + r == 0 then 0 else 2 * p * r / (p + r)

averaged ::
    ((Int, Int, Int, Int) -> Double) ->
    Average ->
    VU.Vector Double ->
    VU.Vector Double ->
    Double
averaged stat avg preds truth =
    case avg of
        Binary pos -> maybe 0 stat (lookup pos cc)
        Macro -> meanOf [stat c | (_, c) <- cc]
        Weighted ->
            let total = sum [sup | (_, (_, _, _, sup)) <- cc]
             in if total == 0
                    then 0
                    else
                        sum [fromIntegral sup * stat c | (_, c@(_, _, _, sup)) <- cc]
                            / fromIntegral total
        Micro ->
            let tp = sum [t | (_, (t, _, _, _)) <- cc]
                fp = sum [x | (_, (_, x, _, _)) <- cc]
                fn = sum [x | (_, (_, _, x, _)) <- cc]
             in stat (tp, fp, fn, 0)
  where
    cc = classCounts preds truth
    meanOf xs = if null xs then 0 else sum xs / fromIntegral (length xs)

-- | Precision with the given averaging.
precision :: Average -> VU.Vector Double -> VU.Vector Double -> Double
precision = averaged precOf

-- | Recall with the given averaging.
recall :: Average -> VU.Vector Double -> VU.Vector Double -> Double
recall = averaged recOf

-- | F1 with the given averaging.
f1 :: Average -> VU.Vector Double -> VU.Vector Double -> Double
f1 = averaged f1Of

-- | Binary ROC-AUC (Mann-Whitney), tied scores sharing their average rank.
rocAuc :: VU.Vector Double -> VU.Vector Double -> Double
rocAuc scores truth
    | nPos == 0 || nNeg == 0 = 0.5
    | otherwise = (rankSum - nPos * (nPos + 1) / 2) / (nPos * nNeg)
  where
    n = VU.length scores
    sorted = VU.modify (VA.sortBy (comparing fst)) (VU.zip scores truth)
    nPos = fromIntegral (VU.length (VU.filter (== 1) truth))
    nNeg = fromIntegral n - nPos
    rankSum = go 0 0
    go !start !acc
        | start >= n = acc
        | otherwise =
            let v = fst (VU.unsafeIndex sorted start)
                end = maybe n (+ start) (VU.findIndex ((/= v) . fst) (VU.drop start sorted))
                rank = fromIntegral (start + 1 + end) / 2
                positives = VU.length (VU.filter ((== 1) . snd) (VU.slice start (end - start) sorted))
             in go end (acc + rank * fromIntegral positives)