feat: Add histogram trees to dataframe-learn

Author: mchav

dataframe-learn/dataframe-learn.cabal

@@ -45,6 +45,7 @@ library internal
                         DataFrame.DecisionTree.Pool
                         DataFrame.DecisionTree.Linear
                         DataFrame.DecisionTree.Tao
+                        DataFrame.DecisionTree.Histogram
                         DataFrame.DecisionTree.Fit
                         DataFrame.LinearSolver
                         DataFrame.LinearSolver.Loss
@@ -126,6 +127,7 @@ test-suite learn-internal
     other-modules:      DataFrameApi
                         Cart
                         NullSplits
+                        HistogramTrees
                         DecisionTree
                         TreePruning
                         Worklist

dataframe-learn/src-internal/DataFrame/DecisionTree/Histogram.hs

@@ -0,0 +1,260 @@
+{-# LANGUAGE BangPatterns #-}
+
+{- | Regression trees grown from per-bin histograms, for gradient boosting.
+Each feature is bucketed once per fit; a node's split search then costs
+O(bins) per feature instead of a sweep over its sorted rows. Splits, null
+handling and leaf values follow 'DataFrame.DecisionTree.Regression': with no
+more distinct values than bins, the trees partition the rows identically.
+-}
+module DataFrame.DecisionTree.Histogram (
+    Binned,
+    binFeatures,
+    TreeLimits (..),
+    fitBinnedTree,
+) where
+
+import Control.Monad.ST (runST)
+import Control.Parallel (par, pseq)
+import Control.Parallel.Strategies (evalList, parList, rseq, using)
+import Data.Maybe (maybeToList)
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import Data.Word (Word16)
+
+import DataFrame.DecisionTree.Cart (
+    CartFeature (..),
+    NullSide (..),
+    sortIndicesByValue,
+    splitMidpoint,
+ )
+import DataFrame.DecisionTree.Types (Tree (..))
+import DataFrame.Internal.Expression (Expr)
+
+-- | A feature's bins: rows in bin @b@ satisfy @value <= upper ! b@.
+data Binned = Binned
+    { bnFeature :: !CartFeature
+    , bnBins :: !(VU.Vector Word16)
+    -- ^ Bin per row; null rows hold 'bnNBins'.
+    , bnNBins :: !Int
+    , bnUpper :: !(VU.Vector Double)
+    }
+
+-- | Bin every feature into at most @maxBins@ bins (capped at 65535).
+binFeatures :: Int -> V.Vector CartFeature -> V.Vector Binned
+binFeatures maxBins = V.map (binFeature (max 2 (min 65535 maxBins)))
+
+binFeature :: Int -> CartFeature -> Binned
+binFeature maxBins f = Binned f bins nb upper
+  where
+    vals = cfValues f
+    isNull i = maybe False (VU.! i) (cfNulls f)
+    sorted =
+        VU.map (vals VU.!) (VU.filter (not . isNull) (sortIndicesByValue vals))
+    runs = runLengths sorted
+    distinct = map fst runs
+    edges
+        | length runs <= maxBins = VU.fromList (zipWith splitMidpoint distinct (drop 1 distinct))
+        | otherwise = VU.fromList (cutEdges maxBins (VU.length sorted) runs)
+    nb = VU.length edges + 1
+    -- The last bin's threshold is the largest value, so "every non-null row
+    -- left" stays a valid null-vs-non-null split, as in the exact tree.
+    upper = VU.snoc edges (if VU.null sorted then 0 else VU.last sorted)
+    bins = VU.generate (VU.length vals) binOf
+    binOf i
+        | isNull i = fromIntegral nb
+        | otherwise = fromIntegral (lowerBound edges (vals VU.! i))
+
+-- | Cut points giving bins of roughly equal row counts, never inside a run.
+cutEdges :: Int -> Int -> [(Double, Int)] -> [Double]
+cutEdges maxBins n = go 0 0
+  where
+    perBin = fromIntegral n / fromIntegral maxBins :: Double
+    go :: Int -> Int -> [(Double, Int)] -> [Double]
+    go !nCuts !acc ((v, c) : rest@((v', _) : _))
+        | nCuts >= maxBins - 1 = []
+        | fromIntegral (acc + c) >= perBin = splitMidpoint v v' : go (nCuts + 1) 0 rest
+        | otherwise = go nCuts (acc + c) rest
+    go _ _ _ = []
+
+runLengths :: VU.Vector Double -> [(Double, Int)]
+runLengths v = go 0
+  where
+    n = VU.length v
+    go i
+        | i >= n = []
+        | otherwise =
+            let x = v VU.! i
+                j = until (\k -> k >= n || v VU.! k /= x) (+ 1) i
+             in (x, j - i) : go j
+
+-- | Number of edges below @x@: the bin @x@ falls in.
+lowerBound :: VU.Vector Double -> Double -> Int
+lowerBound edges x = go 0 (VU.length edges)
+  where
+    go lo hi
+        | lo >= hi = lo
+        | VU.unsafeIndex edges m < x = go (m + 1) hi
+        | otherwise = go lo m
+      where
+        m = (lo + hi) `div` 2
+
+-- | The stopping rules of 'DataFrame.DecisionTree.Regression.RegTreeConfig'.
+data TreeLimits = TreeLimits
+    { tlMaxDepth :: !Int
+    , tlMinSamplesSplit :: !Int
+    , tlMinLeafSize :: !Int
+    , tlMinImpurityDecrease :: !Double
+    }
+
+{- | Per-bin Σw, Σwy, Σwy² and row count, four slots per bin, null bin last.
+Built strictly: left lazy, the histograms of a deep tree pile up as thunks.
+-}
+type Hist = VU.Vector Double
+
+buildHist :: VU.Vector Double -> VU.Vector Double -> VU.Vector Int -> Binned -> Hist
+buildHist w y idxs bn = runST $ do
+    m <- VUM.replicate (4 * (bnNBins bn + 1)) 0
+    VU.forM_ idxs $ \i -> do
+        let s = 4 * fromIntegral (VU.unsafeIndex (bnBins bn) i)
+            wi = VU.unsafeIndex w i
+            yi = VU.unsafeIndex y i
+        VUM.unsafeModify m (+ wi) s
+        VUM.unsafeModify m (+ wi * yi) (s + 1)
+        VUM.unsafeModify m (+ wi * yi * yi) (s + 2)
+        VUM.unsafeModify m (+ 1) (s + 3)
+    VU.unsafeFreeze m
+
+data Node
+    = NLeaf !Double !(VU.Vector Int)
+    | NBranch !(Expr Bool) !Node !Node
+
+{- | Fit a weighted-SSE regression tree on binned features; returns the tree
+and its in-sample predictions.
+-}
+fitBinnedTree ::
+    TreeLimits ->
+    V.Vector Binned ->
+    VU.Vector Double ->
+    Maybe (VU.Vector Double) ->
+    (Tree Double, VU.Vector Double)
+fitBinnedTree lim binned y mw = (toTree root, inSample)
+  where
+    n = VU.length y
+    w = maybe (VU.replicate n 1) id mw
+    allIdx = VU.enumFromN 0 n
+    root = node 0 allIdx (hists allIdx)
+    inSample = VU.update (VU.replicate n 0) (VU.concat (leafRows root))
+    leafRows (NLeaf v idxs) = [VU.map (\i -> (i, v)) idxs]
+    leafRows (NBranch _ l r) = leafRows l ++ leafRows r
+
+    -- Feature-parallel only where a node is big enough to pay for the sparks.
+    hists idxs =
+        V.fromList (map (buildHist w y idxs) (V.toList binned) `using` strategy)
+      where
+        strategy = if VU.length idxs >= 20000 then parList rseq else evalList rseq
+
+    node depth idxs hs
+        | depth >= tlMaxDepth lim || VU.length idxs < tlMinSamplesSplit lim || V.null hs = leaf
+        | otherwise = maybe leaf split (bestSplit lim binned hs)
+      where
+        Totals tw tsy _ _ = histTotals (V.head hs)
+        leaf = NLeaf (if tw == 0 then 0 else tsy / tw) idxs
+        split (fj, b, nullLeft) =
+            forceNode l `par` (forceNode r `pseq` NBranch cond l r)
+          where
+            bn = binned V.! fj
+            goesLeft i =
+                let k = fromIntegral (VU.unsafeIndex (bnBins bn) i)
+                 in if k == bnNBins bn then nullLeft else k <= b
+            lefts = VU.filter goesLeft idxs
+            rights = VU.filter (not . goesLeft) idxs
+            -- Histogram the smaller child; the larger is parent minus it.
+            smallIsLeft = VU.length lefts <= VU.length rights
+            smallH = hists (if smallIsLeft then lefts else rights)
+            bigH = V.zipWith (VU.zipWith (-)) hs smallH
+            (lh, rh) = if smallIsLeft then (smallH, bigH) else (bigH, smallH)
+            l = node (depth + 1) lefts lh
+            r = node (depth + 1) rights rh
+            side = if nullLeft then NullsLeft else NullsRight
+            cond = cfSplit (bnFeature bn) (bnUpper bn VU.! b) side
+
+data Totals = Totals !Double !Double !Double !Double
+
+histTotals :: Hist -> Totals
+histTotals h = go 0 0 0 0 0
+  where
+    slots = VU.length h `div` 4
+    go !k !a !b !c !d
+        | k >= slots = Totals a b c d
+        | otherwise =
+            go
+                (k + 1)
+                (a + h VU.! (4 * k))
+                (b + h VU.! (4 * k + 1))
+                (c + h VU.! (4 * k + 2))
+                (d + h VU.! (4 * k + 3))
+
+{- | Best @(feature, bin, nulls left)@ by SSE reduction. Each bin boundary is
+scored with the node's null rows on the right and then on the left; ties keep
+the earliest candidate, as in the exact sweep.
+-}
+bestSplit :: TreeLimits -> V.Vector Binned -> V.Vector Hist -> Maybe (Int, Int, Bool)
+bestSplit lim binned hs
+    | null candidates = Nothing
+    | red > 0 && red >= tlMinImpurityDecrease lim = Just sp
+    | otherwise = Nothing
+  where
+    Totals totW totSY totSY2 totC = histTotals (V.head hs)
+    nNode = round totC :: Int
+    nodeSSE = sse totSY totSY2 totW
+    candidates = concat (zipWith featBest [0 ..] (V.toList hs))
+    (red, sp) = maximumByFst candidates
+    featBest fj h = [(r, (fj, b, nl)) | (b, nl, r) <- maybeToList (scan 0 0 0 0 0 Nothing)]
+      where
+        nb = bnNBins (binned V.! fj)
+        nW = h VU.! (4 * nb)
+        nSY = h VU.! (4 * nb + 1)
+        nSY2 = h VU.! (4 * nb + 2)
+        nC = round (h VU.! (4 * nb + 3)) :: Int
+        hasNulls = nC > 0
+        lastB = if hasNulls then nb - 1 else nb - 2
+        score nl wl syl syl2
+            | nl >= tlMinLeafSize lim && nNode - nl >= tlMinLeafSize lim && wl > 0 && wr > 0 =
+                Just (nodeSSE - (sse syl syl2 wl + sse (totSY - syl) (totSY2 - syl2) wr))
+            | otherwise = Nothing
+          where
+            wr = totW - wl
+        scan :: Int -> Double -> Double -> Double -> Int -> Maybe (Int, Bool, Double) -> Maybe (Int, Bool, Double)
+        scan !b !wl !syl !syl2 !cl best
+            | b > lastB = best
+            | otherwise = scan (b + 1) wl' syl' syl2' cl' best'
+          where
+            wl' = wl + h VU.! (4 * b)
+            syl' = syl + h VU.! (4 * b + 1)
+            syl2' = syl2 + h VU.! (4 * b + 2)
+            cl' = cl + round (h VU.! (4 * b + 3))
+            right = (,) False <$> score cl' wl' syl' syl2'
+            left
+                | hasNulls = (,) True <$> score (cl' + nC) (wl' + nW) (syl' + nSY) (syl2' + nSY2)
+                | otherwise = Nothing
+            consider bst (dir, r)
+                | maybe True (\(_, _, rb) -> r > rb) bst = Just (b, dir, r)
+                | otherwise = bst
+            best' = foldl consider best (maybeToList right ++ maybeToList left)
+
+-- | Weighted SSE of a node from its Σy, Σy², and total weight.
+sse :: Double -> Double -> Double -> Double
+sse sumY sumSq wt = sumSq - (if wt == 0 then 0 else sumY * sumY / wt)
+
+maximumByFst :: (Ord a) => [(a, b)] -> (a, b)
+maximumByFst = foldr1 (\x@(a, _) y@(b, _) -> if a >= b then x else y)
+
+toTree :: Node -> Tree Double
+toTree (NLeaf v _) = Leaf v
+toTree (NBranch c l r) = Branch c (toTree l) (toTree r)
+
+-- | Force a subtree so the sibling's spark has real work (cf. 'Regression').
+forceNode :: Node -> ()
+forceNode (NLeaf v _) = v `seq` ()
+forceNode (NBranch _ l r) = forceNode l `seq` forceNode r

dataframe-learn/src/DataFrame/Boosting/GBM.hs

@@ -9,7 +9,9 @@
 {- | Gradient boosting of regression trees (Friedman). Trees are fitted to the
 negative gradient of the loss each round and accumulated with a shrinkage
 factor; squared error gives regression, logistic deviance gives binary
-classification. 'predict' is the additive score; 'gbProbaExpr' /
+classification. Split search runs on per-feature histograms (see
+'gbMaxBins'), so a tree costs O(rows) rather than a sweep per node.
+'predict' is the additive score; 'gbProbaExpr' /
 'gbDecisionExpr' give the classification probability / decision.
 -}
 module DataFrame.Boosting.GBM (
@@ -32,7 +34,7 @@ import DataFrame.Errors (DataFrameException (..))
 
 import DataFrame.DecisionTree.Cart (cartFeatures)
 import DataFrame.DecisionTree.Fit (treeToExpr)
-import DataFrame.DecisionTree.Regression (RegFit (..), RegTreeConfig (..), fitRegTree)
+import DataFrame.DecisionTree.Histogram (TreeLimits (..), binFeatures, fitBinnedTree)
 import DataFrame.DecisionTree.Types (Tree)
 import DataFrame.Expression.Operators ((.*.), (.+.), (.>.))
 import DataFrame.Featurize.Internal (targetDoubles)
@@ -50,6 +52,9 @@ data GBConfig = GBConfig
     , gbNEstimators :: !Int
     , gbLearningRate :: !Double
     , gbMaxDepth :: !Int
+    , gbMaxBins :: !Int
+    -- ^ Bins per feature for split finding. Features with at most this many
+    -- distinct values split exactly as on the raw values.
     , gbSeed :: !Int
     }
     deriving (Eq, Show)
@@ -61,6 +66,7 @@ defaultGBConfig =
         , gbNEstimators = 100
         , gbLearningRate = 0.1
         , gbMaxDepth = 3
+        , gbMaxBins = 1024
         , gbSeed = 0
         }
 
@@ -97,16 +103,16 @@ fitGBM cfg target@(Col name) df =
         (VU.fromList (reverse scores))
         usage
   where
-    feats = V.fromList (cartFeatures name df)
+    binned = binFeatures (gbMaxBins cfg) (V.fromList (cartFeatures name df))
     y = targetDoubles target df
     n = VU.length y
     lr = gbLearningRate cfg
-    rtCfg =
-        RegTreeConfig
-            { rtMaxDepth = gbMaxDepth cfg
-            , rtMinSamplesSplit = 2
-            , rtMinLeafSize = 1
-            , rtMinImpurityDecrease = 0.0
+    limits =
+        TreeLimits
+            { tlMaxDepth = gbMaxDepth cfg
+            , tlMinSamplesSplit = 2
+            , tlMinLeafSize = 1
+            , tlMinImpurityDecrease = 0.0
             }
     f0 = case gbLoss cfg of
         SquaredError -> VU.sum y / fromIntegral (max 1 n)
@@ -118,7 +124,7 @@ fitGBM cfg target@(Col name) df =
         | m >= gbNEstimators cfg = (ts, ss, usageAcc)
         | otherwise =
             let (target', weights) = newtonStep (gbLoss cfg) y fScores
-                RegFit tree pred = fitRegTree rtCfg feats target' weights
+                (tree, pred) = fitBinnedTree limits binned target' weights
                 fScores' = VU.zipWith (\f p -> f + lr * p) fScores pred
                 score = lossValue (gbLoss cfg) y fScores'
                 usage' = foldr (\c -> M.insertWith (+) c 1) usageAcc (treeColumns tree)