MetricsTests.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module Learn.MetricsTests (tests) where

import qualified DataFrame as D
import qualified DataFrame.Functions as F
import qualified DataFrame.Internal.Column as DI

import DataFrame.LinearModel
import DataFrame.Metrics
import DataFrame.Metrics.Report
import DataFrame.ModelSelection
import DataFrame.PCA
import DataFrame.Transform

import qualified Data.Vector.Unboxed as VU
import DataFrame.Model (fit, predict)
import Test.HUnit

close :: Double -> Double -> Double -> Bool
close tol a b = abs (a - b) <= tol

preds3, truth3 :: VU.Vector Double
preds3 = VU.fromList [0, 0, 1, 1, 2, 2, 1, 0]
truth3 = VU.fromList [0, 0, 1, 2, 2, 2, 1, 0]

reg :: D.DataFrame
reg =
    D.fromNamedColumns
        [ ("x", DI.fromList ([1 .. 20] :: [Double]))
        ,
            ( "y"
            , DI.fromList ([2 * fromIntegral i + 1 | i <- [1 .. 20 :: Int]] :: [Double])
            )
        ]

testRegressionMetrics :: Test
testRegressionMetrics = TestCase $ do
    let p = VU.fromList [1, 2, 3, 4]
        t = VU.fromList [1, 2, 3, 5]
    assertBool "mse" (close 1e-9 (mse p t) 0.25)
    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)

testMulticlassMetrics :: Test
testMulticlassMetrics = TestCase $ do
    assertBool "accuracy" (close 1e-9 (accuracy preds3 truth3) 0.875)
    assertBool
        "binary precision class 1"
        (close 1e-9 (precision (Binary 1) preds3 truth3) (2 / 3))
    assertBool
        "macro f1 sane"
        (f1 Macro preds3 truth3 > 0.8 && f1 Macro preds3 truth3 <= 1)
    assertBool
        "micro f1 == accuracy"
        (close 1e-9 (f1 Micro preds3 truth3) (accuracy preds3 truth3))

testRocAuc :: Test
testRocAuc = TestCase $ do
    let scores = VU.fromList [0.1, 0.4, 0.35, 0.8]
        truth = VU.fromList [0, 0, 1, 1]
    assertBool "perfect-ish auc high" (rocAuc scores truth >= 0.75)
    assertBool "auc in [0,1]" (let a = rocAuc scores truth in a >= 0 && a <= 1)

testReports :: Test
testReports = TestCase $ do
    let cr = classificationReport preds3 truth3
    assertEqual "report covers 3 classes" 3 (length (crPerClass cr))
    assertBool "report accuracy" (close 1e-9 (crAccuracy cr) 0.875)
    let rr = regressionReport (VU.fromList [1, 2, 3]) (VU.fromList [1, 2, 4])
    assertBool "regression report rmse" (rrRMSE rr > 0)
    assertBool "classification report shows" (not (null (show cr)))
    assertBool "confusion shows" (not (null (show (confusionMatrix preds3 truth3))))

testEvaluateOneLiner :: Test
testEvaluateOneLiner = TestCase $ do
    let m = fit defaultLinearConfig (F.col @Double "y") reg
        score = evaluate rmse (predict m) (F.col @Double "y") reg
    assertBool "evaluate rmse ~ 0 on exact linear fit" (score < 1e-6)
    let r = regressionReportExpr (predict m) (F.col @Double "y") reg
    assertBool "report r2 ~ 1" (close 1e-6 (rrR2 r) 1)

testCrossValidate :: Test
testCrossValidate = TestCase $ do
    let cv =
            crossValidate
                4
                0
                r2
                (F.col @Double "y")
                (predict . fit defaultLinearConfig (F.col @Double "y"))
                reg
    assertBool "all folds high R2" (all (> 0.99) cv)
    assertEqual "four folds" 4 (length cv)

testTransformCompose :: Test
testTransformCompose = TestCase $ do
    let scaler = standardScaler ["x"] reg
        pca = fit (PCAConfig (NComp 1) False) [F.col @Double "x"] reg
        pipeline = scalerTransform scaler <> pcaTransform pca
        out = applyTransform pipeline reg
    assertBool "pipeline produced a frame" (D.columnNames out /= [])
    assertBool "pipeline has pc1 column" ("pc1" `elem` D.columnNames out)

tests :: [Test]
tests =
    [ testRegressionMetrics
    , testMulticlassMetrics
    , testRocAuc
    , testReports
    , testEvaluateOneLiner
    , testCrossValidate
    , testTransformCompose
    ]