diff --git a/dataframe-learn.cabal b/dataframe-learn.cabal --- a/dataframe-learn.cabal +++ b/dataframe-learn.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe-learn -version: 1.0.1.0 +version: 1.0.2.0 synopsis: Decision trees and feature synthesis for the dataframe ecosystem. description: @@ -13,7 +13,7 @@ license-file: LICENSE author: Michael Chavinda maintainer: mschavinda@gmail.com -copyright: (c) 2024-2025 Michael Chavinda +copyright: (c) 2024-2026 Michael Chavinda category: Data tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 @@ -29,12 +29,26 @@ import: warnings exposed-modules: DataFrame.DecisionTree + DataFrame.DecisionTree.Types + DataFrame.DecisionTree.CondVec + DataFrame.DecisionTree.Cart + DataFrame.DecisionTree.Numeric + DataFrame.DecisionTree.Prune + DataFrame.DecisionTree.Predict + DataFrame.DecisionTree.Categorical + DataFrame.DecisionTree.Pool + DataFrame.DecisionTree.Linear + DataFrame.DecisionTree.Tao + DataFrame.DecisionTree.Fit + DataFrame.LinearSolver DataFrame.Synthesis build-depends: base >= 4 && < 5, containers >= 0.6.7 && < 0.9, + parallel ^>= 3.2, dataframe-core ^>= 1.0, - dataframe-operations ^>= 1.0, + dataframe-operations ^>= 1.1, text >= 2.0 && < 3, - vector ^>= 0.13 + vector ^>= 0.13, + vector-algorithms ^>= 0.9 hs-source-dirs: src default-language: Haskell2010 diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -1,1007 +1,29 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -module DataFrame.DecisionTree where - -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame ( - DataFrame (..), - columnNames, - unsafeGetColumn, - ) -import DataFrame.Internal.Expression (Expr (..), eSize, eqExpr, getColumns) -import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Internal.Statistics (percentileOrd') -import DataFrame.Internal.Types -import DataFrame.Operations.Core (nRows) -import DataFrame.Operations.Subset (exclude, filterWhere) - -import Control.Exception (throw) -import Control.Monad (guard) -import Data.Function (on) -#if MIN_VERSION_base(4,20,0) -import Data.List (maximumBy, minimumBy, nub, nubBy, sort, sortBy) -#else -import Data.List (foldl', maximumBy, minimumBy, nub, nubBy, sort, sortBy) -#endif -import Data.Int (Int16, Int32, Int64, Int8) -import qualified Data.Map.Strict as M -import Data.Proxy (Proxy (..)) -import qualified Data.Text as T -import Data.Type.Equality -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import Data.Word (Word16, Word32, Word64, Word8) -import Type.Reflection (SomeTypeRep (..), typeRep) - -import DataFrame.Operators - -{- | Declares which column types support ordering for decision tree splits. - -Use 'orderable' to register a type, and '<>' to combine: - -@ -defaultTreeConfig - { columnOrdering = defaultColumnOrdering <> orderable \@MyCustomType - } -@ --} -newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict) - -instance Semigroup ColumnOrdering where - ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b) - -instance Monoid ColumnOrdering where - mempty = ColumnOrdering M.empty - --- | Register a type as orderable for decision tree splits. -orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering -orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a))) - --- | All standard numeric, text, and primitive types. -defaultColumnOrdering :: ColumnOrdering -defaultColumnOrdering = - mconcat - [ orderable @Int - , orderable @Int8 - , orderable @Int16 - , orderable @Int32 - , orderable @Int64 - , orderable @Word - , orderable @Word8 - , orderable @Word16 - , orderable @Word32 - , orderable @Word64 - , orderable @Integer - , orderable @Double - , orderable @Float - , orderable @Bool - , orderable @Char - , orderable @T.Text - , orderable @String - ] - --- Internal: existential Ord dictionary. -data OrdDict where - OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict - --- Internal: look up Ord for type @a@. -withOrdFrom :: - forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r -withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of - Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Just k - Nothing -> Nothing - Nothing -> Nothing - -data TreeConfig = TreeConfig - { maxTreeDepth :: Int - , minSamplesSplit :: Int - , minLeafSize :: Int - , percentiles :: [Int] - , expressionPairs :: Int - , synthConfig :: SynthConfig - , taoIterations :: Int - , taoConvergenceTol :: Double - , columnOrdering :: ColumnOrdering - } - -data SynthConfig = SynthConfig - { maxExprDepth :: Int - , boolExpansion :: Int - , disallowedCombinations :: [(T.Text, T.Text)] - , complexityPenalty :: Double - , enableStringOps :: Bool - , enableCrossCols :: Bool - , enableArithOps :: Bool - } - deriving (Eq, Show) - -defaultSynthConfig :: SynthConfig -defaultSynthConfig = - SynthConfig - { maxExprDepth = 2 - , boolExpansion = 2 - , disallowedCombinations = [] - , complexityPenalty = 0.05 - , enableStringOps = True - , enableCrossCols = True - , enableArithOps = True - } - -defaultTreeConfig :: TreeConfig -defaultTreeConfig = - TreeConfig - { maxTreeDepth = 4 - , minSamplesSplit = 5 - , minLeafSize = 1 - , percentiles = [0, 10 .. 100] - , expressionPairs = 10 - , synthConfig = defaultSynthConfig - , taoIterations = 10 - , taoConvergenceTol = 1e-6 - , columnOrdering = defaultColumnOrdering - } - -data Tree a - = Leaf !a - | Branch !(Expr Bool) !(Tree a) !(Tree a) - deriving (Show) - -treeDepth :: Tree a -> Int -treeDepth (Leaf _) = 0 -treeDepth (Branch _ l r) = 1 + max (treeDepth l) (treeDepth r) - -treeToExpr :: (Columnable a) => Tree a -> Expr a -treeToExpr (Leaf v) = Lit v -treeToExpr (Branch cond left right) = - F.ifThenElse cond (treeToExpr left) (treeToExpr right) - --- | Fit a TAO decision tree -fitDecisionTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - Expr a -> - DataFrame -> - Expr a -fitDecisionTree cfg (Col target) df = - let - conds = - nubBy eqExpr $ - numericConditions cfg (exclude [target] df) - ++ generateConditionsOld cfg (exclude [target] df) - - initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df - - indices = V.enumFromN 0 (nRows df) - - optimizedTree = taoOptimize @a cfg target conds df indices initialTree - in - pruneExpr (treeToExpr optimizedTree) -fitDecisionTree _ expr _ = error $ "Cannot create tree for compound expression: " ++ show expr - -taoOptimize :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> -- Target column name - [Expr Bool] -> -- Candidate conditions - DataFrame -> -- Full dataset - V.Vector Int -> -- Indices of points reaching the root - Tree a -> -- Current tree - Tree a -taoOptimize cfg target conds df rootIndices initialTree = - go 0 initialTree (computeTreeLoss @a target df rootIndices initialTree) - where - go :: Int -> Tree a -> Double -> Tree a - go iter tree prevLoss - | iter >= taoIterations cfg = pruneDead tree - | otherwise = - let - tree' = taoIteration @a cfg target conds df rootIndices tree - - newLoss = computeTreeLoss @a target df rootIndices tree' - improvement = prevLoss - newLoss - in - if improvement < taoConvergenceTol cfg - then pruneDead tree' - else go (iter + 1) tree' newLoss - -taoIteration :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -taoIteration cfg target conds df rootIndices tree = - let depth = treeDepth tree - in foldl' - (optimizeDepthLevel @a cfg target conds df rootIndices) - tree - [depth, depth - 1 .. 0] -- Bottom to top - -optimizeDepthLevel :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Int -> -- Target depth - Tree a -optimizeDepthLevel cfg target conds df rootIndices tree = optimizeAtDepth @a cfg target conds df rootIndices tree 0 - -optimizeAtDepth :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Int -> - Int -> - Tree a -optimizeAtDepth cfg target conds df indices tree currentDepth targetDepth - | currentDepth == targetDepth = - optimizeNode @a cfg target conds df indices tree - | otherwise = case tree of - Leaf v -> Leaf v - Branch cond left right -> - let - (indicesL, indicesR) = partitionIndices cond df indices - left' = - optimizeAtDepth @a - cfg - target - conds - df - indicesL - left - (currentDepth + 1) - targetDepth - right' = - optimizeAtDepth @a - cfg - target - conds - df - indicesR - right - (currentDepth + 1) - targetDepth - in - Branch cond left' right' - -optimizeNode :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -optimizeNode cfg target conds df indices tree - | V.null indices = tree - | otherwise = case tree of - Leaf _ -> Leaf (majorityValueFromIndices @a target df indices) - Branch oldCond left right -> - let - newCond = findBestSplitTAO @a cfg target conds df indices left right oldCond - - (newIndicesL, newIndicesR) = partitionIndices newCond df indices - in - if V.length newIndicesL < minLeafSize cfg - || V.length newIndicesR < minLeafSize cfg - then Leaf (majorityValueFromIndices @a target df indices) - else Branch newCond left right - -findBestSplitTAO :: - forall a. - (Columnable a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> -- Left subtree (FIXED) - Tree a -> -- Right subtree (FIXED) - Expr Bool -> -- Current condition (fallback) - Expr Bool -findBestSplitTAO cfg target conds df indices leftTree rightTree currentCond - | V.null indices = currentCond - | null validConds = currentCond - | otherwise = - let - carePoints = identifyCarePoints @a target df indices leftTree rightTree - in - if null carePoints - then currentCond - else - let - evalSplit :: Expr Bool -> Int - evalSplit cond = countCarePointErrors cond df carePoints - - evalWithPenalty c = - let errors = evalSplit c - penalty = - floor - ( complexityPenalty (synthConfig cfg) - * fromIntegral (eSize c) - ) - in errors + penalty - - sortedConds = - take (expressionPairs cfg) $ - sortBy (compare `on` evalWithPenalty) validConds - - expandedConds = - boolExprs - df - sortedConds - sortedConds - 0 - (boolExpansion (synthConfig cfg)) - in - if null expandedConds - then currentCond - else minimumBy (compare `on` evalWithPenalty) expandedConds - where - validConds = filter isValidSplit conds - isValidSplit c = - let (t, f) = partitionIndices c df indices - in V.length t >= minLeafSize cfg && V.length f >= minLeafSize cfg - --- | A care point with its index and which direction leads to correct classification -data CarePoint = CarePoint - { cpIndex :: !Int - , cpCorrectDir :: !Direction -- Which child classifies this point correctly - } - deriving (Eq, Show) - -data Direction = GoLeft | GoRight - deriving (Eq, Show) - -{- | Identify care points: points where exactly one subtree classifies correctly - - For each point reaching the node: - 1. Compute what label the left subtree would predict - 2. Compute what label the right subtree would predict - 3. If exactly one matches the true label, it's a care point - 4. Record which direction leads to correct classification --} -identifyCarePoints :: - forall a. - (Columnable a) => - T.Text -> - DataFrame -> - V.Vector Int -> - Tree a -> -- Left subtree - Tree a -> -- Right subtree - [CarePoint] -identifyCarePoints target df indices leftTree rightTree = - case interpret @a df (Col target) of - Left _ -> [] - Right (TColumn column) -> - case toVector @a column of - Left _ -> [] - Right targetVals -> - V.toList $ V.mapMaybe (checkPoint targetVals) indices - where - checkPoint :: V.Vector a -> Int -> Maybe CarePoint - checkPoint targetVals idx = - let - trueLabel = targetVals V.! idx - leftPred = predictWithTree @a target df idx leftTree - rightPred = predictWithTree @a target df idx rightTree - leftCorrect = leftPred == trueLabel - rightCorrect = rightPred == trueLabel - in - case (leftCorrect, rightCorrect) of - (True, False) -> Just $ CarePoint idx GoLeft - (False, True) -> Just $ CarePoint idx GoRight - _ -> Nothing -- Don't-care point (both correct or both wrong) - --- | Predict the label for a single point using a fixed tree -predictWithTree :: - forall a. - (Columnable a) => - T.Text -> - DataFrame -> - Int -> -- Row index - Tree a -> - a -predictWithTree _target _df _idx (Leaf v) = v -predictWithTree target df idx (Branch cond left right) = - case interpret @Bool df cond of - Left _ -> predictWithTree @a target df idx left -- Default to left on error - Right (TColumn column) -> - case toVector @Bool column of - Left _ -> predictWithTree @a target df idx left - Right boolVals -> - if boolVals V.! idx - then predictWithTree @a target df idx left - else predictWithTree @a target df idx right - -countCarePointErrors :: Expr Bool -> DataFrame -> [CarePoint] -> Int -countCarePointErrors cond df carePoints = - case interpret @Bool df cond of - Left _ -> length carePoints - Right (TColumn column) -> - case toVector @Bool column of - Left _ -> length carePoints - Right boolVals -> - length $ filter (isMisclassified boolVals) carePoints - where - isMisclassified :: V.Vector Bool -> CarePoint -> Bool - isMisclassified boolVals cp = - let goesLeft = boolVals V.! cpIndex cp - shouldGoLeft = cpCorrectDir cp == GoLeft - in goesLeft /= shouldGoLeft - -partitionIndices :: - Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int) -partitionIndices cond df indices = - case interpret @Bool df cond of - Left _ -> (indices, V.empty) - Right (TColumn column) -> - case toVector @Bool column of - Left _ -> (indices, V.empty) - Right boolVals -> - V.partition (boolVals V.!) indices - -majorityValueFromIndices :: - forall a. - (Columnable a, Ord a) => - T.Text -> - DataFrame -> - V.Vector Int -> - a -majorityValueFromIndices target df indices = - case interpret @a df (Col target) of - Left e -> throw e - Right (TColumn column) -> - case toVector @a column of - Left e -> throw e - Right vals -> - let counts = - V.foldl' - (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) - M.empty - indices - in if M.null counts - then error "Empty indices in majorityValueFromIndices" - else fst $ maximumBy (compare `on` snd) (M.toList counts) - -computeTreeLoss :: - forall a. - (Columnable a) => - T.Text -> - DataFrame -> - V.Vector Int -> - Tree a -> - Double -computeTreeLoss target df indices tree - | V.null indices = 0 - | otherwise = - case interpret @a df (Col target) of - Left _ -> 1.0 - Right (TColumn column) -> - case toVector @a column of - Left _ -> 1.0 - Right targetVals -> - let - n = V.length indices - errors = - V.length $ - V.filter - (\i -> targetVals V.! i /= predictWithTree @a target df i tree) - indices - in - fromIntegral errors / fromIntegral n - -pruneDead :: Tree a -> Tree a -pruneDead (Leaf v) = Leaf v -pruneDead (Branch cond left right) = - let - left' = pruneDead left - right' = pruneDead right - in - Branch cond left' right' - -pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a -pruneExpr (If cond trueBranch falseBranch) = - let t = pruneExpr trueBranch - f = pruneExpr falseBranch - in if eqExpr t f - then t - else case (t, f) of - (If condInner tInner _, _) | eqExpr cond condInner -> If cond tInner f - (_, If condInner _ fInner) | eqExpr cond condInner -> If cond t fInner - _ -> If cond t f -pruneExpr (Unary op e) = Unary op (pruneExpr e) -pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) -pruneExpr e = e - -buildGreedyTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - Int -> - T.Text -> - [Expr Bool] -> - DataFrame -> - Tree a -buildGreedyTree cfg depth target conds df - | depth <= 0 || nRows df <= minSamplesSplit cfg = - Leaf (majorityValue @a target df) - | otherwise = - case findBestGreedySplit @a cfg target conds df of - Nothing -> Leaf (majorityValue @a target df) - Just bestCond -> - let (dfTrue, dfFalse) = partitionDataFrame bestCond df - in if nRows dfTrue < minLeafSize cfg || nRows dfFalse < minLeafSize cfg - then Leaf (majorityValue @a target df) - else - Branch - bestCond - (buildGreedyTree @a cfg (depth - 1) target conds dfTrue) - (buildGreedyTree @a cfg (depth - 1) target conds dfFalse) - -findBestGreedySplit :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) -findBestGreedySplit cfg target conds df = - let - initialImpurity = calculateGini @a target df - calculateComplexity c = complexityPenalty (synthConfig cfg) * fromIntegral (eSize c) - - evalGain :: Expr Bool -> (Double, Int) - evalGain cond = - let (t, f) = partitionDataFrame cond df - n = fromIntegral @Int @Double (nRows df) - weightT = fromIntegral @Int @Double (nRows t) / n - weightF = fromIntegral @Int @Double (nRows f) / n - newImpurity = - weightT * calculateGini @a target t - + weightF * calculateGini @a target f - in ( (initialImpurity - newImpurity) - calculateComplexity cond - , negate (eSize cond) - ) - - validConds = - filter - ( \c -> - let (t, f) = partitionDataFrame c df - in nRows t >= minLeafSize cfg && nRows f >= minLeafSize cfg - ) - conds - - sortedConditions = - map fst $ - take - (expressionPairs cfg) - ( filter - (\(c, v) -> ((> negate (calculateComplexity c)) . fst) v) - (sortBy (flip compare `on` snd) (map (\c -> (c, evalGain c)) validConds)) - ) - in - if null sortedConditions - then Nothing - else - Just $ - maximumBy - (compare `on` evalGain) - ( boolExprs - df - sortedConditions - sortedConditions - 0 - (boolExpansion (synthConfig cfg)) - ) - --- | Unifies non-nullable and nullable Double expressions for feature generation. -data NumExpr - = NDouble !(Expr Double) - | NMaybeDouble !(Expr (Maybe Double)) - -numExprCols :: NumExpr -> [T.Text] -numExprCols (NDouble e) = getColumns e -numExprCols (NMaybeDouble e) = getColumns e - -numExprEq :: NumExpr -> NumExpr -> Bool -numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2 -numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2 -numExprEq _ _ = False - -combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] -combineNumExprs (NDouble e1) (NDouble e2) = - [ NDouble (e1 .+ e2) - , NDouble (e1 .- e2) - , NDouble (e1 .* e2) - , NDouble - (F.ifThenElse (e2 ./= F.lit (0 :: Double)) (e1 ./ e2) (F.lit (0 :: Double))) - ] -combineNumExprs (NDouble e1) (NMaybeDouble e2) = - [ NMaybeDouble (e1 .+ e2) - , NMaybeDouble (e1 .- e2) - , NMaybeDouble (e1 .* e2) - , NMaybeDouble - ( F.ifThenElse - (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) - (e1 ./ e2) - (F.lit (Nothing :: Maybe Double)) - ) - ] -combineNumExprs (NMaybeDouble e1) (NDouble e2) = - [ NMaybeDouble (e1 .+ e2) - , NMaybeDouble (e1 .- e2) - , NMaybeDouble (e1 .* e2) - , NMaybeDouble - ( F.ifThenElse - (e2 ./= F.lit (0 :: Double)) - (e1 ./ e2) - (F.lit (Nothing :: Maybe Double)) - ) - ] -combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) = - [ NMaybeDouble (e1 .+ e2) - , NMaybeDouble (e1 .- e2) - , NMaybeDouble (e1 .* e2) - , NMaybeDouble - ( F.ifThenElse - (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) - (e1 ./ e2) - (F.lit (Nothing :: Maybe Double)) - ) - ] - -numericConditions :: TreeConfig -> DataFrame -> [Expr Bool] -numericConditions = generateNumericConds - -generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool] -generateNumericConds cfg df = do - expr <- numericExprsWithTerms (synthConfig cfg) df - let thresholds = numericThresholds expr - threshold <- thresholds - numericCondsFromExpr expr threshold - where - numericThresholds (NDouble e) = map (\p -> percentile p e df) (percentiles cfg) - numericThresholds (NMaybeDouble e) = map (\p -> percentile p (F.fromMaybe 0 e) df) (percentiles cfg) - - numericCondsFromExpr (NDouble e) t = - [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] - numericCondsFromExpr (NMaybeDouble e) t = - [ F.fromMaybe False (e .<= F.lit t) - , F.fromMaybe False (e .>= F.lit t) - , F.fromMaybe False (e .< F.lit t) - , F.fromMaybe False (e .> F.lit t) - ] - -numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr] -numericExprsWithTerms cfg df = - concatMap (numericExprs cfg df [] 0) [0 .. maxExprDepth cfg] - -numericCols :: DataFrame -> [NumExpr] -numericCols df = concatMap extract (columnNames df) - where - extract colName = case unsafeGetColumn colName df of - UnboxedColumn Nothing (_ :: VU.Vector b) -> - case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NDouble (Col colName)] - Nothing -> case sIntegral @b of - STrue -> [NDouble (F.toDouble (Col @b colName))] - SFalse -> [] - BoxedColumn (Just _) (_ :: V.Vector b) -> - case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)] - Nothing -> case sIntegral @b of - STrue -> - [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))] - SFalse -> [] - UnboxedColumn (Just _) (_ :: VU.Vector b) -> - case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)] - Nothing -> case sIntegral @b of - STrue -> - [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))] - SFalse -> [] - _ -> [] - -numericExprs :: - SynthConfig -> DataFrame -> [NumExpr] -> Int -> Int -> [NumExpr] -numericExprs cfg df prevExprs depth maxDepth - | depth == 0 = baseExprs ++ numericExprs cfg df baseExprs (depth + 1) maxDepth - | depth >= maxDepth = [] - | otherwise = - combinedExprs ++ numericExprs cfg df combinedExprs (depth + 1) maxDepth - where - baseExprs = numericCols df - combinedExprs - | not (enableArithOps cfg) = [] - | otherwise = do - e1 <- prevExprs - e2 <- baseExprs - let cols = numExprCols e1 <> numExprCols e2 - guard - ( not (numExprEq e1 e2) - && not - ( any - (\(l, r) -> l `elem` cols && r `elem` cols) - (disallowedCombinations cfg) - ) - ) - combineNumExprs e1 e2 - -boolExprs :: - DataFrame -> [Expr Bool] -> [Expr Bool] -> Int -> Int -> [Expr Bool] -boolExprs df baseExprs prevExprs depth maxDepth - | depth == 0 = - baseExprs ++ boolExprs df baseExprs prevExprs (depth + 1) maxDepth - | depth >= maxDepth = [] - | otherwise = - combinedExprs ++ boolExprs df baseExprs combinedExprs (depth + 1) maxDepth - where - combinedExprs = do - e1 <- prevExprs - e2 <- baseExprs - guard (Prelude.not (eqExpr e1 e2)) - [F.and e1 e2, F.or e1 e2] - -generateConditionsOld :: TreeConfig -> DataFrame -> [Expr Bool] -generateConditionsOld cfg df = - let - ords = columnOrdering cfg - genConds :: T.Text -> [Expr Bool] - genConds colName = case unsafeGetColumn colName df of - (BoxedColumn Nothing (column :: V.Vector a)) -> - case withOrdFrom @a ords (map (Lit . (`percentileOrd'` column)) [1, 25, 75, 99]) of - Just ps -> map (\p -> Col @a colName .==. p) ps - Nothing -> [] - (BoxedColumn (Just _) (column :: V.Vector a)) -> case sFloating @a of - STrue -> [] -- handled by numericCols / numericExprs - SFalse -> case sIntegral @a of - STrue -> [] -- handled by numericCols / numericExprs - SFalse -> - case withOrdFrom @a - ords - (map (Lit . Just . (`percentileOrd'` column)) [1, 25, 75, 99]) of - Just ps -> map (\p -> Col @(Maybe a) colName .==. p) ps - Nothing -> [] - (UnboxedColumn _ (_ :: VU.Vector a)) -> [] - - columnConds = - concatMap - colConds - [ (l, r) - | l <- columnNames df - , r <- columnNames df - , not - ( any - (\(l', r') -> sort [l', r'] == sort [l, r]) - (disallowedCombinations (synthConfig cfg)) - ) - ] - where - colConds (!l, !r) = case (unsafeGetColumn l df, unsafeGetColumn r df) of - ( BoxedColumn Nothing (_col1 :: V.Vector a) - , BoxedColumn Nothing (_ :: V.Vector b) - ) -> - case testEquality (typeRep @a) (typeRep @b) of - Nothing -> [] - Just Refl -> [Col @a l .==. Col @a r] - (UnboxedColumn _ (_ :: VU.Vector a), UnboxedColumn _ (_ :: VU.Vector b)) -> [] - ( BoxedColumn (Just _) (_ :: V.Vector a) - , BoxedColumn (Just _) (_ :: V.Vector b) - ) -> case testEquality (typeRep @a) (typeRep @b) of - Nothing -> [] - Just Refl -> case testEquality (typeRep @a) (typeRep @T.Text) of - Nothing -> - case withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r] of - Just leExprs -> - leExprs ++ [Col @(Maybe a) l .==. Col @(Maybe a) r] - Nothing -> [Col @(Maybe a) l .==. Col @(Maybe a) r] - Just Refl -> [Col @(Maybe a) l .==. Col @(Maybe a) r] - _ -> [] - in - concatMap genConds (columnNames df) ++ columnConds - -partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) -partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) - -calculateGini :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double -calculateGini target df = - let n = fromIntegral $ nRows df - counts = getCounts @a target df - numClasses = fromIntegral $ M.size counts - probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts) - in if n == 0 then 0 else 1 - sum (map (^ (2 :: Int)) probs) - -majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a -majorityValue target df = - let counts = getCounts @a target df - in if M.null counts - then error "Empty DataFrame in leaf" - else fst $ maximumBy (compare `on` snd) (M.toList counts) - -getCounts :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int -getCounts target df = - case interpret @a df (Col target) of - Left e -> throw e - Right (TColumn column) -> - case toVector @a column of - Left e -> throw e - Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals) - -percentile :: Int -> Expr Double -> DataFrame -> Double -percentile p expr df = - case interpret @Double df expr of - Left _ -> 0 - Right (TColumn column) -> - case toVector @Double column of - Left _ -> 0 - Right vals -> - let sorted = V.fromList $ sort $ V.toList vals - n = V.length sorted - idx = min (n - 1) $ max 0 $ (p * n) `div` 100 - in if n == 0 then 0 else sorted V.! idx - -buildTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - Int -> - T.Text -> - [Expr Bool] -> - DataFrame -> - Expr a -buildTree cfg depth target conds df = - let - tree = buildGreedyTree @a cfg depth target conds df - indices = V.enumFromN 0 (nRows df) - optimized = taoOptimize @a cfg target conds df indices tree - in - pruneExpr (treeToExpr optimized) - -findBestSplit :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) -findBestSplit = findBestGreedySplit @a - -pruneTree :: forall a. (Columnable a) => Expr a -> Expr a -pruneTree = pruneExpr - --- | A tree where each leaf stores a class-probability distribution. -type ProbTree a = Tree (M.Map a Double) - --- | Compute normalised class probabilities from a subset of training rows. -probsFromIndices :: - forall a. - (Columnable a, Ord a) => - T.Text -> - DataFrame -> - V.Vector Int -> - M.Map a Double -probsFromIndices target df indices = - case interpret @a df (Col target) of - Left _ -> M.empty - Right (TColumn column) -> - case toVector @a column of - Left _ -> M.empty - Right vals -> - let counts = - V.foldl' - (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) - M.empty - indices - total = fromIntegral (V.length indices) :: Double - in M.map (\c -> fromIntegral c / total) counts - -{- | Annotate a fitted 'Tree a' with class distributions by routing the - training data through it. The split conditions are preserved; only the - leaf values change from a majority label to a probability map. --} -buildProbTree :: - forall a. - (Columnable a, Ord a) => - Tree a -> - T.Text -> - DataFrame -> - V.Vector Int -> - ProbTree a -buildProbTree (Leaf _) target df indices = - Leaf (probsFromIndices @a target df indices) -buildProbTree (Branch cond left right) target df indices = - let (indicesL, indicesR) = partitionIndices cond df indices - in Branch - cond - (buildProbTree @a left target df indicesL) - (buildProbTree @a right target df indicesR) - -{- | Fit a TAO decision tree and return one @Expr Double@ per class. - - Each @(c, e)@ pair in the result map means: evaluate @e@ on a 'DataFrame' - row to get the predicted probability of class @c@. You can insert these - as new columns with 'derive' or evaluate them with 'interpret'. - - Example: - @ - let pes = fitProbTree \@T.Text cfg (Col \"species\") trainDf - -- pes M.! \"setosa\" :: Expr Double - df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes - @ --} -fitProbTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - Expr a -> -- target column, e.g. @Col \"label\"@ - DataFrame -> - M.Map a (Expr Double) -fitProbTree cfg (Col target) df = - let - conds = - nubBy eqExpr $ - numericConditions cfg (exclude [target] df) - ++ generateConditionsOld cfg (exclude [target] df) - initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df - indices = V.enumFromN 0 (nRows df) - optimizedTree = taoOptimize @a cfg target conds df indices initialTree - pruned = pruneDead optimizedTree - in - probExprs (buildProbTree @a pruned target df indices) -fitProbTree _ expr _ = - error $ "Cannot create prob tree for compound expression: " ++ show expr - -{- | Convert a 'ProbTree' into one 'Expr Double' per class. - - Each @(c, e)@ pair means: evaluate @e@ on a 'DataFrame' row to get the - predicted probability of class @c@. You can insert these as new columns - with 'derive' or evaluate them with 'interpret'. - - Example: - @ - let pt = fitProbTree \@T.Text cfg (Col \"species\") trainDf - pes = probExprs pt - -- pes M.! \"setosa\" :: Expr Double - df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes - @ --} -probExprs :: - forall a. - (Columnable a, Ord a) => - ProbTree a -> - M.Map a (Expr Double) -probExprs tree = - let classes = nub (allClasses tree) - in M.fromList [(c, classExpr c tree) | c <- classes] - where - allClasses :: ProbTree a -> [a] - allClasses (Leaf m) = M.keys m - allClasses (Branch _ l r) = allClasses l ++ allClasses r - - classExpr :: a -> ProbTree a -> Expr Double - classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m) - classExpr c (Branch cond l r) = - F.ifThenElse cond (classExpr c l) (classExpr c r) +{- | Decision-tree training on DataFrames: a faithful CART tree refined by Tree +Alternating Optimization (TAO). This module re-exports the implementation, +which is split across the @DataFrame.DecisionTree.*@ modules. +-} +module DataFrame.DecisionTree ( + module DataFrame.DecisionTree.Types, + module DataFrame.DecisionTree.CondVec, + module DataFrame.DecisionTree.Cart, + module DataFrame.DecisionTree.Numeric, + module DataFrame.DecisionTree.Categorical, + module DataFrame.DecisionTree.Pool, + module DataFrame.DecisionTree.Predict, + module DataFrame.DecisionTree.Linear, + module DataFrame.DecisionTree.Tao, + module DataFrame.DecisionTree.Prune, + module DataFrame.DecisionTree.Fit, +) where + +import DataFrame.DecisionTree.Cart +import DataFrame.DecisionTree.Categorical +import DataFrame.DecisionTree.CondVec +import DataFrame.DecisionTree.Fit +import DataFrame.DecisionTree.Linear +import DataFrame.DecisionTree.Numeric +import DataFrame.DecisionTree.Pool +import DataFrame.DecisionTree.Predict +import DataFrame.DecisionTree.Prune +import DataFrame.DecisionTree.Tao +import DataFrame.DecisionTree.Types diff --git a/src/DataFrame/DecisionTree/Cart.hs b/src/DataFrame/DecisionTree/Cart.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Cart.hs @@ -0,0 +1,230 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | sklearn-faithful CART initializer used to seed TAO. One-hot encodes +-- categoricals, splits on exact (unsmoothed) Gini over midpoint thresholds +-- (@<=@ routes left), and emits a @Tree@ predicting identically to +-- @DecisionTreeClassifier(criterion='gini')@ on continuous features. +module DataFrame.DecisionTree.Cart ( + CartFeature (..), + CartNode (..), + sortIndicesByValue, + buildCartTree, + cartFeatures, + cartTargetLabels, +) where + +import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operations.Core (nRows) +import DataFrame.Operators + +import Data.Either (fromRight) +import Data.Function (on) +import Data.List (foldl') +import qualified Data.Map.Strict as M +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Merge as VA +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | A one-hot feature column: per-row Double values plus the sklearn LEFT +-- predicate (@x <= threshold@) over the ORIGINAL DataFrame. +data CartFeature = CartFeature + { cfValues :: !(VU.Vector Double) + , cfPred :: !(Double -> Expr Bool) + } + +-- | Pre-'Tree' CART node: a leaf class id, or a split on feature @j@. +data CartNode = CLeaf !Int | CSplit !Int !Double !CartNode !CartNode + +-- | Immutable per-fit context for the CART recursion. +data CartCtx = CartCtx + { ctxFeats :: !(V.Vector CartFeature) + , ctxNFeats :: !Int + , ctxCodes :: !(VU.Vector Int) + , ctxNClasses :: !Int + , ctxMaxDepth :: !Int + , ctxMinLeaf :: !Int + } + +-- | Indices @0..n-1@ stably sorted by their value (ascending), ties keeping +-- ascending index. In-place unboxed merge sort — no boxed-list allocation. +sortIndicesByValue :: VU.Vector Double -> VU.Vector Int +sortIndicesByValue vs = + VU.create $ do + mv <- VU.thaw (VU.enumFromN 0 (VU.length vs)) + VA.sortBy (compare `on` (vs VU.!)) mv + pure mv + +buildCartTree :: forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> Tree a +buildCartTree cfg target df = + cartToTree feats classes (buildCartNode ctx 0 (VU.enumFromN 0 nAll) featSorted) + where + nAll = nRows df + feats = V.fromList (cartFeatures target df) + featSorted = V.map (sortIndicesByValue . cfValues) feats + labels = cartLabels @a df target + classes = cartClasses labels + ctx = + CartCtx + feats + (V.length feats) + (classCodes classes labels) + (V.length classes) + (maxTreeDepth cfg) + (max 1 (minLeafSize cfg)) + +cartLabels :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a +cartLabels df target = case interpret @a df (Col target) of + Right (TColumn column) -> fromRight err (toVector @a column) + _ -> err + where + err = error "buildCartTree: cannot interpret target column" + +cartClasses :: (Ord a) => V.Vector a -> V.Vector a +cartClasses = V.fromList . Set.toList . Set.fromList . V.toList + +classCodes :: (Ord a) => V.Vector a -> V.Vector a -> VU.Vector Int +classCodes classes labels = VU.generate (V.length labels) (\i -> M.findWithDefault 0 (labels V.! i) ix) + where + ix = M.fromList (zip (V.toList classes) [0 ..]) + +cartToTree :: V.Vector CartFeature -> V.Vector a -> CartNode -> Tree a +cartToTree feats classes = go + where + go (CLeaf cid) = Leaf (classes V.! cid) + go (CSplit fj thr l r) = Branch (cfPred (feats V.! fj) thr) (go l) (go r) + +classCounts :: CartCtx -> VU.Vector Int -> VU.Vector Int +classCounts ctx idxs = + VU.accumulate (+) (VU.replicate (ctxNClasses ctx) 0) (VU.map (\i -> (ctxCodes ctx VU.! i, 1)) idxs) + +isPure :: VU.Vector Int -> Bool +isPure counts = VU.length (VU.filter (> 0) counts) <= 1 + +buildCartNode :: CartCtx -> Int -> VU.Vector Int -> V.Vector (VU.Vector Int) -> CartNode +buildCartNode ctx depth idxs sortedByFeat + | VU.length idxs < 2 || depth >= ctxMaxDepth ctx || isPure counts = leaf + | otherwise = maybe leaf (splitNode ctx depth idxs sortedByFeat) (bestSplit ctx sortedByFeat counts n) + where + n = VU.length idxs + counts = classCounts ctx idxs + leaf = CLeaf (VU.maxIndex counts) + +splitNode :: CartCtx -> Int -> VU.Vector Int -> V.Vector (VU.Vector Int) -> (Int, Double) -> CartNode +splitNode ctx depth idxs sortedByFeat (fj, thr) = + CSplit fj thr (rec leftIdx leftSorted) (rec rightIdx rightSorted) + where + vals = cfValues (ctxFeats ctx V.! fj) + leftIdx = VU.filter (\i -> vals VU.! i <= thr) idxs + rightIdx = VU.filter (\i -> vals VU.! i > thr) idxs + leftSorted = V.map (VU.filter (\i -> vals VU.! i <= thr)) sortedByFeat + rightSorted = V.map (VU.filter (\i -> vals VU.! i > thr)) sortedByFeat + rec = buildCartNode ctx (depth + 1) + +-- | Minimum weighted-child-Gini @(feature, threshold)@; the first feature wins +-- ties; 'Nothing' when no feature has a leaf-size-respecting threshold. +bestSplit :: CartCtx -> V.Vector (VU.Vector Int) -> VU.Vector Int -> Int -> Maybe (Int, Double) +bestSplit ctx sortedByFeat counts n = + fmap (\(_, j, t) -> (j, t)) (foldl' consider Nothing [0 .. ctxNFeats ctx - 1]) + where + total = VU.toList counts + consider acc fj = case sweepFeature ctx total (sortedByFeat V.! fj) (ctxFeats ctx V.! fj) n of + Just (g, thr) | maybe True (\(gB, _, _) -> g < gB) acc -> Just (g, fj, thr) + _ -> acc + +-- | Accumulator while sweeping a feature's sorted rows: best @(gini, thr)@ so +-- far, per-class left counts, rows moved left, and the previous value seen. +data Sweep = Sweep + { swBest :: !(Maybe (Double, Double)) + , swLeft :: ![Int] + , swMoved :: !Int + , swPrev :: !Double + } + +sweepFeature :: CartCtx -> [Int] -> VU.Vector Int -> CartFeature -> Int -> Maybe (Double, Double) +sweepFeature ctx total si feat n = + swBest (foldl' step (Sweep Nothing (replicate (ctxNClasses ctx) 0) 0 (0 / 0)) [0 .. VU.length si - 1]) + where + vals = cfValues feat + step s k = advance ctx total n (vals VU.! i) (ctxCodes ctx VU.! i) s + where + i = si VU.! k + +advance :: CartCtx -> [Int] -> Int -> Double -> Int -> Sweep -> Sweep +advance ctx total n v c s = + Sweep (considerThreshold ctx total n v s) (bumpClass c (swLeft s)) (swMoved s + 1) v + +considerThreshold :: CartCtx -> [Int] -> Int -> Double -> Sweep -> Maybe (Double, Double) +considerThreshold ctx total n v s + | swMoved s >= ctxMinLeaf ctx + , n - swMoved s >= ctxMinLeaf ctx + , v > swPrev s + 1e-7 = + keepBetter (swBest s) (weightedGini total (swLeft s) (swMoved s) n) ((swPrev s + v) / 2) + | otherwise = swBest s + +keepBetter :: Maybe (Double, Double) -> Double -> Double -> Maybe (Double, Double) +keepBetter best g thr = case best of + Just (wb, _) | wb <= g -> best + _ -> Just (g, thr) + +weightedGini :: [Int] -> [Int] -> Int -> Int -> Double +weightedGini total leftAcc nl n = + (fromIntegral nl * giniImpurity leftAcc nl + fromIntegral nr * giniImpurity rightAcc nr) + / fromIntegral n + where + nr = n - nl + rightAcc = zipWith (-) total leftAcc + +-- | Gini impurity @1 - Σ (c/m)²@ of a class-count list of total @m@. +giniImpurity :: [Int] -> Int -> Double +giniImpurity _ 0 = 0 +giniImpurity cs m = 1 - sum [let p = fromIntegral c / fromIntegral m in p * p | c <- cs] + +bumpClass :: Int -> [Int] -> [Int] +bumpClass c = zipWith (\j x -> if j == c then x + 1 else x) [0 ..] + +-- | One-hot features in @pd.get_dummies(drop_first=False)@ column order. +cartFeatures :: T.Text -> DataFrame -> [CartFeature] +cartFeatures target df = concatMap (featuresOfColumn df) (filter (/= target) (columnNames df)) + +featuresOfColumn :: DataFrame -> T.Text -> [CartFeature] +featuresOfColumn df c = case unsafeGetColumn c df of + UnboxedColumn _ (v :: VU.Vector b) -> numericFeature @b c v + BoxedColumn _ (v :: V.Vector b) -> oneHotFeatures @b (nRows df) c v + +numericFeature :: forall b. (Columnable b, VU.Unbox b) => T.Text -> VU.Vector b -> [CartFeature] +numericFeature c v = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [CartFeature v (\t -> F.col @Double c .<=. F.lit t)] + Nothing -> case sIntegral @b of + STrue -> [CartFeature (VU.map fromIntegral v) (\t -> F.toDouble (F.col @b c) .<=. F.lit t)] + SFalse -> [] + +oneHotFeatures :: forall b. (Columnable b) => Int -> T.Text -> V.Vector b -> [CartFeature] +oneHotFeatures nAll c v = case testEquality (typeRep @b) (typeRep @T.Text) of + Just Refl -> [oneHot nAll c v cat | cat <- Set.toList (Set.fromList (V.toList v))] + Nothing -> [] + +oneHot :: Int -> T.Text -> V.Vector T.Text -> T.Text -> CartFeature +oneHot nAll c v cat = + CartFeature (VU.generate nAll (\i -> if v V.! i == cat then 1 else 0)) (const (F.col @T.Text c ./=. F.lit cat)) + +-- | Target column as string labels (matches pandas @y.astype(str)@). +cartTargetLabels :: T.Text -> DataFrame -> V.Vector T.Text +cartTargetLabels target df = case unsafeGetColumn target df of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @T.Text) of + Just Refl -> v + Nothing -> V.map (T.pack . show) v + UnboxedColumn _ (v :: VU.Vector b) -> V.map (T.pack . show) (V.convert v) diff --git a/src/DataFrame/DecisionTree/Categorical.hs b/src/DataFrame/DecisionTree/Categorical.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Categorical.hs @@ -0,0 +1,281 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Categorical split candidates: Breiman prefixes for binary targets, +-- subset/singleton enumeration otherwise, and cross-column equality. Each +-- value-list yields an OR-of-equalities condition (as an expression or a +-- directly-read membership truth vector). +module DataFrame.DecisionTree.Categorical ( + TargetInfo (..), + mkTargetInfo, + distinctValuesUpTo, + validBoxedValues, + orEqs, + subsetSplits, + subsetLists, + singletonSplits, + singletonLists, + breimanPrefixSplits, + breimanPrefixLists, + catValueLists, + membershipVec, + crossColumnConds, + discreteConditions, + discreteCondVecs, +) where + +import DataFrame.DecisionTree.CondVec (CondVec (..), materializeCondVec) +import DataFrame.DecisionTree.Types (ColumnOrdering, SynthConfig (..), TreeConfig (..), withOrdFrom) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operators + +import Data.Either (fromRight) +import Data.Function (on) +import Data.List (inits, sort, sortBy, subsequences) +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe, mapMaybe) +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | Valid-slot view of a nullable boxed column (null slots hold crash-thunks). +validBoxedValues :: Bitmap -> V.Vector a -> V.Vector a +validBoxedValues bm = V.ifilter (\i _ -> bitmapTestBit bm i) + +-- | Target-column summary driving the categorical generator: binary vs +-- multi-class, the deterministic positive class, and the raw label vector. +data TargetInfo target = TargetInfo + { tiIsBinary :: !Bool + , tiPositiveClass :: !(Maybe target) + , tiValues :: !(V.Vector target) + } + +-- | Compute 'TargetInfo' once per fit. The positive class for binary targets +-- is the lexicographically-first distinct value, for deterministic pools. +mkTargetInfo :: forall target. (Columnable target, Ord target) => T.Text -> DataFrame -> Maybe (TargetInfo target) +mkTargetInfo target df = case interpret @target df (Col target) of + Right (TColumn column) -> either (const Nothing) (Just . targetInfoFromValues) (toVector @target column) + _ -> Nothing + +targetInfoFromValues :: (Ord target) => V.Vector target -> TargetInfo target +targetInfoFromValues vals = TargetInfo isBinary posClass vals + where + distinct = Set.toAscList (Set.fromList (V.toList vals)) + isBinary = length distinct == 2 + posClass = case distinct of + (p : _) | isBinary -> Just p + _ -> Nothing + +-- | Distinct values, capped: @Right vs@ (sorted) under the cap, else @Left@ +-- the count-so-far so the caller routes to the high-cardinality path. +distinctValuesUpTo :: (Ord a) => Int -> V.Vector a -> Either Int [a] +distinctValuesUpTo cap values = go Set.empty 0 + where + n = V.length values + go !s !i + | i >= n = Right (Set.toAscList s) + | Set.size s > cap = Left (Set.size s) + | otherwise = go (Set.insert (V.unsafeIndex values i) s) (i + 1) + +-- | OR-of-equalities for a value-list, shared by the expression and +-- truth-vector discrete paths so they stay byte-identical. +orEqs :: (a -> Expr Bool) -> [a] -> Expr Bool +orEqs eqLit = foldr1 (.||.) . map eqLit + +subsetSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] +subsetSplits eqLit = map (orEqs eqLit) . subsetLists + +-- | Proper non-empty, non-full subsets of the values. +subsetLists :: [a] -> [[a]] +subsetLists vs = drop 1 (init (subsequences vs)) + +singletonSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] +singletonSplits = map + +singletonLists :: [a] -> [[a]] +singletonLists = map (: []) + +breimanPrefixSplits :: (Ord a, Ord target) => target -> V.Vector a -> V.Vector target -> [a] -> (a -> Expr Bool) -> [Expr Bool] +breimanPrefixSplits pc values targetVals distinctVals eqLit = + map (orEqs eqLit) (breimanPrefixLists pc values targetVals distinctVals) + +-- | Breiman's binary-target split set: sort levels by Laplace-smoothed +-- positive rate, then take every contiguous non-trivial prefix. +breimanPrefixLists :: (Ord a, Ord target) => target -> V.Vector a -> V.Vector target -> [a] -> [[a]] +breimanPrefixLists pc values targetVals distinctVals = + nonTrivialPrefixes (sortByRate (levelCounts pc values targetVals) distinctVals) + +levelCounts :: (Ord a, Eq target) => target -> V.Vector a -> V.Vector target -> M.Map a (Int, Int) +levelCounts pc values targetVals = V.ifoldl' add M.empty values + where + add acc i v = M.insertWith plus v (indicator (V.unsafeIndex targetVals i == pc), 1) acc + plus (p1, n1) (p2, n2) = (p1 + p2, n1 + n2) + indicator b = if b then 1 else 0 + +laplaceRate :: (Ord a) => M.Map a (Int, Int) -> a -> Double +laplaceRate counts v = case M.lookup v counts of + Nothing -> 0.5 + Just (pos, n) -> (fromIntegral pos + 1) / (fromIntegral n + 2) + +sortByRate :: (Ord a) => M.Map a (Int, Int) -> [a] -> [a] +sortByRate counts = sortBy (compare `on` (\v -> (laplaceRate counts v, v))) + +nonTrivialPrefixes :: [a] -> [[a]] +nonTrivialPrefixes = tail . init . inits + +-- | Value-lists a categorical column contributes; shared by the expression and +-- truth-vector paths so both enumerate identical candidates in the same order. +catValueLists :: (Ord a, Ord target) => Bool -> Maybe target -> V.Vector target -> Int -> V.Vector a -> [[a]] +catValueLists isBinary posClass targetVals subsetCap values + | V.null values = [] + | isBinary, Just pc <- posClass = binaryLists pc targetVals values + | otherwise = multiclassLists subsetCap values + +binaryLists :: (Ord a, Ord target) => target -> V.Vector target -> V.Vector a -> [[a]] +binaryLists pc targetVals values + | length distinct < 2 = [] + | otherwise = breimanPrefixLists pc values targetVals distinct + where + distinct = fromRight (ascDistinct values) (distinctValuesUpTo 64 values) + +multiclassLists :: (Ord a) => Int -> V.Vector a -> [[a]] +multiclassLists subsetCap values = case distinctValuesUpTo subsetCap values of + Right vs | length vs >= 2 -> subsetLists vs + Right _ -> [] + Left _ -> singletonLists (ascDistinct values) + +ascDistinct :: (Ord a) => V.Vector a -> [a] +ascDistinct = Set.toAscList . Set.fromList . V.toList + +-- | Truth vector of @col ∈ values@ read directly from the column; equal to +-- interpreting @orEqs (== v) values@ because the values are distinct. +membershipVec :: (Ord a) => V.Vector a -> [a] -> VU.Vector Bool +membershipVec colVals vs = + let !s = Set.fromList vs + in VU.generate (V.length colVals) (\i -> Set.member (colVals `V.unsafeIndex` i) s) + +-- | Per-fit categorical generation context bundling the target summary and +-- the column-ordering registry. +data CatCtx target = CatCtx + { ccBinary :: !Bool + , ccPos :: !(Maybe target) + , ccTargets :: !(V.Vector target) + , ccSubsetCap :: !Int + , ccOrds :: !ColumnOrdering + } + +catCtx :: TargetInfo target -> TreeConfig -> CatCtx target +catCtx ti cfg = + CatCtx (tiIsBinary ti) (tiPositiveClass ti) (tiValues ti) (maxCategoricalSubsetCardinality (synthConfig cfg)) (columnOrdering cfg) + +catValueListsFor :: (Ord a, Ord target) => CatCtx target -> V.Vector a -> [[a]] +catValueListsFor ctx = catValueLists (ccBinary ctx) (ccPos ctx) (ccTargets ctx) (ccSubsetCap ctx) + +-- | True for numeric columns (handled by the numeric pool, not here). +isNumericKind :: forall a. (Columnable a) => Bool +isNumericKind = case sFloating @a of + STrue -> True + SFalse -> case sIntegral @a of + STrue -> True + SFalse -> False + +-- | All equality-based candidate splits from non-numeric columns: per-column +-- categorical conditions plus cross-column equality/order conditions. +discreteConditions :: forall target. (Columnable target, Ord target) => TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool] +discreteConditions targetInfo cfg df = + concatMap (columnConds (catCtx targetInfo cfg) df) (columnNames df) ++ crossColumnConds cfg df + +columnConds :: (Columnable target, Ord target) => CatCtx target -> DataFrame -> T.Text -> [Expr Bool] +columnConds ctx df colName = case unsafeGetColumn colName df of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColConds ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> nullableColConds ctx colName bm column + UnboxedColumn _ (_ :: VU.Vector a) -> [] + +nonNullColConds :: forall a target. (Columnable a, Ord target) => CatCtx target -> T.Text -> V.Vector a -> [Expr Bool] +nonNullColConds ctx colName column = + fromMaybe [] (withOrdFrom @a (ccOrds ctx) (map (orEqs (eqExprFor @a colName)) (catValueListsFor ctx column))) + +nullableColConds :: forall a target. (Columnable a, Ord target) => CatCtx target -> T.Text -> Bitmap -> V.Vector a -> [Expr Bool] +nullableColConds ctx colName bm column + | isNumericKind @a || V.null valid = [] + | otherwise = fromMaybe [] (withOrdFrom @a (ccOrds ctx) (map (orEqs (eqJustFor @a colName)) (catValueListsFor ctx valid))) + where + valid = validBoxedValues bm column + +eqExprFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool +eqExprFor colName v = Col @a colName .==. Lit v + +eqJustFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool +eqJustFor colName v = Col @(Maybe a) colName .==. Lit (Just v) + +-- | Cross-column equality/order conditions over pairs of same-typed columns. +crossColumnConds :: TreeConfig -> DataFrame -> [Expr Bool] +crossColumnConds cfg df = concatMap (pairConds (columnOrdering cfg) df) (allowedPairs cfg df) + +allowedPairs :: TreeConfig -> DataFrame -> [(T.Text, T.Text)] +allowedPairs cfg df = + [(l, r) | l <- columnNames df, r <- columnNames df, l /= r, not (isDisallowedPair cfg l r)] + +isDisallowedPair :: TreeConfig -> T.Text -> T.Text -> Bool +isDisallowedPair cfg l r = + any (\(l', r') -> sort [l', r'] == sort [l, r]) (disallowedCombinations (synthConfig cfg)) + +pairConds :: ColumnOrdering -> DataFrame -> (T.Text, T.Text) -> [Expr Bool] +pairConds ords df (l, r) = case (unsafeGetColumn l df, unsafeGetColumn r df) of + (BoxedColumn Nothing (_ :: V.Vector a), BoxedColumn Nothing (_ :: V.Vector b)) -> strictPairConds @a @b l r + (BoxedColumn (Just _) (_ :: V.Vector a), BoxedColumn (Just _) (_ :: V.Vector b)) -> nullablePairConds @a @b ords l r + _ -> [] + +strictPairConds :: forall a b. (Columnable a, Columnable b) => T.Text -> T.Text -> [Expr Bool] +strictPairConds l r = case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> [Col @a l .==. Col @a r] + Nothing -> [] + +nullablePairConds :: forall a b. (Columnable a, Columnable b) => ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] +nullablePairConds ords l r = case testEquality (typeRep @a) (typeRep @b) of + Nothing -> [] + Just Refl -> nullableEqOrLe @a ords l r + +nullableEqOrLe :: forall a. (Columnable a) => ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] +nullableEqOrLe ords l r + | isTextType @a = eqOnly + | otherwise = maybe eqOnly (++ eqOnly) (withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r]) + where + eqOnly = [Col @(Maybe a) l .==. Col @(Maybe a) r] + +isTextType :: forall a. (Columnable a) => Bool +isTextType = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> True + Nothing -> False + +-- | 'discreteConditions' materialized with shared per-column reads: the +-- non-nullable categorical path builds truth vectors directly from one read +-- per column; nullable and cross-column fall back to interpret. +discreteCondVecs :: forall target. (Columnable target, Ord target) => TargetInfo target -> TreeConfig -> DataFrame -> [CondVec] +discreteCondVecs targetInfo cfg df = + concatMap (columnCondVecs (catCtx targetInfo cfg) df) (columnNames df) + ++ mapMaybe (materializeCondVec df) (crossColumnConds cfg df) + +columnCondVecs :: (Columnable target, Ord target) => CatCtx target -> DataFrame -> T.Text -> [CondVec] +columnCondVecs ctx df colName = case unsafeGetColumn colName df of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColCondVecs ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> mapMaybe (materializeCondVec df) (nullableColConds ctx colName bm column) + UnboxedColumn _ (_ :: VU.Vector a) -> [] + +nonNullColCondVecs :: forall a target. (Columnable a, Ord target) => CatCtx target -> T.Text -> V.Vector a -> [CondVec] +nonNullColCondVecs ctx colName column = + fromMaybe [] (withOrdFrom @a (ccOrds ctx) (map (membershipCondVec colName column) (catValueListsFor ctx column))) + +membershipCondVec :: forall a. (Columnable a, Ord a) => T.Text -> V.Vector a -> [a] -> CondVec +membershipCondVec colName column vs = CondVec (orEqs (eqExprFor @a colName) vs) (membershipVec column vs) diff --git a/src/DataFrame/DecisionTree/CondVec.hs b/src/DataFrame/DecisionTree/CondVec.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/CondVec.hs @@ -0,0 +1,155 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Cached condition truth vectors and the per-fit cache keyed by structural +-- form. A condition's truth over a fixed DataFrame is invariant for a whole +-- fit, so it is materialized once and reused. +module DataFrame.DecisionTree.CondVec ( + CondVec (..), + materializeCondVec, + CondCache, + condCacheKey, + condCacheFromVecs, + addTreeCondsToCache, + lookupCondVec, + partitionByVec, + countErrorsByVec, + consolidateThreshold, + combineAndVec, + combineOrVec, +) where + +import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..), Tree (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (BinaryOp (binaryName), Expr (..), eqExpr, normalize) +import DataFrame.Internal.Interpreter (interpret) + +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | A boolean condition paired with its truth vector over the full DataFrame. +data CondVec = CondVec + { cvExpr :: !(Expr Bool) + , cvVec :: !(VU.Vector Bool) + } + +-- | Interpret a condition once over the DataFrame; 'Nothing' on a +-- type/interpret failure so the candidate is silently dropped. +materializeCondVec :: DataFrame -> Expr Bool -> Maybe CondVec +materializeCondVec df cond = case interpret @Bool df cond of + Left _ -> Nothing + Right (TColumn column) -> CondVec cond <$> eitherToMaybe (toVector @Bool @VU.Vector column) + +eitherToMaybe :: Either e a -> Maybe a +eitherToMaybe = either (const Nothing) Just + +-- | Full-DataFrame truth vectors keyed by structural form, read-only once +-- built. Seeded for free from the candidate pool plus the initial tree so the +-- predict/loss passes index a vector instead of re-interpreting per node. +type CondCache = M.Map T.Text (VU.Vector Bool) + +-- | Structural key matching the candidate-dedup key, so a tree branch whose +-- condition came from the pool hits the cache (equal keys ⟹ equal vector). +condCacheKey :: Expr Bool -> T.Text +condCacheKey = T.pack . show . normalize + +-- | Seed a cache from already-materialized candidate 'CondVec's (no interpret). +condCacheFromVecs :: [CondVec] -> CondCache +condCacheFromVecs cvs = M.fromList [(condCacheKey (cvExpr cv), cvVec cv) | cv <- cvs] + +-- | Add a tree's branch-condition vectors to a cache (one interpret per +-- distinct, not-yet-cached condition). +addTreeCondsToCache :: DataFrame -> Tree a -> CondCache -> CondCache +addTreeCondsToCache df = go + where + go (Leaf _) c = c + go (Branch cond l r) c = go r (go l (insertCond df cond c)) + +insertCond :: DataFrame -> Expr Bool -> CondCache -> CondCache +insertCond df cond c + | M.member k c = c + | otherwise = maybe c (\cv -> M.insert k (cvVec cv) c) (materializeCondVec df cond) + where + k = condCacheKey cond + +-- | A condition's truth vector: a cache hit, else interpret over the +-- DataFrame. 'Nothing' mirrors the interpret-failure fallback (route left). +lookupCondVec :: CondCache -> DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) +lookupCondVec cache df cond = case M.lookup (condCacheKey cond) cache of + hit@(Just _) -> hit + Nothing -> cvVec <$> materializeCondVec df cond + +-- | Partition row indices by a truth vector: @True@ → left, @False@ → right. +partitionByVec :: VU.Vector Bool -> V.Vector Int -> (V.Vector Int, V.Vector Int) +partitionByVec boolVals = V.partition (boolVals VU.!) + +-- | Count care points the truth vector routes to the wrong child. +countErrorsByVec :: VU.Vector Bool -> [CarePoint] -> Int +countErrorsByVec boolVals = length . filter misrouted + where + misrouted cp = (boolVals VU.! cpIndex cp) /= (cpCorrectDir cp == GoLeft) + +-- | A same-column same-direction Double threshold comparison, with a rebuild +-- function to swap in a new threshold. +data ThreshCmp = ThreshCmp + { tcCol :: !T.Text + , tcName :: !T.Text + , tcThr :: !Double + , tcRebuild :: Double -> Expr Bool + } + +asDoubleThreshold :: Expr Bool -> Maybe ThreshCmp +asDoubleThreshold (Binary op (Col c :: Expr cc) (Lit (t :: tt))) = + case (testEquality (typeRep @cc) (typeRep @Double), testEquality (typeRep @tt) (typeRep @Double)) of + (Just Refl, Just Refl) -> Just (ThreshCmp c (binaryName op) t (Binary op (Col c) . Lit)) + _ -> Nothing +asDoubleThreshold _ = Nothing + +directionalNames :: [T.Text] +directionalNames = ["lt", "leq", "gt", "geq"] + +-- | Tighter (AND) or looser (OR) of two same-direction thresholds: @<@/@<=@ +-- are left-half-spaces (AND = min), @>@/@>=@ are right-half-spaces (AND = max). +chooseThreshold :: Bool -> T.Text -> Double -> Double -> Double +chooseThreshold isAnd name t1 t2 + | leftDir = if isAnd then min t1 t2 else max t1 t2 + | otherwise = if isAnd then max t1 t2 else min t1 t2 + where + leftDir = name == "lt" || name == "leq" + +-- | Collapse two same-column same-direction strict-Double comparisons into one +-- comparison (the @True@ argument selects AND, @False@ OR); 'Nothing' otherwise. +consolidateThreshold :: Bool -> Expr Bool -> Expr Bool -> Maybe (Expr Bool) +consolidateThreshold isAnd ea eb = do + a <- asDoubleThreshold ea + b <- asDoubleThreshold eb + if tcCol a == tcCol b && tcName a == tcName b && tcName a `elem` directionalNames + then Just (tcRebuild a (chooseThreshold isAnd (tcName a) (tcThr a) (tcThr b))) + else Nothing + +-- | AND-combine two cached conditions: idempotence and threshold consolidation +-- first, else the generic @F.and@; the vector is always the elementwise AND. +combineAndVec :: CondVec -> CondVec -> CondVec +combineAndVec a b + | eqExpr (cvExpr a) (cvExpr b) = a + | otherwise = CondVec expr (VU.zipWith (&&) (cvVec a) (cvVec b)) + where + expr = fromMaybe (F.and (cvExpr a) (cvExpr b)) (consolidateThreshold True (cvExpr a) (cvExpr b)) + +-- | OR-combine two cached conditions (see 'combineAndVec'; AND/OR direction +-- differs in 'consolidateThreshold'). +combineOrVec :: CondVec -> CondVec -> CondVec +combineOrVec a b + | eqExpr (cvExpr a) (cvExpr b) = a + | otherwise = CondVec expr (VU.zipWith (||) (cvVec a) (cvVec b)) + where + expr = fromMaybe (F.or (cvExpr a) (cvExpr b)) (consolidateThreshold False (cvExpr a) (cvExpr b)) diff --git a/src/DataFrame/DecisionTree/Fit.hs b/src/DataFrame/DecisionTree/Fit.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Fit.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Top-level fitting: assemble the candidate pool, seed from CART, run TAO, +-- and convert the result to an expression. Also the probability-tree variant +-- ('fitProbTree') that annotates leaves with class distributions. +module DataFrame.DecisionTree.Fit ( + treeToExpr, + fitDecisionTree, + buildTree, + pruneTree, + partitionDataFrame, + calculateGini, + majorityValue, + getCounts, + percentile, + ProbTree, + probsFromIndices, + buildProbTree, + fitProbTree, + probExprs, +) where + +import DataFrame.DecisionTree.Cart (buildCartTree) +import DataFrame.DecisionTree.Categorical (TargetInfo (..), discreteConditions, discreteCondVecs, mkTargetInfo) +import DataFrame.DecisionTree.CondVec (CondVec) +import DataFrame.DecisionTree.Numeric (numericCondVecs, numericConditions) +import DataFrame.DecisionTree.Pool (dedupCVByExpr, nubByExpr) +import DataFrame.DecisionTree.Predict (partitionIndices) +import DataFrame.DecisionTree.Prune (pruneDead, pruneExpr) +import DataFrame.DecisionTree.Tao (taoOptimize, taoOptimizeCV) +import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Operations.Core (nRows) +import DataFrame.Operations.Subset (exclude, filterWhere) + +import Control.Exception (throw) +import Data.Function (on) +import Data.List (foldl', maximumBy, nub, sort) +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V + +-- | Convert a fitted tree to a nested-conditional expression. +treeToExpr :: (Columnable a) => Tree a -> Expr a +treeToExpr (Leaf v) = Lit v +treeToExpr (Branch cond left right) = F.ifThenElse cond (treeToExpr left) (treeToExpr right) + +-- | Fit a TAO decision tree (CART-seeded) and return it as an expression. +fitDecisionTree :: forall a. (Columnable a, Ord a) => TreeConfig -> Expr a -> DataFrame -> Expr a +fitDecisionTree cfg (Col target) df = + pruneExpr (treeToExpr (taoOptimizeCV @a cfg target condVecs df indices initialTree)) + where + condVecs = candidatePool @a cfg target df + initialTree = buildCartTree @a cfg target df + indices = V.enumFromN 0 (nRows df) +fitDecisionTree _ expr _ = error ("Cannot create tree for compound expression: " ++ show expr) + +-- | The deduplicated numeric + discrete candidate pool for a target column. +candidatePool :: forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> [CondVec] +candidatePool cfg target df = dedupCVByExpr (numericCVs ++ discreteCVs) + where + dfNoTarget = exclude [target] df + numericCVs = numericCondVecs cfg dfNoTarget df + discreteCVs = discreteCondVecs (targetInfoOrEmpty @a target df) cfg dfNoTarget + +targetInfoOrEmpty :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> TargetInfo a +targetInfoOrEmpty target df = fromMaybe (TargetInfo False Nothing V.empty) (mkTargetInfo @a target df) + +-- | Fit a tree at a given depth from a raw condition list (CART + TAO + prune). +buildTree :: forall a. (Columnable a, Ord a) => TreeConfig -> Int -> T.Text -> [Expr Bool] -> DataFrame -> Expr a +buildTree cfg depth target conds df = + pruneExpr (treeToExpr (taoOptimize @a cfg target conds df indices tree)) + where + tree = buildCartTree @a cfg{maxTreeDepth = depth} target df + indices = V.enumFromN 0 (nRows df) + +pruneTree :: forall a. (Columnable a) => Expr a -> Expr a +pruneTree = pruneExpr + +partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) +partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) + +-- | Laplace-smoothed Gini impurity of the target distribution. +calculateGini :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double +calculateGini target df + | n == 0 = 0 + | otherwise = 1 - sum (map (^ (2 :: Int)) probs) + where + n = fromIntegral (nRows df) + counts = getCounts @a target df + numClasses = fromIntegral (M.size counts) + probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts) + +majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a +majorityValue target df + | M.null counts = error "Empty DataFrame in leaf" + | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) + where + counts = getCounts @a target df + +getCounts :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int +getCounts target df = case interpret @a df (Col target) of + Left e -> throw e + Right (TColumn column) -> case toVector @a column of + Left e -> throw e + Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals) + +-- | The @p@-th percentile of an expression's values (@0@ on failure/empty). +percentile :: Int -> Expr Double -> DataFrame -> Double +percentile p expr df = case interpret @Double df expr of + Right (TColumn column) -> either (const 0) (percentileOfVec p) (toVector @Double column) + _ -> 0 + +percentileOfVec :: Int -> V.Vector Double -> Double +percentileOfVec p vals + | n == 0 = 0 + | otherwise = sorted V.! min (n - 1) (max 0 ((p * n) `div` 100)) + where + sorted = V.fromList (sort (V.toList vals)) + n = V.length sorted + +-- | A tree whose leaves hold class-probability distributions. +type ProbTree a = Tree (M.Map a Double) + +-- | Normalised class probabilities over a subset of training rows. +probsFromIndices :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> M.Map a Double +probsFromIndices target df indices = case interpret @a df (Col target) of + Right (TColumn column) -> either (const M.empty) (normaliseCounts indices) (toVector @a column) + _ -> M.empty + +normaliseCounts :: (Ord a) => V.Vector Int -> V.Vector a -> M.Map a Double +normaliseCounts indices vals = M.map (\c -> fromIntegral c / total) counts + where + counts = V.foldl' (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty indices + total = fromIntegral (V.length indices) :: Double + +-- | Re-label a fitted tree's leaves with class distributions, routing the +-- training data through the (unchanged) split conditions. +buildProbTree :: forall a. (Columnable a, Ord a) => Tree a -> T.Text -> DataFrame -> V.Vector Int -> ProbTree a +buildProbTree (Leaf _) target df indices = Leaf (probsFromIndices @a target df indices) +buildProbTree (Branch cond left right) target df indices = + Branch cond (buildProbTree @a left target df l) (buildProbTree @a right target df r) + where + (l, r) = partitionIndices cond df indices + +-- | Fit a TAO tree and return one probability expression per class. +fitProbTree :: forall a. (Columnable a, Ord a) => TreeConfig -> Expr a -> DataFrame -> M.Map a (Expr Double) +fitProbTree cfg (Col target) df = probExprs (buildProbTree @a pruned target df indices) + where + conds = nubByExpr (numericConditions cfg dfNoTarget ++ discreteConditions (targetInfoOrEmpty @a target df) cfg dfNoTarget) + dfNoTarget = exclude [target] df + indices = V.enumFromN 0 (nRows df) + pruned = pruneDead (taoOptimize @a cfg target conds df indices (buildCartTree @a cfg target df)) +fitProbTree _ expr _ = error ("Cannot create prob tree for compound expression: " ++ show expr) + +-- | Convert a 'ProbTree' into one @Expr Double@ per class. +probExprs :: forall a. (Columnable a, Ord a) => ProbTree a -> M.Map a (Expr Double) +probExprs tree = M.fromList [(c, classExpr c tree) | c <- nub (allClasses tree)] + +allClasses :: ProbTree a -> [a] +allClasses (Leaf m) = M.keys m +allClasses (Branch _ l r) = allClasses l ++ allClasses r + +classExpr :: (Ord a) => a -> ProbTree a -> Expr Double +classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m) +classExpr c (Branch cond l r) = F.ifThenElse cond (classExpr c l) (classExpr c r) diff --git a/src/DataFrame/DecisionTree/Linear.hs b/src/DataFrame/DecisionTree/Linear.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Linear.hs @@ -0,0 +1,133 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Oblique split candidates: fit an L1-regularised logistic hyperplane to the +-- care points (class-balanced) and convert it to a boolean condition, rejecting +-- all-zero and degenerate (single-side) hyperplanes. +module DataFrame.DecisionTree.Linear ( + bestLinearCandidate, + fitLinearCandidate, + careRowsFromFeatures, + careLabels, + featName, + imputeMean, + materializeFeatureForCare, +) where + +import DataFrame.DecisionTree.Numeric (NumExpr (..), numericCols) +import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..), TreeConfig (..)) +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr, getColumns) +import DataFrame.Internal.Interpreter (interpret) +import qualified DataFrame.LinearSolver as LS + +import Data.Maybe (catMaybes, fromMaybe, mapMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +-- | Best oblique candidate, or 'Nothing' when the linear path is disabled or +-- there are too few care points to fit on. +bestLinearCandidate :: TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) +bestLinearCandidate cfg df carePoints + | not (useLinearSolver cfg) = Nothing + | length carePoints < minCarePointsForLinear cfg = Nothing + | otherwise = fitLinearCandidate cfg df carePoints + +-- | Fit an L1 logistic regression to the care points and convert the resulting +-- hyperplane to a condition, or 'Nothing' when no numeric features exist or the +-- fitted model is all-zero or degenerate. +fitLinearCandidate :: TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) +fitLinearCandidate cfg df carePoints = case materializedFeatures df carePoints of + [] -> Nothing + mats -> linearFromFeatures cfg carePoints mats + +materializedFeatures :: DataFrame -> [CarePoint] -> [(T.Text, VU.Vector Double)] +materializedFeatures df carePoints = mapMaybe (materializeFeatureForCare df carePoints) (numericCols df) + +linearFromFeatures :: TreeConfig -> [CarePoint] -> [(T.Text, VU.Vector Double)] -> Maybe (Expr Bool) +linearFromFeatures cfg carePoints mats + | VU.all (== 0) weights = Nothing + | degenerateHyperplane rows weights (LS.lmIntercept model) = Nothing + | otherwise = Just (LS.modelToExpr model) + where + rows = careRowsFromFeatures (length carePoints) mats + labels = careLabels carePoints + model = LS.fitL1Logistic (solverConfigFor cfg labels) rows labels (V.fromList (map fst mats)) + weights = LS.lmWeights model + +solverConfigFor :: TreeConfig -> VU.Vector Double -> LS.SolverConfig +solverConfigFor cfg labels = (linearSolverConfig cfg){LS.scSampleWeights = classBalancedWeights labels} + +-- | Class-balanced sklearn-form weights @w_i = N / (2 · N_class)@ (mean 1), or +-- 'Nothing' in the degenerate one-class case (uniform weighting). +classBalancedWeights :: VU.Vector Double -> Maybe (VU.Vector Double) +classBalancedWeights labels + | nPos > 0 && nNeg > 0 = Just (VU.generate nCare weightAt) + | otherwise = Nothing + where + nCare = VU.length labels + nPos = VU.length (VU.filter (> 0) labels) + nNeg = nCare - nPos + weightAt i + | VU.unsafeIndex labels i > 0 = fromIntegral nCare / (2 * fromIntegral nPos) + | otherwise = fromIntegral nCare / (2 * fromIntegral nNeg) + +-- | A hyperplane is degenerate when every care row scores on the same side of +-- zero (equivalent to an invalid split, caught upstream). +degenerateHyperplane :: V.Vector (VU.Vector Double) -> VU.Vector Double -> Double -> Bool +degenerateHyperplane rows weights bias = + nCare > 0 && (VU.minimum scores > 0 || VU.maximum scores < 0) + where + nCare = V.length rows + scores = VU.generate nCare (\i -> VU.sum (VU.zipWith (*) weights (V.unsafeIndex rows i)) + bias) + +-- | Per-care-point feature rows from materialized columns (each of length +-- @nCare@, so indexing is in range). +careRowsFromFeatures :: Int -> [(T.Text, VU.Vector Double)] -> V.Vector (VU.Vector Double) +careRowsFromFeatures nCare mats = + V.generate nCare (\i -> VU.generate nFeat (\j -> snd (matsVec V.! j) VU.! i)) + where + matsVec = V.fromList mats + nFeat = V.length matsVec + +-- | Solver labels: @+1@ when 'GoLeft' is correct, @-1@ otherwise. +careLabels :: [CarePoint] -> VU.Vector Double +careLabels carePoints = VU.fromList [if cpCorrectDir cp == GoLeft then 1.0 else -1.0 | cp <- carePoints] + +-- | First column referenced by an expression, or a placeholder when none. +featName :: Expr b -> T.Text +featName expr = case getColumns expr of + (c : _) -> c + [] -> "" + +-- | Replace missing values with the mean of present ones; 'Nothing' when +-- nothing is present so the caller can drop the feature. +imputeMean :: [Maybe Double] -> Maybe (VU.Vector Double) +imputeMean careRaw = case catMaybes careRaw of + [] -> Nothing + present -> Just (VU.fromList [fromMaybe (mean present) mv | mv <- careRaw]) + where + mean xs = sum xs / fromIntegral (length xs) + +interpretDoubleVals :: DataFrame -> Expr Double -> Maybe (V.Vector Double) +interpretDoubleVals df expr = case interpret @Double df expr of + Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) + _ -> Nothing + +interpretMaybeDoubleVals :: DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) +interpretMaybeDoubleVals df expr = case interpret @(Maybe Double) df expr of + Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) + _ -> Nothing + +-- | Materialize a 'NumExpr' over the care rows; 'Nothing' on interpret failure +-- or (nullable) when no care point has a present value, else mean-imputed. +materializeFeatureForCare :: DataFrame -> [CarePoint] -> NumExpr -> Maybe (T.Text, VU.Vector Double) +materializeFeatureForCare df carePoints (NDouble expr) = do + vals <- interpretDoubleVals df expr + Just (featName expr, VU.fromList [vals V.! cpIndex cp | cp <- carePoints]) +materializeFeatureForCare df carePoints (NMaybeDouble expr) = do + vals <- interpretMaybeDoubleVals df expr + imputed <- imputeMean [vals V.! cpIndex cp | cp <- carePoints] + Just (featName expr, imputed) diff --git a/src/DataFrame/DecisionTree/Numeric.hs b/src/DataFrame/DecisionTree/Numeric.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Numeric.hs @@ -0,0 +1,220 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Numeric split candidates: per-column Double expressions, arithmetic +-- expansion, and threshold conditions. 'numericCondVecs' materializes the +-- pool with a single interpret per distinct expression, deriving every +-- threshold/operator truth vector by direct comparison. +module DataFrame.DecisionTree.Numeric ( + NumExpr (..), + numExprCols, + numExprEq, + combineNumExprs, + numericConditions, + generateNumericConds, + percentilesOf, + numericCondVecs, + numericExprsWithTerms, + numericCols, +) where + +import DataFrame.DecisionTree.CondVec (CondVec (..)) +import DataFrame.DecisionTree.Types (SynthConfig (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..), eqExpr, getColumns, normalize) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operators + +import Data.List (sort) +import Data.Maybe (fromMaybe) +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | A numeric feature expression, non-nullable or nullable. +data NumExpr + = NDouble !(Expr Double) + | NMaybeDouble !(Expr (Maybe Double)) + +numExprCols :: NumExpr -> [T.Text] +numExprCols (NDouble e) = getColumns e +numExprCols (NMaybeDouble e) = getColumns e + +numExprEq :: NumExpr -> NumExpr -> Bool +numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2 +numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2 +numExprEq _ _ = False + +-- | Safe division: @0@ (or @Nothing@) where the divisor is zero. +safeDivD :: Expr Double -> Expr Double -> Expr Double +safeDivD a b = F.ifThenElse (b ./= F.lit (0 :: Double)) (a ./ b) (F.lit (0 :: Double)) + +safeDivMaybe :: Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double) +safeDivMaybe nonZero q = F.ifThenElse nonZero q (F.lit (Nothing :: Maybe Double)) + +-- | Arithmetic combinations (@+@, @-@, @*@, safe @/@) of two numeric exprs. +combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] +combineNumExprs (NDouble e1) (NDouble e2) = + map NDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivD e1 e2] +combineNumExprs (NDouble e1) (NMaybeDouble e2) = + map NMaybeDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2)] +combineNumExprs (NMaybeDouble e1) (NDouble e2) = + map NMaybeDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivMaybe (e2 ./= F.lit (0 :: Double)) (e1 ./ e2)] +combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) = + map NMaybeDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2)] + +numericConditions :: TreeConfig -> DataFrame -> [Expr Bool] +numericConditions = generateNumericConds + +generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool] +generateNumericConds cfg df = do + expr <- numericExprsWithTerms (synthConfig cfg) df + threshold <- numericThresholds cfg df expr + condsFromExpr expr threshold + +numericThresholds :: TreeConfig -> DataFrame -> NumExpr -> [Double] +numericThresholds cfg df (NDouble e) = thresholdsForExpr cfg df e +numericThresholds cfg df (NMaybeDouble e) = thresholdsForExpr cfg df (F.fromMaybe 0 e) + +thresholdsForExpr :: TreeConfig -> DataFrame -> Expr Double -> [Double] +thresholdsForExpr cfg df e = + maybe [] (percentilesOf (percentiles cfg) . V.toList) (interpretDoubleCol df e) + +condsFromExpr :: NumExpr -> Double -> [Expr Bool] +condsFromExpr (NDouble e) t = [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] +condsFromExpr (NMaybeDouble e) t = map (F.fromMaybe False) [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] + +-- | Percentile thresholds for a value list: sort once, index each percentile. +-- Shared by 'generateNumericConds' and 'numericCondVecs' for identical results. +percentilesOf :: [Int] -> [Double] -> [Double] +percentilesOf ps valsList + | n == 0 = [] + | otherwise = map (\p -> sortedV V.! min (n - 1) (max 0 (p * n `div` 100))) ps + where + !sortedV = V.fromList (sort valsList) + !n = V.length sortedV + +interpretDoubleCol :: DataFrame -> Expr Double -> Maybe (V.Vector Double) +interpretDoubleCol df e = case interpret @Double df e of + Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) + _ -> Nothing + +interpretMaybeDoubleCol :: DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) +interpretMaybeDoubleCol df e = case interpret @(Maybe Double) df e of + Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) + _ -> Nothing + +-- | Materialize the numeric pool with one interpret per distinct expression, +-- deriving each threshold/operator truth vector by direct comparison. +-- Byte-identical to materializing 'numericConditions' one at a time, but +-- avoids re-interpreting each LHS per threshold and operator. +numericCondVecs :: TreeConfig -> DataFrame -> DataFrame -> [CondVec] +numericCondVecs cfg dfGen df = concatMap forExpr (numericExprsWithTerms (synthConfig cfg) dfGen) + where + forExpr (NDouble e) = maybe [] (condsForDouble cfg e) (interpretDoubleCol df e) + forExpr (NMaybeDouble e) = maybe [] (condsForMaybe cfg e) (interpretMaybeDoubleCol df e) + +condsForDouble :: TreeConfig -> Expr Double -> V.Vector Double -> [CondVec] +condsForDouble cfg e vals = concatMap (doubleCondsAt e vals (V.length vals)) ts + where + ts = percentilesOf (percentiles cfg) (V.toList vals) + +doubleCondsAt :: Expr Double -> V.Vector Double -> Int -> Double -> [CondVec] +doubleCondsAt e vals n t = + [ CondVec (e .<= F.lit t) (gen (<= t)) + , CondVec (e .>= F.lit t) (gen (>= t)) + , CondVec (e .< F.lit t) (gen (< t)) + , CondVec (e .> F.lit t) (gen (> t)) + ] + where + gen p = VU.generate n (\i -> p (vals V.! i)) + +condsForMaybe :: TreeConfig -> Expr (Maybe Double) -> V.Vector (Maybe Double) -> [CondVec] +condsForMaybe cfg e mvals = concatMap (maybeCondsAt e mvals (V.length mvals)) ts + where + ts = percentilesOf (percentiles cfg) (map (fromMaybe 0) (V.toList mvals)) + +maybeCondsAt :: Expr (Maybe Double) -> V.Vector (Maybe Double) -> Int -> Double -> [CondVec] +maybeCondsAt e mvals n t = + [ CondVec (F.fromMaybe False (e .<= F.lit t)) (gen (<= t)) + , CondVec (F.fromMaybe False (e .>= F.lit t)) (gen (>= t)) + , CondVec (F.fromMaybe False (e .< F.lit t)) (gen (< t)) + , CondVec (F.fromMaybe False (e .> F.lit t)) (gen (> t)) + ] + where + gen p = VU.generate n (\i -> maybe False p (mvals V.! i)) + +-- | Arithmetic candidate expansion, generated already-deduped: each round +-- combines @frontier × base@ and admits only normalized-novel candidates. +-- Produces @base@ plus @maxExprDepth-1@ combination rounds. +numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr] +numericExprsWithTerms cfg df + | not (enableArithOps cfg) = base + | otherwise = base ++ expandRounds cfg base (max 0 (maxExprDepth cfg - 1)) base seen0 + where + base = numericCols df + seen0 = Set.fromList (map keyNum base) + +keyNum :: NumExpr -> String +keyNum (NDouble e) = show (normalize e) +keyNum (NMaybeDouble e) = show (normalize e) + +isDisallowed :: SynthConfig -> NumExpr -> NumExpr -> Bool +isDisallowed cfg e1 e2 = + any (\(l, r) -> l `elem` cols && r `elem` cols) (disallowedCombinations cfg) + where + cols = numExprCols e1 <> numExprCols e2 + +roundProducts :: SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr] +roundProducts cfg frontier base = + [c | e1 <- frontier, e2 <- base, not (numExprEq e1 e2), not (isDisallowed cfg e1 e2), c <- combineNumExprs e1 e2] + +expandRounds :: SynthConfig -> [NumExpr] -> Int -> [NumExpr] -> Set.Set String -> [NumExpr] +expandRounds _ _ 0 _ _ = [] +expandRounds cfg base d frontier seen + | null admitted = [] + | otherwise = admitted ++ expandRounds cfg base (d - 1) admitted seen' + where + (admitted, seen') = admitNovel seen (roundProducts cfg frontier base) + +admitNovel :: Set.Set String -> [NumExpr] -> ([NumExpr], Set.Set String) +admitNovel seen0 = go seen0 [] + where + go seen acc [] = (reverse acc, seen) + go seen acc (c : cs) + | keyNum c `Set.member` seen = go seen acc cs + | otherwise = go (Set.insert (keyNum c) seen) (c : acc) cs + +numericCols :: DataFrame -> [NumExpr] +numericCols df = concatMap (numExprsOfColumn df) (columnNames df) + +numExprsOfColumn :: DataFrame -> T.Text -> [NumExpr] +numExprsOfColumn df colName = case unsafeGetColumn colName df of + UnboxedColumn Nothing (_ :: VU.Vector b) -> strictNumeric @b colName + BoxedColumn (Just _) (_ :: V.Vector b) -> nullableNumeric @b colName + UnboxedColumn (Just _) (_ :: VU.Vector b) -> nullableNumeric @b colName + _ -> [] + +strictNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] +strictNumeric c = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [NDouble (Col c)] + Nothing -> case sIntegral @b of + STrue -> [NDouble (F.toDouble (Col @b c))] + SFalse -> [] + +nullableNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] +nullableNumeric c = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [NMaybeDouble (Col @(Maybe b) c)] + Nothing -> case sIntegral @b of + STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) c))] + SFalse -> [] diff --git a/src/DataFrame/DecisionTree/Pool.hs b/src/DataFrame/DecisionTree/Pool.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Pool.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Candidate-pool scoring and boolean expansion: penalized scoring, diverse +-- top-K selection, AND/OR saturation, and structural/truth-vector dedup. The +-- per-node scoring scans run in parallel chunks. +module DataFrame.DecisionTree.Pool ( + evalWithPenaltyVec, + primaryColExpr, + primaryColCV, + takeDiverse, + candidateParChunk, + bestDiscreteCandidate, + boolExprsVec, + DedupMode (..), + saturateCandidates, + roundProducts, + admitKeys, + admitVecs, + dedupCVByExpr, + nubByExpr, +) where + +import DataFrame.DecisionTree.CondVec (CondVec (..), combineAndVec, combineOrVec, countErrorsByVec) +import DataFrame.DecisionTree.Types (CarePoint, SynthConfig (..), TreeConfig (..)) +import DataFrame.Internal.Expression (Expr, compareExpr, eSize, eqExpr, getColumns, normalize) + +import Control.Parallel.Strategies (parListChunk, rdeepseq, using) +import Data.Function (on) +import Data.List (minimumBy, sortBy) +import qualified Data.Map.Strict as M +import qualified Data.Set as Set +import qualified Data.Text as T +import qualified Data.Vector.Unboxed as VU + +-- | Penalized score of a candidate: care-point errors plus a complexity +-- penalty, tie-broken by expression size. +evalWithPenaltyVec :: TreeConfig -> [CarePoint] -> CondVec -> (Int, Int) +evalWithPenaltyVec cfg carePoints cv = (countErrorsByVec (cvVec cv) carePoints + penalty, sz) + where + sz = eSize (cvExpr cv) + penalty = floor (complexityPenalty (synthConfig cfg) * fromIntegral sz) + +-- | First referenced column of a condition (a sentinel for literal-only ones), +-- used by 'takeDiverse' to enforce per-column diversity. +primaryColExpr :: Expr Bool -> T.Text +primaryColExpr e = case getColumns e of + [] -> "" + (c : _) -> c + +primaryColCV :: CondVec -> T.Text +primaryColCV = primaryColExpr . cvExpr + +-- | Keep the first @k@ of an already-sorted list, admitting at most @quota@ per +-- primary column (@Nothing@ disables the per-column cap). +takeDiverse :: Int -> Maybe Int -> (a -> T.Text) -> [a] -> [a] +takeDiverse k Nothing _ = take k +takeDiverse k (Just quota) primary = go M.empty 0 + where + go !_ !_ [] = [] + go !seen !n (x : xs) + | n >= k = [] + | M.findWithDefault 0 col seen >= quota = go seen n xs + | otherwise = x : go (M.insertWith (+) col 1 seen) (n + 1) xs + where + !col = primary x + +-- | Chunk size for the parallel per-node candidate scans; tuned by an -N +-- sweep, not correctness-affecting. +candidateParChunk :: Int +candidateParChunk = 64 + +-- | Decorate candidates with their penalty in parallel chunks, forcing only +-- the @(Int, Int)@ key so the order (hence later sorts/minima) is preserved. +decorate :: (CondVec -> (Int, Int)) -> [CondVec] -> [((Int, Int), CondVec)] +decorate penaltyCV xs = zip (map penaltyCV xs `using` parListChunk candidateParChunk rdeepseq) xs + +-- | The diverse top-@expressionPairs@ valid candidates by penalty. +sortedTopK :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> [CondVec] +sortedTopK cfg penaltyCV validCondVecs = + map snd (takeDiverse (expressionPairs cfg) (perColumnQuota (synthConfig cfg)) (primaryColCV . snd) sorted) + where + sorted = sortBy (compare `on` fst) (decorate penaltyCV validCondVecs) + +-- | Lowest-penalty candidate after boolean saturation of the diverse top-K. +bestDiscreteCandidate :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> Maybe CondVec +bestDiscreteCandidate _ _ [] = Nothing +bestDiscreteCandidate cfg penaltyCV validCondVecs = + case saturateCandidates Structural (boolExpansion (synthConfig cfg)) (sortedTopK cfg penaltyCV validCondVecs) of + [] -> Nothing + xs -> Just (snd (minimumBy (compare `on` fst) (decorate penaltyCV xs))) + +-- | AND/OR expansion of cached conditions to depth @maxDepth@ (each +-- combination is a single vector op, not an interpret). +boolExprsVec :: [CondVec] -> [CondVec] -> Int -> Int -> [CondVec] +boolExprsVec baseExprs prevExprs depth maxDepth + | depth == 0 = baseExprs ++ boolExprsVec baseExprs prevExprs (depth + 1) maxDepth + | depth >= maxDepth = [] + | otherwise = combined ++ boolExprsVec baseExprs combined (depth + 1) maxDepth + where + combined = roundProducts prevExprs baseExprs + +data DedupMode = Structural | TruthVector + deriving (Eq, Show) + +-- | Saturate the pool with AND/OR combinations, deduplicating structurally +-- (byte-identical, first occurrence kept) or by truth vector (opt-in). +saturateCandidates :: DedupMode -> Int -> [CondVec] -> [CondVec] +saturateCandidates Structural maxDepth base = base' ++ go 1 base' seen0 + where + (base', seen0) = admitKeys Set.empty base + go !depth frontier seen + | depth >= maxDepth || null frontier = [] + | otherwise = let (admitted, seen') = admitKeys seen (roundProducts frontier base) in admitted ++ go (depth + 1) admitted seen' +saturateCandidates TruthVector maxDepth base = M.elems (go 1 frontier0 reps0) + where + (reps0, frontier0) = admitVecs M.empty base + go !depth frontier reps + | depth >= maxDepth || null frontier = reps + | otherwise = let (reps', admitted) = admitVecs reps (roundProducts frontier base) in go (depth + 1) admitted reps' + +-- | One combination round: @frontier × base@ via AND then OR, skipping +-- self-pairs (mirrors 'boolExprsVec' for byte-identical structural output). +roundProducts :: [CondVec] -> [CondVec] -> [CondVec] +roundProducts frontier base = + [c | e1 <- frontier, e2 <- base, not (eqExpr (cvExpr e1) (cvExpr e2)), c <- [combineAndVec e1 e2, combineOrVec e1 e2]] + +-- | Admit candidates with a not-yet-seen normalized form, preserving order. +admitKeys :: Set.Set String -> [CondVec] -> ([CondVec], Set.Set String) +admitKeys = go [] + where + go acc seen [] = (reverse acc, seen) + go acc !seen (c : cs) + | structuralKey c `Set.member` seen = go acc seen cs + | otherwise = go (c : acc) (Set.insert (structuralKey c) seen) cs + +structuralKey :: CondVec -> String +structuralKey = show . normalize . cvExpr + +-- | Admit candidates by distinct truth vector, keeping the smallest-expression +-- representative per vector. +admitVecs :: M.Map (VU.Vector Bool) CondVec -> [CondVec] -> (M.Map (VU.Vector Bool) CondVec, [CondVec]) +admitVecs = go [] + where + go acc reps [] = (reps, reverse acc) + go acc !reps (c : cs) = case M.lookup (cvVec c) reps of + Nothing -> go (c : acc) (M.insert (cvVec c) c reps) cs + Just r -> go acc (M.insert (cvVec c) (smaller r c) reps) cs + +smaller :: CondVec -> CondVec -> CondVec +smaller a b = case compare (eSize (cvExpr a)) (eSize (cvExpr b)) of + LT -> a + GT -> b + EQ -> if compareExpr (cvExpr a) (cvExpr b) /= GT then a else b + +-- | Deduplicate 'CondVec's by normalized 'cvExpr', keeping the first. +dedupCVByExpr :: [CondVec] -> [CondVec] +dedupCVByExpr = go Set.empty + where + go _ [] = [] + go seen (cv : cvs) + | structuralKey cv `Set.member` seen = go seen cvs + | otherwise = cv : go (Set.insert (structuralKey cv) seen) cvs + +-- | Deduplicate expressions by normalized form, keeping the first. +nubByExpr :: [Expr Bool] -> [Expr Bool] +nubByExpr = go Set.empty + where + go _ [] = [] + go seen (e : es) + | k `Set.member` seen = go seen es + | otherwise = e : go (Set.insert k seen) es + where + k = show (normalize e) diff --git a/src/DataFrame/DecisionTree/Predict.hs b/src/DataFrame/DecisionTree/Predict.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Predict.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Prediction, care-point identification, node validity, and tree loss. The +-- batched, cache-aware variants resolve each branch condition's truth vector +-- once per call instead of once per row. +module DataFrame.DecisionTree.Predict ( + predictWithTree, + predictManyWithTree, + predictManyWithTreeCached, + identifyCarePoints, + identifyCarePointsCached, + countCarePointErrors, + partitionIndices, + partitionIndicesCached, + majorityValueFromIndices, + computeTreeLoss, + computeTreeLossCached, + isValidAtNode, +) where + +import DataFrame.DecisionTree.CondVec (CondCache, countErrorsByVec, lookupCondVec) +import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..), Tree (..), TreeConfig (..)) +import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) + +import Control.Exception (throw) +import Control.Monad.ST (ST) +import Data.Function (on) +import Data.List (maximumBy) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU + +-- | A condition's truth vector over the DataFrame, or 'Nothing' on a +-- type/interpret failure (callers default such rows to the left child). +branchBool :: DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) +branchBool df cond = case interpret @Bool df cond of + Right (TColumn column) -> either (const Nothing) Just (toVector @Bool @VU.Vector column) + _ -> Nothing + +-- | The target column as a label vector, or 'Nothing' on failure. +interpretLabelCol :: forall a. (Columnable a) => DataFrame -> T.Text -> Maybe (V.Vector a) +interpretLabelCol df target = case interpret @a df (Col target) of + Right (TColumn column) -> either (const Nothing) Just (toVector @a column) + _ -> Nothing + +-- | Predict the label for a single row by walking a fixed tree (@True@ → left). +predictWithTree :: forall a. (Columnable a) => T.Text -> DataFrame -> Int -> Tree a -> a +predictWithTree _ _ _ (Leaf v) = v +predictWithTree target df idx (Branch cond left right) = + predictWithTree @a target df idx (childFor cond left right idx df) + +childFor :: Expr Bool -> Tree a -> Tree a -> Int -> DataFrame -> Tree a +childFor cond left right idx df = case branchBool df cond of + Nothing -> left + Just boolVals -> if boolVals VU.! idx then left else right + +predictManyWithTree :: forall a. (Columnable a) => Tree a -> DataFrame -> V.Vector Int -> V.Vector a +predictManyWithTree = predictManyWithTreeCached @a M.empty + +-- | 'predictManyWithTree' resolving each branch condition through a 'CondCache'. +-- Each condition is read at most once per call rather than once per row. +predictManyWithTreeCached :: forall a. (Columnable a) => CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a +predictManyWithTreeCached cache tree df indices = V.create $ do + mv <- VM.new (V.length indices) + fill mv (V.zip (V.enumFromN 0 (V.length indices)) indices) tree + pure mv + where + fill :: VM.MVector s a -> V.Vector (Int, Int) -> Tree a -> ST s () + fill mv prs (Leaf v) = V.mapM_ (\(p, _) -> VM.write mv p v) prs + fill mv prs (Branch cond left right) = case lookupCondVec cache df cond of + Nothing -> fill mv prs left + Just boolVals -> fillSplit mv (V.partition (\(_, i) -> boolVals VU.! i) prs) left right + + fillSplit :: VM.MVector s a -> (V.Vector (Int, Int), V.Vector (Int, Int)) -> Tree a -> Tree a -> ST s () + fillSplit mv (leftPrs, rightPrs) left right = fill mv leftPrs left >> fill mv rightPrs right + +identifyCarePoints :: forall a. (Columnable a) => T.Text -> DataFrame -> V.Vector Int -> Tree a -> Tree a -> [CarePoint] +identifyCarePoints = identifyCarePointsCached @a M.empty + +-- | Rows the parent must route to a specific child for the (fixed) subtrees to +-- classify correctly; a 'CondCache' avoids re-interpreting subtree conditions. +identifyCarePointsCached :: forall a. (Columnable a) => CondCache -> T.Text -> DataFrame -> V.Vector Int -> Tree a -> Tree a -> [CarePoint] +identifyCarePointsCached cache target df indices leftTree rightTree = + maybe [] carePoints (interpretLabelCol @a df target) + where + leftPreds = predictManyWithTreeCached cache leftTree df indices + rightPreds = predictManyWithTreeCached cache rightTree df indices + carePoints targetVals = V.toList (V.imapMaybe (checkPoint targetVals leftPreds rightPreds) indices) + +checkPoint :: (Eq a) => V.Vector a -> V.Vector a -> V.Vector a -> Int -> Int -> Maybe CarePoint +checkPoint targetVals leftPreds rightPreds k idx = + case (leftPreds V.! k == trueLabel, rightPreds V.! k == trueLabel) of + (True, False) -> Just (CarePoint idx GoLeft) + (False, True) -> Just (CarePoint idx GoRight) + _ -> Nothing + where + trueLabel = targetVals V.! idx + +-- | Care points a free condition misroutes (uncached; for the linear path). +countCarePointErrors :: Expr Bool -> DataFrame -> [CarePoint] -> Int +countCarePointErrors cond df carePoints = + maybe (length carePoints) (`countErrorsByVec` carePoints) (branchBool df cond) + +partitionIndices :: Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int) +partitionIndices = partitionIndicesCached M.empty + +-- | 'partitionIndices' resolving the condition through a 'CondCache'; a miss +-- routes every index left (matching the uncached fallback). +partitionIndicesCached :: CondCache -> Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int) +partitionIndicesCached cache cond df indices = case lookupCondVec cache df cond of + Nothing -> (indices, V.empty) + Just boolVals -> V.partition (boolVals VU.!) indices + +-- | A split is valid at a node when both children keep at least 'minLeafSize'. +isValidAtNode :: TreeConfig -> DataFrame -> V.Vector Int -> Expr Bool -> Bool +isValidAtNode cfg df indices c = + V.length t >= minLeafSize cfg && V.length f >= minLeafSize cfg + where + (t, f) = partitionIndices c df indices + +majorityValueFromIndices :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> a +majorityValueFromIndices target df indices = majorityOf (countLabels (labelColOrThrow @a df target) indices) + +labelColOrThrow :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a +labelColOrThrow df target = case interpret @a df (Col target) of + Left e -> throw e + Right (TColumn column) -> either throw id (toVector @a column) + +countLabels :: (Ord a) => V.Vector a -> V.Vector Int -> M.Map a Int +countLabels vals = V.foldl' (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty + +majorityOf :: M.Map a Int -> a +majorityOf counts + | M.null counts = error "Empty indices in majorityValueFromIndices" + | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) + +computeTreeLoss :: forall a. (Columnable a) => T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double +computeTreeLoss = computeTreeLossCached @a M.empty + +-- | 0/1 loss of a tree over @indices@, with a 'CondCache' for the predictions. +computeTreeLossCached :: forall a. (Columnable a) => CondCache -> T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double +computeTreeLossCached cache target df indices tree + | V.null indices = 0 + | otherwise = maybe 1.0 (treeLoss cache tree df indices) (interpretLabelCol @a df target) + +treeLoss :: (Columnable a) => CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a -> Double +treeLoss cache tree df indices targetVals = + fromIntegral (countMismatches targetVals indices preds) / fromIntegral (V.length indices) + where + preds = predictManyWithTreeCached cache tree df indices + +countMismatches :: (Eq a) => V.Vector a -> V.Vector Int -> V.Vector a -> Int +countMismatches targetVals indices preds = + V.length (V.ifilter (\k _ -> targetVals V.! (indices V.! k) /= preds V.! k) preds) diff --git a/src/DataFrame/DecisionTree/Prune.hs b/src/DataFrame/DecisionTree/Prune.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Prune.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Post-convergence simplification of a fitted tree and its expression form: +-- drop branches forced by path-condition entailment, collapse identical +-- siblings, and fold redundant nested conditionals. +module DataFrame.DecisionTree.Prune ( + pruneDead, + treeEq, + pruneExpr, +) where + +import DataFrame.DecisionTree.Types (Tree (..)) +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr (..), eqExpr) +import DataFrame.Internal.Simplify (PredFact, entails, factFalse, factTrue) + +-- | Drop branches whose test is forced by the path conditions reaching them, +-- and collapse @Branch c t t@ to @t@. Sound for the decidable threshold subset; +-- other tests are left untouched. +pruneDead :: forall a. (Columnable a) => Tree a -> Tree a +pruneDead = go [] + where + go :: [PredFact] -> Tree a -> Tree a + go _ (Leaf v) = Leaf v + go facts (Branch cond left right) = case entails facts cond of + Just True -> go facts left + Just False -> go facts right + Nothing -> reconcile cond (go (addFact (factTrue cond) facts) left) (go (addFact (factFalse cond) facts) right) + +reconcile :: (Columnable a) => Expr Bool -> Tree a -> Tree a -> Tree a +reconcile cond left right + | treeEq left right = left + | otherwise = Branch cond left right + +addFact :: Maybe PredFact -> [PredFact] -> [PredFact] +addFact (Just f) fs = f : fs +addFact Nothing fs = fs + +treeEq :: (Columnable a) => Tree a -> Tree a -> Bool +treeEq (Leaf x) (Leaf y) = x == y +treeEq (Branch c1 l1 r1) (Branch c2 l2 r2) = eqExpr c1 c2 && treeEq l1 l2 && treeEq r1 r2 +treeEq _ _ = False + +-- | Recursively fold @If@ expressions whose branches coincide or nest the same +-- condition; leave other expressions structurally unchanged. +pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a +pruneExpr (If cond t0 f0) = collapseIf cond (pruneExpr t0) (pruneExpr f0) +pruneExpr (Unary op e) = Unary op (pruneExpr e) +pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) +pruneExpr e = e + +collapseIf :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a +collapseIf cond t f + | eqExpr t f = t + | If ci ti _ <- t, eqExpr cond ci = If cond ti f + | If ci _ fi <- f, eqExpr cond ci = If cond t fi + | otherwise = If cond t f diff --git a/src/DataFrame/DecisionTree/Tao.hs b/src/DataFrame/DecisionTree/Tao.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Tao.hs @@ -0,0 +1,168 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Tree Alternating Optimization: hold the tree fixed and re-optimize one node +-- at a time, bottom-up, minimizing care-point misroutes. Sibling subtrees at a +-- depth level are independent and optimized in parallel. +module DataFrame.DecisionTree.Tao ( + taoOptimize, + taoOptimizeCV, + taoIteration, + taoIterationCV, + optimizeNode, + findBestSplitTAO, +) where + +import DataFrame.DecisionTree.CondVec +import DataFrame.DecisionTree.Linear (bestLinearCandidate) +import DataFrame.DecisionTree.Pool (bestDiscreteCandidate, candidateParChunk, evalWithPenaltyVec) +import DataFrame.DecisionTree.Predict +import DataFrame.DecisionTree.Prune (pruneDead) +import DataFrame.DecisionTree.Types +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr) + +import Control.Parallel (par, pseq) +import Control.Parallel.Strategies (parListChunk, rdeepseq, using) +import Data.Function (on) +import Data.List (foldl', minimumBy) +import Data.Maybe (catMaybes, mapMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +-- | The constant per-fit context threaded through the node-optimization +-- recursion (the cache is rebuilt each iteration). +data TaoEnv = TaoEnv + { teCache :: !CondCache + , teCfg :: !TreeConfig + , teTarget :: !T.Text + , teConds :: ![CondVec] + , teDf :: !DataFrame + } + +-- | Public TAO entry point over raw conditions; materializes each once. +taoOptimize :: forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> V.Vector Int -> Tree a -> Tree a +taoOptimize cfg target conds df = + taoOptimizeCV @a cfg target (mapMaybe (materializeCondVec df) conds) df + +-- | TAO outer loop over pre-evaluated candidates: iterate until the iteration +-- budget or convergence tolerance is reached, then prune dead branches. +taoOptimizeCV :: forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> [CondVec] -> DataFrame -> V.Vector Int -> Tree a -> Tree a +taoOptimizeCV cfg target condVecs df rootIndices initialTree = + go 0 initialTree (lossWith baseCache initialTree) + where + baseCache = condCacheFromVecs condVecs + lossWith cache = computeTreeLossCached @a cache target df rootIndices + go iter tree prevLoss + | iter >= taoIterations cfg = pruneDead tree + | prevLoss - newLoss < taoConvergenceTol cfg = pruneDead tree' + | otherwise = go (iter + 1) tree' newLoss + where + cache = addTreeCondsToCache df tree baseCache + tree' = taoIterationCV @a cache cfg target condVecs df rootIndices tree + newLoss = lossWith cache tree' + +-- | Public single-iteration entry point. +taoIteration :: forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> V.Vector Int -> Tree a -> Tree a +taoIteration cfg target conds df rootIndices tree = + let condVecs = mapMaybe (materializeCondVec df) conds + cache = addTreeCondsToCache df tree (condCacheFromVecs condVecs) + in taoIterationCV @a cache cfg target condVecs df rootIndices tree + +-- | One bottom-to-top sweep: re-optimize every node level by level. +taoIterationCV :: forall a. (Columnable a, Ord a) => CondCache -> TreeConfig -> T.Text -> [CondVec] -> DataFrame -> V.Vector Int -> Tree a -> Tree a +taoIterationCV cache cfg target condVecs df rootIndices tree = + foldl' (optimizeDepthLevel env rootIndices) tree [treeDepth tree, treeDepth tree - 1 .. 0] + where + env = TaoEnv cache cfg target condVecs df + +optimizeDepthLevel :: forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Int -> Tree a +optimizeDepthLevel env rootIndices tree = optimizeAtDepth @a env rootIndices tree 0 + +optimizeAtDepth :: forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Int -> Int -> Tree a +optimizeAtDepth env indices tree currentDepth targetDepth + | currentDepth == targetDepth = optimizeNode @a env indices tree + | otherwise = case tree of + Leaf v -> Leaf v + Branch cond left right -> optimizeChildren @a env indices cond left right currentDepth targetDepth + +-- | Optimize the two subtrees over their disjoint index sets, scoring the left +-- in parallel with the right (the cache is read-only, so this is pure). +optimizeChildren :: forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Int -> Int -> Tree a +optimizeChildren env indices cond left right currentDepth targetDepth = + forceTreeWork left' `par` (forceTreeWork right' `pseq` Branch cond left' right') + where + (indicesL, indicesR) = partitionIndicesCached (teCache env) cond (teDf env) indices + left' = optimizeAtDepth @a env indicesL left (currentDepth + 1) targetDepth + right' = optimizeAtDepth @a env indicesR right (currentDepth + 1) targetDepth + +-- | Force a subtree's optimization work to WHNF so the parallel scheduler has +-- something substantial to evaluate; pure and value-preserving. +forceTreeWork :: Tree a -> () +forceTreeWork (Leaf v) = v `seq` () +forceTreeWork (Branch c l r) = c `seq` forceTreeWork l `seq` forceTreeWork r + +-- | Re-optimize one node: pick its best split, or collapse to a leaf when the +-- node is empty or the chosen split underflows 'minLeafSize'. +optimizeNode :: forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Tree a +optimizeNode env indices tree + | V.null indices = tree + | otherwise = case tree of + Leaf _ -> leaf + Branch oldCond left right -> rebuiltBranch env indices oldCond left right leaf + where + leaf = Leaf (majorityValueFromIndices @a (teTarget env) (teDf env) indices) + +rebuiltBranch :: forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Tree a -> Tree a +rebuiltBranch env indices oldCond left right leaf + | underflows = leaf + | otherwise = Branch newCond left right + where + newCond = findBestSplitTAO @a env indices left right oldCond + (l, r) = partitionIndicesCached (teCache env) newCond (teDf env) indices + underflows = V.length l < minLeafSize (teCfg env) || V.length r < minLeafSize (teCfg env) + +-- | The lowest-penalty replacement condition for a node, falling back to the +-- current condition when no valid candidate beats it. +findBestSplitTAO :: forall a. (Columnable a) => TaoEnv -> V.Vector Int -> Tree a -> Tree a -> Expr Bool -> Expr Bool +findBestSplitTAO env indices leftTree rightTree currentCond + | V.null indices || null carePoints = currentCond + | pureReplacementLinear cfg, Just c <- linearCandidate, isValidAtNode cfg (teDf env) indices c = c + | otherwise = bestOfPool penaltyCV currentCond pool + where + cfg = teCfg env + carePoints = identifyCarePointsCached @a (teCache env) (teTarget env) (teDf env) indices leftTree rightTree + penaltyCV = evalWithPenaltyVec cfg carePoints + linearCandidate = bestLinearCandidate cfg (teDf env) carePoints + valid = filterValidCandidates cfg indices (teConds env) + pool = candidatePool env indices currentCond (bestDiscreteCandidate cfg penaltyCV valid) linearCandidate + +bestOfPool :: (CondVec -> (Int, Int)) -> Expr Bool -> [CondVec] -> Expr Bool +bestOfPool _ currentCond [] = currentCond +bestOfPool penaltyCV _ pool = cvExpr (minimumBy (compare `on` penaltyCV) pool) + +-- | Validity-filtered candidates the node could split on: both children must +-- keep at least 'minLeafSize'. Scored in parallel chunks, order preserved. +filterValidCandidates :: TreeConfig -> V.Vector Int -> [CondVec] -> [CondVec] +filterValidCandidates cfg indices condVecs = map snd (filter fst (zip validity condVecs)) + where + validity = map (validAtNode cfg indices) condVecs `using` parListChunk candidateParChunk rdeepseq + +validAtNode :: TreeConfig -> V.Vector Int -> CondVec -> Bool +validAtNode cfg indices cv = nTrue >= minLeaf && (V.length indices - nTrue) >= minLeaf + where + minLeaf = minLeafSize cfg + nTrue = V.foldl' (\ !acc i -> if cvVec cv VU.! i then acc + 1 else acc) (0 :: Int) indices + +-- | The candidate pool to minimize over: the current condition, the best +-- discrete candidate, and the linear candidate, each kept only if valid. +candidatePool :: TaoEnv -> V.Vector Int -> Expr Bool -> Maybe CondVec -> Maybe (Expr Bool) -> [CondVec] +candidatePool env indices currentCond discreteCV linearCandidate = + filter (isValidAtNode (teCfg env) (teDf env) indices . cvExpr) (catMaybes [currentCV, discreteCV, linearCV]) + where + currentCV = CondVec currentCond <$> lookupCondVec (teCache env) (teDf env) currentCond + linearCV = linearCandidate >>= materializeCondVec (teDf env) diff --git a/src/DataFrame/DecisionTree/Types.hs b/src/DataFrame/DecisionTree/Types.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/DecisionTree/Types.hs @@ -0,0 +1,185 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Shared types, configuration and ordering machinery for the decision-tree +-- learner. Imported by every other @DataFrame.DecisionTree.*@ module. +module DataFrame.DecisionTree.Types ( + Tree (..), + treeDepth, + TreeConfig (..), + SynthConfig (..), + defaultTreeConfig, + defaultSynthConfig, + ColumnOrdering (..), + orderable, + defaultColumnOrdering, + withOrdFrom, + CarePoint (..), + Direction (..), + ttrace, +) where + +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr (..)) +import qualified DataFrame.LinearSolver as LS + +import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.Map.Strict as M +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import Data.Word (Word16, Word32, Word64, Word8) +import qualified Debug.Trace as Trace +import System.Environment (lookupEnv) +import System.IO.Unsafe (unsafePerformIO) +import Type.Reflection (SomeTypeRep (..), typeRep) + +-- | A fitted tree: a leaf value, or an internal node testing a boolean +-- expression with @True@ routing left. +data Tree a + = Leaf !a + | Branch !(Expr Bool) !(Tree a) !(Tree a) + deriving (Show) + +treeDepth :: Tree a -> Int +treeDepth (Leaf _) = 0 +treeDepth (Branch _ l r) = 1 + max (treeDepth l) (treeDepth r) + +-- | A row the parent node must route to a specific child for the subtrees to +-- classify it correctly (the TAO objective is the count of misroutes). +data CarePoint = CarePoint + { cpIndex :: !Int + , cpCorrectDir :: !Direction + } + deriving (Eq, Show) + +data Direction = GoLeft | GoRight + deriving (Eq, Show) + +data TreeConfig = TreeConfig + { maxTreeDepth :: Int + , minSamplesSplit :: Int + , minLeafSize :: Int + , percentiles :: [Int] + , expressionPairs :: Int + , synthConfig :: SynthConfig + , taoIterations :: Int + , taoConvergenceTol :: Double + , columnOrdering :: ColumnOrdering + , useLinearSolver :: Bool + , linearSolverConfig :: LS.SolverConfig + , minCarePointsForLinear :: Int + , pureReplacementLinear :: Bool + } + +data SynthConfig = SynthConfig + { maxExprDepth :: Int + , boolExpansion :: Int + , disallowedCombinations :: [(T.Text, T.Text)] + , complexityPenalty :: Double + , enableStringOps :: Bool + , enableCrossCols :: Bool + , enableArithOps :: Bool + , maxCategoricalSubsetCardinality :: Int + , perColumnQuota :: Maybe Int + } + deriving (Eq, Show) + +defaultSynthConfig :: SynthConfig +defaultSynthConfig = + SynthConfig + { maxExprDepth = 2 + , boolExpansion = 2 + , disallowedCombinations = [] + , complexityPenalty = 0.05 + , enableStringOps = True + , enableCrossCols = True + , enableArithOps = True + , maxCategoricalSubsetCardinality = 4 + , perColumnQuota = Just 3 + } + +defaultTreeConfig :: TreeConfig +defaultTreeConfig = + TreeConfig + { maxTreeDepth = 4 + , minSamplesSplit = 5 + , minLeafSize = 1 + , percentiles = [0, 10 .. 100] + , expressionPairs = 10 + , synthConfig = defaultSynthConfig + , taoIterations = 10 + , taoConvergenceTol = 1e-6 + , columnOrdering = defaultColumnOrdering + , useLinearSolver = True + , linearSolverConfig = LS.defaultSolverConfig + , minCarePointsForLinear = 10 + , pureReplacementLinear = False + } + +-- | Which column types support ordering for splits. Register a type with +-- 'orderable' and combine with @<>@. +newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict) + +instance Semigroup ColumnOrdering where + ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b) + +instance Monoid ColumnOrdering where + mempty = ColumnOrdering M.empty + +-- | Register a type as orderable for decision-tree splits. +orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering +orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a))) + +-- | All standard numeric, text, and primitive types. +defaultColumnOrdering :: ColumnOrdering +defaultColumnOrdering = mconcat (numericOrderings ++ otherOrderings) + +numericOrderings :: [ColumnOrdering] +numericOrderings = + [ orderable @Int + , orderable @Int8 + , orderable @Int16 + , orderable @Int32 + , orderable @Int64 + , orderable @Word + , orderable @Word8 + , orderable @Word16 + , orderable @Word32 + , orderable @Word64 + , orderable @Integer + , orderable @Double + , orderable @Float + ] + +otherOrderings :: [ColumnOrdering] +otherOrderings = + [orderable @Bool, orderable @Char, orderable @T.Text, orderable @String] + +-- | Existential @Ord@ dictionary keyed by type representation. +data OrdDict where + OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict + +-- | Run @k@ with the @Ord a@ instance recovered from the ordering registry, +-- or 'Nothing' when @a@ is not registered. +withOrdFrom :: + forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r +withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of + Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Just k + Nothing -> Nothing + Nothing -> Nothing + +{-# NOINLINE taoTraceEnabled #-} +taoTraceEnabled :: Bool +taoTraceEnabled = unsafePerformIO (fmap (== Just "1") (lookupEnv "TAO_TRACE")) + +-- | Emit a trace line when @TAO_TRACE=1@; a no-op otherwise. +ttrace :: String -> a -> a +ttrace msg x + | taoTraceEnabled = Trace.trace ("[TAO] " ++ msg) x + | otherwise = x diff --git a/src/DataFrame/LinearSolver.hs b/src/DataFrame/LinearSolver.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/LinearSolver.hs @@ -0,0 +1,430 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | L1-regularized logistic regression used as the per-node split solver in +'DataFrame.DecisionTree'. Produces a sparse oblique hyperplane that can be +compiled to an 'Expr Bool' over numeric columns. +-} +module DataFrame.LinearSolver ( + -- * Model + LinearModel (..), + + -- * Configuration + SolverConfig (..), + defaultSolverConfig, + + -- * Solver + fitL1Logistic, + + -- * Expr conversion + modelToExpr, + + -- * Internals (exposed for testing) + standardize, + softThreshold, + sigmoid, + dotProduct, +) where + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Operators ((.*.), (.+.), (.>.)) + +import Control.Monad.ST (ST, runST) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +{- | A fitted linear classifier: predicts the positive class when +@sum (weights .* features) + intercept > 0@. Weights of exactly @0@ mark +features dropped by the L1 penalty (filtered out by 'modelToExpr'). +-} +data LinearModel = LinearModel + { lmWeights :: !(VU.Vector Double) + , lmIntercept :: !Double + , lmFeatureNames :: !(V.Vector T.Text) + } + deriving (Eq, Show) + +-- | Hyper-parameters for the FISTA solver. +data SolverConfig = SolverConfig + { scL1Lambda :: !Double + -- ^ Strength of the L1 penalty on weights (intercept is not regularized). + , scL2Lambda :: !Double + {- ^ Strength of the L2 penalty @(λ₂/2)·|w|²@ (Elastic Net; Zou & Hastie + 2005). Combined with @scL1Lambda@ this gives the standard elastic-net + objective @λ₁·|w|₁ + (λ₂/2)·|w|²@. At @scL2Lambda = 0@ the solver + reduces to pure L1 (the original behaviour). The Friedman/Hastie/ + Tibshirani 2010 glmnet proximal step under step size @1/L@ is + @softThreshold(z, λ₁/L) / (1 + λ₂/L)@ with @L = (d+1)/4 + λ₂@. + -} + , scMaxIter :: !Int + -- ^ Maximum number of FISTA iterations. + , scTol :: !Double + -- ^ Convergence tolerance on the weight delta (L-inf norm). + , scSampleWeights :: !(Maybe (VU.Vector Double)) + {- ^ Optional per-row sample weights, length @n@. @Nothing@ is uniform + weight 1 (legacy behaviour, A1-A18 path). The 1/N gradient + normalisation is preserved by convention: weights should have mean + 1 (i.e. @Σ w_i = N@) so the existing Lipschitz bound stays valid. + See 'fitLinearCandidate' in 'DataFrame.DecisionTree' for the + class-balanced construction @w_i = N / (2 · N_class(label_i))@. + -} + } + deriving (Eq, Show) + +defaultSolverConfig :: SolverConfig +defaultSolverConfig = + SolverConfig + { scL1Lambda = 0.005 + , scL2Lambda = 0.005 + , scMaxIter = 200 + , scTol = 1.0e-4 + , scSampleWeights = Nothing + } + +{- | Fit L1-regularized binary logistic regression by FISTA. Rows are feature +vectors of equal length; labels are in @{\-1,+1}@. Features are standardized +internally and weights de-standardized, so the model applies to raw values. +-} +fitL1Logistic :: + SolverConfig -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + V.Vector T.Text -> + LinearModel +{-# INLINEABLE fitL1Logistic #-} +fitL1Logistic cfg rows labels featureNames + | n == 0 || d == 0 = zeroModel + | otherwise = + let (!means, !stds, !variances) = columnStats rows + !keep = keptIndices variances + in if VU.null keep + then zeroModel + else + let !meansKept = gatherBy keep means + !stdsKept = gatherBy keep stds + !xKept = V.map (standardizeRowKept keep means stds) rows + -- Elastic-Net Lipschitz: standard logistic bound + -- @(d+1)/4@ plus the L2 part's Hessian-norm + -- contribution @λ₂·I@ (operator norm @λ₂@). + !lipschitz = + fromIntegral (VU.length keep + 1) / 4 + + scL2Lambda cfg + (!wStdKept, !bStd) = + fistaLoop + (scL1Lambda cfg) + (scL2Lambda cfg) + lipschitz + (scMaxIter cfg) + (scTol cfg) + (scSampleWeights cfg) + xKept + labels + (VU.replicate (VU.length keep) 0) + 0 + !wRawKept = VU.zipWith (/) wStdKept stdsKept + !bRaw = bStd - VU.sum (VU.zipWith (*) wRawKept meansKept) + in LinearModel (expandWeights d keep wRawKept) bRaw featureNames + where + !n = V.length rows + !d = V.length featureNames + zeroModel = LinearModel (VU.replicate d 0) 0 featureNames + +{- | Indices of columns whose variance clears the near-constant threshold. +Columns below it are dropped before fitting; their weight ends up @0@. +-} +keptIndices :: VU.Vector Double -> VU.Vector Int +keptIndices variances = + VU.fromList + [ j + | j <- [0 .. VU.length variances - 1] + , VU.unsafeIndex variances j >= 1.0e-12 + ] + +{- | Gather the entries of @v@ at @idxs@, preserving order. unsafeIndex is +safe: every index in @idxs@ is in range by construction. +-} +gatherBy :: VU.Vector Int -> VU.Vector Double -> VU.Vector Double +gatherBy idxs v = VU.map (VU.unsafeIndex v) idxs + +{- | Standardize one row to the kept columns only (subtract column mean, divide +by column std). unsafeIndex is safe: rows share the column layout. +-} +standardizeRowKept :: + VU.Vector Int -> + VU.Vector Double -> + VU.Vector Double -> + VU.Vector Double -> + VU.Vector Double +standardizeRowKept keep means stds row = VU.map standardizeAt keep + where + standardizeAt j = + (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j + +{- | Scatter kept-column weights back into a full-width vector, with @0@ for +the dropped (near-constant) columns. +-} +expandWeights :: Int -> VU.Vector Int -> VU.Vector Double -> VU.Vector Double +expandWeights d keep wKept = VU.create $ do + mv <- VUM.replicate d 0 + VU.iforM_ keep $ \k j -> VUM.unsafeWrite mv j (VU.unsafeIndex wKept k) + pure mv + +{- | Convert a fitted model to an 'Expr Bool' over its feature columns, +dropping zero-weight features. With no non-zero weights it returns the +constant @Lit (intercept > 0)@. +-} +modelToExpr :: LinearModel -> Expr Bool +modelToExpr m = + case nonZero of + [] -> F.lit (b > 0) + (w0, n0) : rest -> score rest (term w0 n0) .>. F.lit (0 :: Double) + where + b = lmIntercept m + nonZero = + [ (w, n) + | (w, n) <- zip (VU.toList (lmWeights m)) (V.toList (lmFeatureNames m)) + , w /= 0 + ] + term w n = F.lit w .*. (Col n :: Expr Double) + score rest first = foldl (\acc (w, n) -> acc .+. term w n) first rest .+. F.lit b + +{- | Per-column @(means, stds, variances)@ of a feature matrix. Cheaper than +'standardize' when only the statistics are needed. unsafeIndex within is +safe: all rows share width @d@. +-} +columnStats :: + V.Vector (VU.Vector Double) -> + (VU.Vector Double, VU.Vector Double, VU.Vector Double) +columnStats x + | V.null x = (VU.empty, VU.empty, VU.empty) + | otherwise = + let !d = VU.length (V.unsafeHead x) + !invN = 1 / fromIntegral (V.length x) + !means = columnMeans d invN x + !variances = columnVariances d invN means x + !stds = VU.map (\v -> if v < 1e-12 then 1 else sqrt v) variances + in (means, stds, variances) + +-- | Mean of each of the @d@ columns; @invN@ is @1 / nRows@. +columnMeans :: Int -> Double -> V.Vector (VU.Vector Double) -> VU.Vector Double +columnMeans d invN x = runST $ do + acc <- VUM.replicate d 0 + V.forM_ x $ \row -> + VU.iforM_ row $ \j v -> VUM.unsafeModify acc (+ v) j + scaleInPlace invN acc + VU.unsafeFreeze acc + +-- | Variance of each of the @d@ columns about the supplied @means@. +columnVariances :: + Int -> + Double -> + VU.Vector Double -> + V.Vector (VU.Vector Double) -> + VU.Vector Double +columnVariances d invN means x = runST $ do + acc <- VUM.replicate d 0 + V.forM_ x $ \row -> + VU.iforM_ row $ \j v -> + let !c = v - VU.unsafeIndex means j in VUM.unsafeModify acc (+ c * c) j + scaleInPlace invN acc + VU.unsafeFreeze acc + +-- | Multiply every element of a mutable vector by @factor@ in place. +scaleInPlace :: Double -> VUM.MVector s Double -> ST s () +scaleInPlace factor mv = go 0 + where + go !j + | j >= VUM.length mv = pure () + | otherwise = VUM.unsafeModify mv (* factor) j >> go (j + 1) + +{- | Standardize each column to zero mean and unit variance, also returning +@(means, stds, variances)@. Near-constant columns get std @1@; callers use +the raw variances to detect and drop them (see 'fitL1Logistic'). +-} +standardize :: + V.Vector (VU.Vector Double) -> + ( V.Vector (VU.Vector Double) + , VU.Vector Double + , VU.Vector Double + , VU.Vector Double + ) +standardize x + | V.null x = (x, VU.empty, VU.empty, VU.empty) + | otherwise = + let (!means, !stds, !variances) = columnStats x + !d = VU.length (V.unsafeHead x) + standardizeRow row = + VU.generate d $ \j -> + (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j + in (V.map standardizeRow x, means, stds, variances) + +{- | Proximal operator for the L1 norm: shrink @v@ toward zero by @lambda@, +clamping at zero. +-} +softThreshold :: Double -> Double -> Double +softThreshold lambda v + | v > lambda = v - lambda + | v < -lambda = v + lambda + | otherwise = 0 + +-- | Numerically stable logistic sigmoid. +sigmoid :: Double -> Double +sigmoid z + | z >= 0 = 1 / (1 + exp (-z)) + | otherwise = let ez = exp z in ez / (1 + ez) + +{- | Dot product of two unboxed vectors. Caller must ensure equal length; +lengths are not checked. +-} +dotProduct :: VU.Vector Double -> VU.Vector Double -> Double +dotProduct u v = VU.sum (VU.zipWith (*) u v) + +{- | Gradient of the average binary logistic loss at @(w, b)@ for labels in +@{\-1,+1}@. Returns @(gradW, gradB)@. + +Sample-weighted variant: when @sampleWeights@ is @Just ws@ the per-row +contribution is multiplied by @ws[i]@. With weights of mean 1 +(i.e. @Σ w_i = N@; the class-balanced convention used by +'fitLinearCandidate'), the @1/N@ normalisation is preserved exactly. +-} +logisticGradient :: + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +logisticGradient sampleWeights features labels w b = (gradW, gradB) + where + !invN = 1 / fromIntegral (V.length features) + !coeffs = rowCoeffs sampleWeights features labels w b invN + !gradW = accumulateGradW (VU.length w) features coeffs + !gradB = VU.sum coeffs + +{- | Per-row loss coefficient. Without sample weights: +@c_i = -y_i * sigmoid(-y_i * margin_i) / N@. With @Just ws@, each row's +contribution is additionally multiplied by @ws[i]@. + +unsafeIndex is safe: @i@ ranges over @[0,n-1]@ and @labels@ / +@sampleWeights@ both have length @n@ by construction. +-} +rowCoeffs :: + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + Double -> + VU.Vector Double +rowCoeffs sampleWeights features labels w b invN = + VU.generate (V.length features) $ \i -> + let !yi = VU.unsafeIndex labels i + !row = V.unsafeIndex features i + !margin = yi * (dotProduct w row + b) + !base = -(yi * sigmoid (-margin) * invN) + in case sampleWeights of + Nothing -> base + Just ws -> base * VU.unsafeIndex ws i + +{- | Accumulate the weight gradient in one pass over every (row, feature) +pair, scattering into a length-@d@ mutable vector. +-} +accumulateGradW :: + Int -> V.Vector (VU.Vector Double) -> VU.Vector Double -> VU.Vector Double +accumulateGradW d features coeffs = runST $ do + mv <- VUM.replicate d 0 + V.iforM_ features $ \i row -> + let !c = VU.unsafeIndex coeffs i + in VU.iforM_ row $ \j v -> VUM.unsafeModify mv (+ c * v) j + VU.unsafeFreeze mv + +{- | Inner FISTA loop over standardized features. Returns the final @(w, b)@; +the caller is responsible for de-standardization. + +@lambda1@ and @lambda2@ are the L1 / L2 penalty strengths; @lp@ is the +Lipschitz constant of the smooth part @(d+1)/4 + λ₂@. The Elastic-Net +proximal step is applied per FHT 2010 glmnet §2.6: +@prox(z) = softThreshold(z, λ₁/lp) / (1 + λ₂/lp)@. +-} +fistaLoop :: + Double -> + Double -> + Double -> + Int -> + Double -> + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +fistaLoop lambda1 lambda2 lp maxIter tol sampleWeights features labels w0 b0 = + go 0 w0 b0 w0 b0 1.0 + where + !shrink = lambda1 / lp + !ridgeDenom = 1 + lambda2 / lp + !stepInv = 1 / lp + proxStep = fistaProxStep sampleWeights features labels shrink ridgeDenom stepInv + go !iter !xWPrev !xBPrev !yW !yB !t + | iter >= maxIter = (xWPrev, xBPrev) + | iter > 0 && delta < tol = (xW, xB) + | otherwise = go (iter + 1) xW xB yWNew yBNew tNew + where + (!xW, !xB) = proxStep yW yB + !delta = if VU.null xW then 0 else deltaInf xWPrev xW + (!yWNew, !yBNew, !tNew) = fistaMomentum t xWPrev xBPrev xW xB + +{- | One fused FISTA prox step: gradient step plus the Elastic-Net +proximal operator (soft-threshold then ridge shrinkage), without +materializing the intermediate trial weights. + +The Elastic-Net prox of @g(w) = λ₁·|w|₁ + (λ₂/2)·|w|²@ at step @1/lp@ is +@softThreshold(z, λ₁/lp) / (1 + λ₂/lp)@ (FHT 2010 glmnet §2.6; Beck & +Teboulle 2009 §4). The intercept is unregularised (no L1 or L2 applied). +-} +fistaProxStep :: + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + Double -> + Double -> + Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +fistaProxStep sampleWeights features labels shrink ridgeDenom stepInv yW yB = + let (gW, gB) = logisticGradient sampleWeights features labels yW yB + !wNew = + VU.zipWith + (\yi gi -> softThreshold shrink (yi - gi * stepInv) / ridgeDenom) + yW + gW + !bNew = yB - gB * stepInv + in (wNew, bNew) + +{- | Nesterov momentum extrapolation: new look-ahead point @(yW, yB)@ and the +updated step size @t@. +-} +fistaMomentum :: + Double -> + VU.Vector Double -> + Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double, Double) +fistaMomentum t xWPrev xBPrev xW xB = + let !tNew = (1 + sqrt (1 + 4 * t * t)) / 2 + !mom = (t - 1) / tNew + !yW = VU.zipWith (\new old -> new + mom * (new - old)) xW xWPrev + !yB = xB + mom * (xB - xBPrev) + in (yW, yB, tNew) + +{- | L-inf norm of the weight delta. unsafeIndex is safe: both vectors share +the same length by construction. +-} +{-# INLINE deltaInf #-} +deltaInf :: VU.Vector Double -> VU.Vector Double -> Double +deltaInf xWPrev = VU.ifoldl' (\acc i x -> max acc (abs (x - VU.unsafeIndex xWPrev i))) 0