Model.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

{- | sklearn-style standalone tree estimators returning inspectable records
(depth, leaf count, per-feature split usage). 'fit' trains a classifier (from a
'TreeConfig') or a regressor (from a 'RegTreeConfig'); 'predict' is the compiled
tree expression, and the record exposes the raw 'Tree' too. The bare
'DataFrame.DecisionTree.Fit.fitDecisionTree' remains for callers that only want
the classifier @Expr@.
-}
module DataFrame.DecisionTree.Model (
    module DataFrame.Model,
    DecisionTreeClassifier (..),
    DecisionTreeRegressor (..),
) where

import qualified Data.Map.Strict as M
import qualified Data.Text as T

import qualified Data.Vector as V

import DataFrame.DecisionTree.Cart (cartFeatures)
import DataFrame.DecisionTree.Fit (fitDecisionTree, treeToExpr)
import DataFrame.DecisionTree.Regression (RegTreeConfig, fitRegTreeOn)
import DataFrame.DecisionTree.Types (Tree (..), TreeConfig)
import DataFrame.Featurize.Internal (targetDoubles)
import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (Expr (..), getColumns)
import DataFrame.Model

-- | A fitted classification tree with structural diagnostics.
data DecisionTreeClassifier a = DecisionTreeClassifier
    { dtcExpr :: !(Expr a)
    , dtcDepth :: !Int
    , dtcNLeaves :: !Int
    , dtcFeatureUsage :: !(M.Map T.Text Int)
    }
    deriving (Show)

-- | A fitted regression tree with structural diagnostics.
data DecisionTreeRegressor = DecisionTreeRegressor
    { dtrTree :: !(Tree Double)
    , dtrExpr :: !(Expr Double)
    , dtrDepth :: !Int
    , dtrNLeaves :: !Int
    , dtrFeatureUsage :: !(M.Map T.Text Int)
    }
    deriving (Show)

instance (Columnable a, Ord a) => Fit TreeConfig (Expr a) where
    type ModelOf TreeConfig (Expr a) = (DecisionTreeClassifier a)
    fit cfg target df =
        DecisionTreeClassifier
            e
            (exprDepth e)
            (exprLeaves e)
            (usageCounts (exprUsage e))
      where
        e = fitDecisionTree cfg target df

instance Predict (DecisionTreeClassifier a) where
    type Prediction (DecisionTreeClassifier a) = Expr a
    predict = dtcExpr

instance Fit RegTreeConfig (Expr Double) where
    type ModelOf RegTreeConfig (Expr Double) = DecisionTreeRegressor
    fit cfg target df =
        DecisionTreeRegressor
            t
            e
            (exprDepth e)
            (exprLeaves e)
            (usageCounts (exprUsage e))
      where
        t = case target of
            Col name ->
                fitRegTreeOn
                    cfg
                    (V.fromList (cartFeatures name df))
                    (targetDoubles target df)
                    Nothing
            _ ->
                error
                    ("fit @DecisionTreeRegressor: target must be a column, got " ++ show target)
        e = treeToExpr t

instance Predict DecisionTreeRegressor where
    type Prediction DecisionTreeRegressor = Expr Double
    predict = dtrExpr

usageCounts :: [T.Text] -> M.Map T.Text Int
usageCounts = foldr (\c -> M.insertWith (+) c 1) M.empty

exprUsage :: Expr a -> [T.Text]
exprUsage (If c t e) = getColumns c ++ exprUsage t ++ exprUsage e
exprUsage _ = []

exprLeaves :: Expr a -> Int
exprLeaves (If _ t e) = exprLeaves t + exprLeaves e
exprLeaves _ = 1

exprDepth :: Expr a -> Int
exprDepth (If _ t e) = 1 + max (exprDepth t) (exprDepth e)
exprDepth _ = 0