dataframe-learn (empty) → 1.0.0.0
raw patch · 4 files changed
+1550/−0 lines, 4 filesdep +basedep +containersdep +dataframe-core
Dependencies added: base, containers, dataframe-core, dataframe-operations, text, vector
Files
- LICENSE +20/−0
- dataframe-learn.cabal +39/−0
- src/DataFrame/DecisionTree.hs +1007/−0
- src/DataFrame/Synthesis.hs +484/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2026 Michael Chavinda++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ dataframe-learn.cabal view
@@ -0,0 +1,39 @@+cabal-version: 2.4+name: dataframe-learn+version: 1.0.0.0++synopsis: Decision trees and feature synthesis for the dataframe ecosystem.+description:+ @DataFrame.DecisionTree@ — decision-tree training on DataFrames.+ @DataFrame.Synthesis@ — feature synthesis. Built on top of+ @dataframe-operations@.++bug-reports: https://github.com/mchav/dataframe/issues+license: MIT+license-file: LICENSE+author: Michael Chavinda+maintainer: mschavinda@gmail.com+copyright: (c) 2024-2025 Michael Chavinda+category: Data+tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2++common warnings+ ghc-options:+ -Wincomplete-patterns+ -Wincomplete-uni-patterns+ -Wunused-imports+ -Wunused-local-binds++library+ import: warnings+ exposed-modules:+ DataFrame.DecisionTree+ DataFrame.Synthesis+ build-depends: base >= 4 && < 5,+ containers >= 0.6.7 && < 0.9,+ dataframe-core ^>= 1.0,+ dataframe-operations ^>= 1.0,+ text >= 2.0 && < 3,+ vector ^>= 0.13+ hs-source-dirs: src+ default-language: Haskell2010
+ src/DataFrame/DecisionTree.hs view
@@ -0,0 +1,1007 @@+{-# 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)
+ src/DataFrame/Synthesis.hs view
@@ -0,0 +1,484 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Synthesis where++import qualified DataFrame.Functions as F+import DataFrame.Internal.Column+import DataFrame.Internal.DataFrame (+ DataFrame (..),+ columnNames,+ )+import DataFrame.Internal.Expression (+ Expr (..),+ eSize,+ eqExpr,+ )+import DataFrame.Internal.Interpreter (interpret)+import DataFrame.Internal.Statistics+import DataFrame.Operations.Core (columnAsDoubleVector)+import qualified DataFrame.Operations.Statistics as Stats+import DataFrame.Operations.Subset (exclude)++import Control.Exception (throw)+import Data.Function+import qualified Data.List as L+import qualified Data.Map as M+import Data.Maybe (listToMaybe)+import qualified Data.Text as T+import Data.Type.Equality+import qualified Data.Vector.Unboxed as VU+import qualified DataFrame.Operations.Core as D+import DataFrame.Operators+import Debug.Trace (trace)+import Type.Reflection (typeRep)++generateConditions ::+ TypedColumn Double -> [Expr Bool] -> [Expr Double] -> DataFrame -> [Expr Bool]+generateConditions labels conds ps df =+ let+ newConds =+ [ p .<= q+ | p <- filter (not . isLiteral) ps+ , q <- ps+ , Prelude.not (eqExpr p q)+ ]+ ++ [ F.not p+ | p <- conds+ ]+ expandedConds =+ conds+ ++ newConds+ ++ [p .&& q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)]+ ++ [p .|| q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)]+ in+ pickTopNBool df labels (deduplicate df expandedConds)++generatePrograms ::+ Bool ->+ [Expr Bool] ->+ [Expr Double] ->+ [Expr Double] ->+ [Expr Double] ->+ [Expr Double]+generatePrograms _ _ vars' constants [] = vars' ++ constants+generatePrograms includeConds conds vars constants ps =+ let+ existingPrograms = ps ++ vars ++ constants+ in+ existingPrograms+ ++ [ transform p+ | p <- ps ++ vars+ , Prelude.not (isConditional p)+ , transform <-+ [ sqrt+ , abs+ , log . (+ Lit 1)+ , exp+ , sin+ , cos+ , F.relu+ , signum+ ]+ ]+ ++ [ F.pow p i+ | p <- existingPrograms+ , Prelude.not (isConditional p)+ , i <- [2 .. 6]+ ]+ ++ [ p + q+ | (i, p) <- zip [(0 :: Int) ..] existingPrograms+ , (j, q) <- zip [(0 :: Int) ..] existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , i >= j+ ]+ ++ [ p - q+ | (i, p) <- zip [(0 :: Int) ..] existingPrograms+ , (j, q) <- zip [(0 :: Int) ..] existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , i /= j+ ]+ ++ ( if includeConds+ then+ [ F.min p q+ | (i, p) <- zip [(0 :: Int) ..] existingPrograms+ , (j, q) <- zip [(0 :: Int) ..] existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , Prelude.not (eqExpr p q)+ , i > j+ ]+ ++ [ F.max p q+ | (i, p) <- zip [(0 :: Int) ..] existingPrograms+ , (j, q) <- zip [(0 :: Int) ..] existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , Prelude.not (eqExpr p q)+ , i > j+ ]+ ++ [ F.ifThenElse cond r s+ | cond <- conds+ , r <- existingPrograms+ , s <- existingPrograms+ , Prelude.not (isConditional r || isConditional s)+ , Prelude.not (eqExpr r s)+ ]+ else []+ )+ ++ [ p * q+ | (i, p) <- zip [(0 :: Int) ..] existingPrograms+ , (j, q) <- zip [(0 :: Int) ..] existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , i >= j+ ]+ ++ [ p / q+ | p <- existingPrograms+ , q <- existingPrograms+ , Prelude.not (isLiteral p && isLiteral q)+ , Prelude.not (isConditional p || isConditional q)+ , Prelude.not (eqExpr p q)+ ]++isLiteral :: Expr a -> Bool+isLiteral (Lit _) = True+isLiteral _ = False++isConditional :: Expr a -> Bool+isConditional (If{}) = True+isConditional _ = False++deduplicate ::+ forall a.+ (Columnable a) =>+ DataFrame ->+ [Expr a] ->+ [(Expr a, TypedColumn a)]+deduplicate df = go [] . L.nubBy eqExpr . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2))+ where+ go _ [] = []+ go seen (x : xs)+ | hasInvalid = go seen xs+ | res `elem` seen = go seen xs+ | otherwise = (x, res) : go (res : seen) xs+ where+ res = case interpret @a df x of+ Left e -> throw e+ Right v -> v+ hasInvalid = case res of+ (TColumn (UnboxedColumn _ (column :: VU.Vector b))) -> case testEquality (typeRep @Double) (typeRep @b) of+ Just Refl -> VU.any (\n -> isNaN n || isInfinite n) column+ Nothing -> False+ _ -> False++-- | Checks if two programs generate the same outputs given all the same inputs.+equivalent :: DataFrame -> Expr Double -> Expr Double -> Bool+equivalent df p1 p2 = case (==) <$> interpret df p1 <*> interpret df p2 of+ Left e -> throw e+ Right v -> v++synthesizeFeatureExpr ::+ -- | Target expression+ T.Text ->+ BeamConfig ->+ DataFrame ->+ Either String (Expr Double)+synthesizeFeatureExpr target cfg df =+ let+ df' = exclude [target] df+ t = case interpret df (Col target) of+ Left e -> throw e+ Right v -> v+ in+ case beamSearch+ df'+ cfg+ t+ (percentiles df')+ []+ [] of+ Nothing -> Left "No programs found"+ Just p -> Right p++f1FromBinary :: VU.Vector Double -> VU.Vector Double -> Maybe Double+f1FromBinary trues preds =+ let (!tp, !fp, !fn) =+ VU.foldl' step (0 :: Int, 0 :: Int, 0 :: Int) $+ VU.zip (VU.map (> 0) preds) (VU.map (> 0) trues)+ in f1FromCounts tp fp fn+ where+ step (!tp, !fp, !fn) (!p, !t) =+ case (p, t) of+ (True, True) -> (tp + 1, fp, fn)+ (True, False) -> (tp, fp + 1, fn)+ (False, True) -> (tp, fp, fn + 1)+ (False, False) -> (tp, fp, fn)++f1FromCounts :: Int -> Int -> Int -> Maybe Double+f1FromCounts tp fp fn =+ let tp' = fromIntegral tp+ fp' = fromIntegral fp+ fn' = fromIntegral fn+ precision = if tp' + fp' == 0 then 0 else tp' / (tp' + fp')+ recall = if tp' + fn' == 0 then 0 else tp' / (tp' + fn')+ in if precision + recall == 0+ then Nothing+ else Just (2 * precision * recall / (precision + recall))++fitClassifier ::+ -- | Target expression+ T.Text ->+ -- | Depth of search (Roughly, how many terms in the final expression)+ Int ->+ -- | Beam size - the number of candidate expressions to consider at a time.+ Int ->+ DataFrame ->+ Either String (Expr Int)+fitClassifier target d b df =+ let+ df' = exclude [target] df+ t = case interpret df (Col target) of+ Left e -> throw e+ Right v -> v+ in+ case beamSearch+ df'+ (BeamConfig d b F1 True)+ t+ (percentiles df' ++ [Lit 1, Lit 0, Lit (-1)])+ []+ [] of+ Nothing -> Left "No programs found"+ Just p -> Right (F.ifThenElse (p .> (0 :: Expr Double)) 1 0)++percentiles :: DataFrame -> [Expr Double]+percentiles df =+ let+ doubleColumns =+ map+ (either throw id . ((`columnAsDoubleVector` df) . Col @Double))+ (columnNames df)+ in+ concatMap+ (\c -> map (Lit . roundTo2SigDigits . (`percentile'` c)) [1, 25, 75, 99])+ doubleColumns+ ++ map (Lit . roundTo2SigDigits . variance') doubleColumns+ ++ map (Lit . roundTo2SigDigits . sqrt . variance') doubleColumns++roundToSigDigits :: Int -> Double -> Double+roundToSigDigits n x+ | x == 0 = 0+ | otherwise =+ let magnitude = floor (logBase 10 (abs x))+ scale = 10 ** fromIntegral (n - 1 - magnitude)+ in fromIntegral (round (x * scale) :: Int) / scale++roundTo2SigDigits :: Double -> Double+roundTo2SigDigits = roundToSigDigits 2++fitRegression ::+ -- | Target expression+ T.Text ->+ -- | Depth of search (Roughly, how many terms in the final expression)+ Int ->+ -- | Beam size - the number of candidate expressions to consider at a time.+ Int ->+ DataFrame ->+ Either String (Expr Double)+fitRegression target d b df =+ let+ df' = exclude [target] df+ targetMean = Stats.mean (Col @Double target) df+ t = case interpret df (Col target) of+ Left e -> throw e+ Right v -> v+ cfg = BeamConfig d b MeanSquaredError True+ constants =+ percentiles df'+ ++ [Lit targetMean]+ ++ [ F.pow p i+ | i <- [1 .. 6]+ , p <- [Lit 10, Lit 1, Lit 0.1]+ ]+ in+ case beamSearch df' cfg t constants [] [] of+ Nothing -> Left "No programs found"+ Just p -> Right p++data LossFunction+ = PearsonCorrelation+ | MutualInformation+ | MeanSquaredError+ | F1++getLossFunction ::+ LossFunction -> (VU.Vector Double -> VU.Vector Double -> Maybe Double)+getLossFunction f = case f of+ MutualInformation ->+ ( \l r ->+ mutualInformationBinned+ (Prelude.max 10 (ceiling (sqrt (fromIntegral (VU.length l) :: Double))))+ l+ r+ )+ PearsonCorrelation -> (\l r -> (^ (2 :: Int)) <$> correlation' l r)+ MeanSquaredError -> (\l r -> fmap negate (meanSquaredError l r))+ F1 -> f1FromBinary++data BeamConfig = BeamConfig+ { searchDepth :: Int+ , beamLength :: Int+ , lossFunction :: LossFunction+ , includeConditionals :: Bool+ }++defaultBeamConfig :: BeamConfig+defaultBeamConfig = BeamConfig 2 100 PearsonCorrelation False++beamSearch ::+ DataFrame ->+ -- | Parameters of the beam search.+ BeamConfig ->+ -- | Examples+ TypedColumn Double ->+ -- | Constants+ [Expr Double] ->+ -- | Conditions+ [Expr Bool] ->+ -- | Programs+ [Expr Double] ->+ Maybe (Expr Double)+beamSearch df cfg outputs constants conds programs+ | searchDepth cfg == 0 = case ps of+ [] -> Nothing+ (x : _) -> Just x+ | otherwise =+ beamSearch+ df+ (cfg{searchDepth = searchDepth cfg - 1})+ outputs+ constants+ conditions+ (generatePrograms (includeConditionals cfg) conditions vars constants ps)+ where+ vars = map Col names+ conditions = generateConditions outputs conds (vars ++ constants) df+ ps = pickTopN df outputs cfg $ deduplicate df programs+ names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df++pickTopN ::+ DataFrame ->+ TypedColumn Double ->+ BeamConfig ->+ [(Expr Double, TypedColumn a)] ->+ [Expr Double]+pickTopN _ _ _ [] = []+pickTopN df (TColumn column) cfg ps =+ let+ l = case toVector @Double @VU.Vector column of+ Left e -> throw e+ Right v -> v+ ordered =+ Prelude.take+ (beamLength cfg)+ ( map fst $+ L.sortBy+ ( \(_, c2) (_, c1) ->+ if maybe False isInfinite c1+ || maybe False isInfinite c2+ || maybe False isNaN c1+ || maybe False isNaN c2+ then LT+ else compare c1 c2+ )+ ( map+ (\(e, res) -> (e, getLossFunction (lossFunction cfg) l (asDoubleVector res)))+ ps+ )+ )+ asDoubleVector c =+ let+ (TColumn col') = c+ in+ case toVector @Double @VU.Vector col' of+ Left e -> throw e+ Right v -> VU.convert v+ interpretDoubleVector e' =+ let+ (TColumn col') = case interpret df e' of+ Left err -> throw err+ Right v -> v+ in+ case toVector @Double @VU.Vector col' of+ Left err -> throw err+ Right v -> VU.convert v+ in+ trace+ ( "Best loss: "+ ++ show+ ( getLossFunction (lossFunction cfg) l . interpretDoubleVector+ <$> listToMaybe ordered+ )+ ++ " "+ ++ (if null ordered then "empty" else show (listToMaybe ordered))+ )+ ordered++pickTopNBool ::+ DataFrame ->+ TypedColumn Double ->+ [(Expr Bool, TypedColumn Bool)] ->+ [Expr Bool]+pickTopNBool _ _ [] = []+pickTopNBool _df (TColumn column) ps =+ let+ l = case toVector @Double @VU.Vector column of+ Left e -> throw e+ Right v -> v+ ordered =+ Prelude.take+ 10+ ( map fst $+ L.sortBy+ ( \(_, c2) (_, c1) ->+ if maybe False isInfinite c1+ || maybe False isInfinite c2+ || maybe False isNaN c1+ || maybe False isNaN c2+ then LT+ else compare c1 c2+ )+ ( map+ (\(e, res) -> (e, getLossFunction MutualInformation l (asDoubleVector res)))+ ps+ )+ )+ asDoubleVector c =+ let+ (TColumn col') = c+ in+ case toVector @Bool @VU.Vector col' of+ Left e -> throw e+ Right v -> VU.map (fromIntegral @Int @Double . fromEnum) v+ in+ ordered++satisfiesExamples :: DataFrame -> TypedColumn Double -> Expr Double -> Bool+satisfiesExamples df column expr =+ let+ result = case interpret df expr of+ Left e -> throw e+ Right v -> v+ in+ result == column