dataframe-learn-2.0.0.0: tests-internal/Cart.hs

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

{- | Agreement tests: 'buildCartTree' must predict identically to sklearn's
@DecisionTreeClassifier(random_state=0, max_depth=4)@ on shared folds (golden
fixtures from @bench/export_cart_fixtures.py@); missing fixtures SKIP.
-}
module Cart (tests) where

import Control.Exception (SomeException, try)
import Data.Aeson (FromJSON (..), eitherDecode, withObject, (.:))
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as T
import qualified Data.Vector as V
import Test.HUnit

import DataFrame.DecisionTree (
    TreeConfig (..),
    defaultTreeConfig,
 )
import DataFrame.DecisionTree.Cart (buildCartTree)
import DataFrame.DecisionTree.Predict (predictManyWithTree)
import qualified DataFrame.Operations.Subset as DSub
import qualified DataFrameApi as D

data Fold = Fold ![Int] ![Int]
instance FromJSON Fold where
    parseJSON = withObject "fold" $ \o -> Fold <$> o .: "train" <*> o .: "test"

newtype Folds = Folds [Fold]
instance FromJSON Folds where
    parseJSON = withObject "folds" $ \o -> Folds <$> o .: "folds"

data Fixture = Fixture ![Int] ![T.Text]
instance FromJSON Fixture where
    parseJSON = withObject "fixture" $ \o -> Fixture <$> o .: "test_index" <*> o .: "test_pred"

-- sklearn cart_d4 params: max_depth 4, min_samples_leaf 1 (min_samples_split
-- is fixed at 2 inside buildCartTree).
cartCfg :: TreeConfig
cartCfg = defaultTreeConfig{maxTreeDepth = 4, minLeafSize = 1}

cartCases :: [(String, Int)]
cartCases =
    [("wine", i) | i <- [0 .. 4]] ++ [("bcw", i) | i <- [0 .. 4]] ++ [("adult", 0)]

tests :: [Test]
tests =
    [ TestLabel ("cart: " ++ n ++ " fold " ++ show i) (TestCase (runCase n i))
    | (n, i) <- cartCases
    ]

readJson :: (FromJSON a) => FilePath -> IO (Either String a)
readJson fp = do
    e <- try (BL.readFile fp) :: IO (Either SomeException BL.ByteString)
    pure $ case e of
        Left _ -> Left "missing"
        Right raw -> eitherDecode raw

-- wine is tie-free so sklearn is deterministic (exact match); bcw/adult have
-- equal-gain ties sklearn breaks via a seeded feature permutation, so we only
-- report the match fraction rather than chase its RNG.
runCase :: String -> Int -> IO ()
runCase name i = do
    efx <- readJson ("tests/fixtures/cart/" ++ name ++ "_fold" ++ show i ++ ".json")
    case efx of
        Left "missing" ->
            putStrLn
                ( "  [skip] cart "
                    ++ name
                    ++ " fold "
                    ++ show i
                    ++ ": fixture missing (run bench/export_cart_fixtures.py)"
                )
        Left e -> assertFailure ("fixture parse (" ++ name ++ "): " ++ e)
        Right (Fixture _ predExpected) -> do
            efolds <- readJson ("data/folds/" ++ name ++ ".json")
            case efolds of
                Left e -> assertFailure ("folds parse (" ++ name ++ "): " ++ e)
                Right (Folds fs) -> do
                    df <- D.readCsv ("data/uci/" ++ name ++ "_clean.csv")
                    let Fold trainIdx testIdx = fs !! i
                        trainDf = DSub.selectRows trainIdx df
                        tree = buildCartTree @Int cartCfg "target" trainDf
                        preds =
                            map
                                (T.pack . show)
                                (V.toList (predictManyWithTree tree df (V.fromList testIdx)))
                    if name == "wine"
                        then assertEqual ("cart " ++ name ++ " fold " ++ show i) predExpected preds
                        else do
                            let n = length predExpected
                                m = length (filter id (zipWith (==) predExpected preds))
                            putStrLn
                                ( "  [diagnostic] cart "
                                    ++ name
                                    ++ " fold "
                                    ++ show i
                                    ++ ": "
                                    ++ show m
                                    ++ "/"
                                    ++ show n
                                    ++ " predictions match sklearn(random_state=0) (remainder = sklearn's seeded equal-gain tie-break)"
                                )