dataframe-learn 1.0.1.0 → 1.0.2.0

raw patch · 14 files changed

+2417/−1011 lines, 14 filesdep +paralleldep +vector-algorithmsdep ~dataframe-operationsPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: parallel, vector-algorithms

Dependency ranges changed: dataframe-operations

API changes (from Hackage documentation)

Files

• dataframe-learn.cabal +18/−4
• src/DataFrame/DecisionTree.hs +29/−1007
• src/DataFrame/DecisionTree/Cart.hs +230/−0
• src/DataFrame/DecisionTree/Categorical.hs +281/−0
• src/DataFrame/DecisionTree/CondVec.hs +155/−0
• src/DataFrame/DecisionTree/Fit.hs +174/−0
• src/DataFrame/DecisionTree/Linear.hs +133/−0
• src/DataFrame/DecisionTree/Numeric.hs +220/−0
• src/DataFrame/DecisionTree/Pool.hs +174/−0
• src/DataFrame/DecisionTree/Predict.hs +161/−0
• src/DataFrame/DecisionTree/Prune.hs +59/−0
• src/DataFrame/DecisionTree/Tao.hs +168/−0
• src/DataFrame/DecisionTree/Types.hs +185/−0
• src/DataFrame/LinearSolver.hs +430/−0

@@ -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

@@ -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

@@ -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)

@@ -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)

@@ -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))

@@ -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)

@@ -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+    [] -> "<feat>"++-- | 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)

@@ -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 -> []

@@ -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+    [] -> "<noncol>"+    (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)

@@ -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)

@@ -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

@@ -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)

@@ -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

@@ -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