{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# 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 (..),
    discreteCondVecs,
    discreteConditions,
    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 DataFrame.Errors (DataFrameException (..))
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 :: forall a. Columnable a => Tree a -> Expr a
treeToExpr (Leaf a
v) = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
v
treeToExpr (Branch Expr Bool
cond Tree a
left Tree a
right) = Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
F.ifThenElse Expr Bool
cond (Tree a -> Expr a
forall a. Columnable a => Tree a -> Expr a
treeToExpr Tree a
left) (Tree a -> Expr a
forall a. Columnable a => Tree a -> Expr a
treeToExpr Tree a
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 :: forall a.
(Columnable a, Ord a) =>
TreeConfig -> Expr a -> DataFrame -> Expr a
fitDecisionTree TreeConfig
cfg (Col Text
target) DataFrame
df =
    Expr a -> Expr a
forall a. Columnable a => Expr a -> Expr a
pruneExpr
        (Tree a -> Expr a
forall a. Columnable a => Tree a -> Expr a
treeToExpr (forall a.
(Columnable a, Ord a) =>
TreeConfig
-> Text -> [CondVec] -> DataFrame -> Vector Int -> Tree a -> Tree a
taoOptimizeCV @a TreeConfig
cfg Text
target [CondVec]
condVecs DataFrame
df Vector Int
indices Tree a
initialTree))
  where
    condVecs :: [CondVec]
condVecs = forall a.
(Columnable a, Ord a) =>
TreeConfig -> Text -> DataFrame -> [CondVec]
candidatePool @a TreeConfig
cfg Text
target DataFrame
df
    initialTree :: Tree a
initialTree = forall a.
(Columnable a, Ord a) =>
TreeConfig -> Text -> DataFrame -> Tree a
buildCartTree @a TreeConfig
cfg Text
target DataFrame
df
    indices :: Vector Int
indices = Int -> Int -> Vector Int
forall a. Num a => a -> Int -> Vector a
V.enumFromN Int
0 (DataFrame -> Int
nRows DataFrame
df)
fitDecisionTree TreeConfig
_ Expr a
expr DataFrame
_ =
    DataFrameException -> Expr a
forall a e. Exception e => e -> a
throw (Text -> DataFrameException
NonColumnReferenceException (Text
"fitDecisionTree: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Expr a -> String
forall a. Show a => a -> String
show Expr a
expr)))

-- | The deduplicated numeric + discrete candidate pool for a target column.
candidatePool ::
    forall a.
    (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> [CondVec]
candidatePool :: forall a.
(Columnable a, Ord a) =>
TreeConfig -> Text -> DataFrame -> [CondVec]
candidatePool TreeConfig
cfg Text
target DataFrame
df = [CondVec] -> [CondVec]
dedupCVByExpr ([CondVec]
numericCVs [CondVec] -> [CondVec] -> [CondVec]
forall a. [a] -> [a] -> [a]
++ [CondVec]
discreteCVs)
  where
    dfNoTarget :: DataFrame
dfNoTarget = [Text] -> DataFrame -> DataFrame
exclude [Text
target] DataFrame
df
    numericCVs :: [CondVec]
numericCVs = TreeConfig -> DataFrame -> DataFrame -> [CondVec]
numericCondVecs TreeConfig
cfg DataFrame
dfNoTarget DataFrame
df
    discreteCVs :: [CondVec]
discreteCVs = TargetInfo a -> TreeConfig -> DataFrame -> [CondVec]
forall target.
(Columnable target, Ord target) =>
TargetInfo target -> TreeConfig -> DataFrame -> [CondVec]
discreteCondVecs (forall a.
(Columnable a, Ord a) =>
Text -> DataFrame -> TargetInfo a
targetInfoOrEmpty @a Text
target DataFrame
df) TreeConfig
cfg DataFrame
dfNoTarget

targetInfoOrEmpty ::
    forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> TargetInfo a
targetInfoOrEmpty :: forall a.
(Columnable a, Ord a) =>
Text -> DataFrame -> TargetInfo a
targetInfoOrEmpty Text
target DataFrame
df = TargetInfo a -> Maybe (TargetInfo a) -> TargetInfo a
forall a. a -> Maybe a -> a
fromMaybe (Bool -> Maybe a -> Vector a -> TargetInfo a
forall target.
Bool -> Maybe target -> Vector target -> TargetInfo target
TargetInfo Bool
False Maybe a
forall a. Maybe a
Nothing Vector a
forall a. Vector a
V.empty) (forall target.
(Columnable target, Ord target) =>
Text -> DataFrame -> Maybe (TargetInfo target)
mkTargetInfo @a Text
target DataFrame
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 :: forall a.
(Columnable a, Ord a) =>
TreeConfig -> Int -> Text -> [Expr Bool] -> DataFrame -> Expr a
buildTree TreeConfig
cfg Int
depth Text
target [Expr Bool]
conds DataFrame
df =
    Expr a -> Expr a
forall a. Columnable a => Expr a -> Expr a
pruneExpr (Tree a -> Expr a
forall a. Columnable a => Tree a -> Expr a
treeToExpr (forall a.
(Columnable a, Ord a) =>
TreeConfig
-> Text
-> [Expr Bool]
-> DataFrame
-> Vector Int
-> Tree a
-> Tree a
taoOptimize @a TreeConfig
cfg Text
target [Expr Bool]
conds DataFrame
df Vector Int
indices Tree a
tree))
  where
    tree :: Tree a
tree = forall a.
(Columnable a, Ord a) =>
TreeConfig -> Text -> DataFrame -> Tree a
buildCartTree @a TreeConfig
cfg{maxTreeDepth = depth} Text
target DataFrame
df
    indices :: Vector Int
indices = Int -> Int -> Vector Int
forall a. Num a => a -> Int -> Vector a
V.enumFromN Int
0 (DataFrame -> Int
nRows DataFrame
df)

pruneTree :: forall a. (Columnable a) => Expr a -> Expr a
pruneTree :: forall a. Columnable a => Expr a -> Expr a
pruneTree = Expr a -> Expr a
forall a. Columnable a => Expr a -> Expr a
pruneExpr

partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame)
partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame)
partitionDataFrame Expr Bool
cond DataFrame
df = (Expr Bool -> DataFrame -> DataFrame
filterWhere Expr Bool
cond DataFrame
df, Expr Bool -> DataFrame -> DataFrame
filterWhere (Expr Bool -> Expr Bool
F.not Expr Bool
cond) DataFrame
df)

-- | Laplace-smoothed Gini impurity of the target distribution.
calculateGini ::
    forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double
calculateGini :: forall a. (Columnable a, Ord a) => Text -> DataFrame -> Double
calculateGini Text
target DataFrame
df
    | Double
n Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = Double
0
    | Bool
otherwise = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
probs)
  where
    n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DataFrame -> Int
nRows DataFrame
df)
    counts :: Map a Int
counts = forall a. (Columnable a, Ord a) => Text -> DataFrame -> Map a Int
getCounts @a Text
target DataFrame
df
    numClasses :: Double
numClasses = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Map a Int -> Int
forall k a. Map k a -> Int
M.size Map a Int
counts)
    probs :: [Double]
probs = (Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
c -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
n Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
numClasses)) (Map a Int -> [Int]
forall k a. Map k a -> [a]
M.elems Map a Int
counts)

majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a
majorityValue :: forall a. (Columnable a, Ord a) => Text -> DataFrame -> a
majorityValue Text
target DataFrame
df
    | Map a Int -> Bool
forall k a. Map k a -> Bool
M.null Map a Int
counts = DataFrameException -> a
forall a e. Exception e => e -> a
throw (Text -> DataFrameException
EmptyDataSetException Text
"majorityValue (tree leaf)")
    | Bool
otherwise = (a, Int) -> a
forall a b. (a, b) -> a
fst (((a, Int) -> (a, Int) -> Ordering) -> [(a, Int)] -> (a, Int)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((a, Int) -> Int) -> (a, Int) -> (a, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, Int) -> Int
forall a b. (a, b) -> b
snd) (Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map a Int
counts))
  where
    counts :: Map a Int
counts = forall a. (Columnable a, Ord a) => Text -> DataFrame -> Map a Int
getCounts @a Text
target DataFrame
df

getCounts ::
    forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int
getCounts :: forall a. (Columnable a, Ord a) => Text -> DataFrame -> Map a Int
getCounts Text
target DataFrame
df = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @a DataFrame
df (Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
target) of
    Left DataFrameException
e -> DataFrameException -> Map a Int
forall a e. Exception e => e -> a
throw DataFrameException
e
    Right (TColumn Column
column) -> case forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a Column
column of
        Left DataFrameException
e -> DataFrameException -> Map a Int
forall a e. Exception e => e -> a
throw DataFrameException
e
        Right Vector a
vals -> (Map a Int -> a -> Map a Int) -> Map a Int -> [a] -> Map a Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map a Int
acc a
x -> (Int -> Int -> Int) -> a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) a
x Int
1 Map a Int
acc) Map a Int
forall k a. Map k a
M.empty (Vector a -> [a]
forall a. Vector a -> [a]
V.toList Vector a
vals)

-- | The @p@-th percentile of an expression's values (@0@ on failure/empty).
percentile :: Int -> Expr Double -> DataFrame -> Double
percentile :: Int -> Expr Double -> DataFrame -> Double
percentile Int
p Expr Double
expr DataFrame
df = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @Double DataFrame
df Expr Double
expr of
    Right (TColumn Column
column) -> (DataFrameException -> Double)
-> (Vector Double -> Double)
-> Either DataFrameException (Vector Double)
-> Double
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Double -> DataFrameException -> Double
forall a b. a -> b -> a
const Double
0) (Int -> Vector Double -> Double
percentileOfVec Int
p) (forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @Double Column
column)
    Either DataFrameException (TypedColumn Double)
_ -> Double
0

percentileOfVec :: Int -> V.Vector Double -> Double
percentileOfVec :: Int -> Vector Double -> Double
percentileOfVec Int
p Vector Double
vals
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Double
0
    | Bool
otherwise = Vector Double
sorted Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
100))
  where
    sorted :: Vector Double
sorted = [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList ([Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort (Vector Double -> [Double]
forall a. Vector a -> [a]
V.toList Vector Double
vals))
    n :: Int
n = Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
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 :: forall a.
(Columnable a, Ord a) =>
Text -> DataFrame -> Vector Int -> Map a Double
probsFromIndices Text
target DataFrame
df Vector Int
indices = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @a DataFrame
df (Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
target) of
    Right (TColumn Column
column) -> (DataFrameException -> Map a Double)
-> (Vector a -> Map a Double)
-> Either DataFrameException (Vector a)
-> Map a Double
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Map a Double -> DataFrameException -> Map a Double
forall a b. a -> b -> a
const Map a Double
forall k a. Map k a
M.empty) (Vector Int -> Vector a -> Map a Double
forall a. Ord a => Vector Int -> Vector a -> Map a Double
normaliseCounts Vector Int
indices) (forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a Column
column)
    Either DataFrameException (TypedColumn a)
_ -> Map a Double
forall k a. Map k a
M.empty

normaliseCounts :: (Ord a) => V.Vector Int -> V.Vector a -> M.Map a Double
normaliseCounts :: forall a. Ord a => Vector Int -> Vector a -> Map a Double
normaliseCounts Vector Int
indices Vector a
vals = (Int -> Double) -> Map a Int -> Map a Double
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\Int
c -> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total) Map a Int
counts
  where
    counts :: Map a Int
counts =
        (Map a Int -> Int -> Map a Int)
-> Map a Int -> Vector Int -> Map a Int
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl'
            (\Map a Int
acc Int
i -> (Int -> Int -> Int) -> a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Vector a
vals Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! Int
i) (Int
1 :: Int) Map a Int
acc)
            Map a Int
forall k a. Map k a
M.empty
            Vector Int
indices
    total :: Double
total = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Int
forall a. Vector a -> Int
V.length Vector Int
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 :: forall a.
(Columnable a, Ord a) =>
Tree a -> Text -> DataFrame -> Vector Int -> ProbTree a
buildProbTree (Leaf a
_) Text
target DataFrame
df Vector Int
indices = Map a Double -> Tree (Map a Double)
forall a. a -> Tree a
Leaf (forall a.
(Columnable a, Ord a) =>
Text -> DataFrame -> Vector Int -> Map a Double
probsFromIndices @a Text
target DataFrame
df Vector Int
indices)
buildProbTree (Branch Expr Bool
cond Tree a
left Tree a
right) Text
target DataFrame
df Vector Int
indices =
    Expr Bool
-> Tree (Map a Double)
-> Tree (Map a Double)
-> Tree (Map a Double)
forall a. Expr Bool -> Tree a -> Tree a -> Tree a
Branch
        Expr Bool
cond
        (forall a.
(Columnable a, Ord a) =>
Tree a -> Text -> DataFrame -> Vector Int -> ProbTree a
buildProbTree @a Tree a
left Text
target DataFrame
df Vector Int
l)
        (forall a.
(Columnable a, Ord a) =>
Tree a -> Text -> DataFrame -> Vector Int -> ProbTree a
buildProbTree @a Tree a
right Text
target DataFrame
df Vector Int
r)
  where
    (Vector Int
l, Vector Int
r) = Expr Bool -> DataFrame -> Vector Int -> (Vector Int, Vector Int)
partitionIndices Expr Bool
cond DataFrame
df Vector Int
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 :: forall a.
(Columnable a, Ord a) =>
TreeConfig -> Expr a -> DataFrame -> Map a (Expr Double)
fitProbTree TreeConfig
cfg (Col Text
target) DataFrame
df = ProbTree a -> Map a (Expr Double)
forall a.
(Columnable a, Ord a) =>
ProbTree a -> Map a (Expr Double)
probExprs (forall a.
(Columnable a, Ord a) =>
Tree a -> Text -> DataFrame -> Vector Int -> ProbTree a
buildProbTree @a Tree a
pruned Text
target DataFrame
df Vector Int
indices)
  where
    conds :: [Expr Bool]
conds =
        [Expr Bool] -> [Expr Bool]
nubByExpr
            ( TreeConfig -> DataFrame -> [Expr Bool]
numericConditions TreeConfig
cfg DataFrame
dfNoTarget
                [Expr Bool] -> [Expr Bool] -> [Expr Bool]
forall a. [a] -> [a] -> [a]
++ TargetInfo a -> TreeConfig -> DataFrame -> [Expr Bool]
forall target.
(Columnable target, Ord target) =>
TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool]
discreteConditions (forall a.
(Columnable a, Ord a) =>
Text -> DataFrame -> TargetInfo a
targetInfoOrEmpty @a Text
target DataFrame
df) TreeConfig
cfg DataFrame
dfNoTarget
            )
    dfNoTarget :: DataFrame
dfNoTarget = [Text] -> DataFrame -> DataFrame
exclude [Text
target] DataFrame
df
    indices :: Vector Int
indices = Int -> Int -> Vector Int
forall a. Num a => a -> Int -> Vector a
V.enumFromN Int
0 (DataFrame -> Int
nRows DataFrame
df)
    pruned :: Tree a
pruned =
        Tree a -> Tree a
forall a. Columnable a => Tree a -> Tree a
pruneDead
            (forall a.
(Columnable a, Ord a) =>
TreeConfig
-> Text
-> [Expr Bool]
-> DataFrame
-> Vector Int
-> Tree a
-> Tree a
taoOptimize @a TreeConfig
cfg Text
target [Expr Bool]
conds DataFrame
df Vector Int
indices (forall a.
(Columnable a, Ord a) =>
TreeConfig -> Text -> DataFrame -> Tree a
buildCartTree @a TreeConfig
cfg Text
target DataFrame
df))
fitProbTree TreeConfig
_ Expr a
expr DataFrame
_ =
    DataFrameException -> Map a (Expr Double)
forall a e. Exception e => e -> a
throw (Text -> DataFrameException
NonColumnReferenceException (Text
"fitProbTree: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Expr a -> String
forall a. Show a => a -> String
show Expr a
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 :: forall a.
(Columnable a, Ord a) =>
ProbTree a -> Map a (Expr Double)
probExprs ProbTree a
tree = [(a, Expr Double)] -> Map a (Expr Double)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(a
c, a -> ProbTree a -> Expr Double
forall a. Ord a => a -> ProbTree a -> Expr Double
classExpr a
c ProbTree a
tree) | a
c <- [a] -> [a]
forall a. Eq a => [a] -> [a]
nub (ProbTree a -> [a]
forall a. ProbTree a -> [a]
allClasses ProbTree a
tree)]

allClasses :: ProbTree a -> [a]
allClasses :: forall a. ProbTree a -> [a]
allClasses (Leaf Map a Double
m) = Map a Double -> [a]
forall k a. Map k a -> [k]
M.keys Map a Double
m
allClasses (Branch Expr Bool
_ Tree (Map a Double)
l Tree (Map a Double)
r) = Tree (Map a Double) -> [a]
forall a. ProbTree a -> [a]
allClasses Tree (Map a Double)
l [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Tree (Map a Double) -> [a]
forall a. ProbTree a -> [a]
allClasses Tree (Map a Double)
r

classExpr :: (Ord a) => a -> ProbTree a -> Expr Double
classExpr :: forall a. Ord a => a -> ProbTree a -> Expr Double
classExpr a
c (Leaf Map a Double
m) = Double -> Expr Double
forall a. Columnable a => a -> Expr a
Lit (Double -> a -> Map a Double -> Double
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Double
0.0 a
c Map a Double
m)
classExpr a
c (Branch Expr Bool
cond Tree (Map a Double)
l Tree (Map a Double)
r) = Expr Bool -> Expr Double -> Expr Double -> Expr Double
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
F.ifThenElse Expr Bool
cond (a -> Tree (Map a Double) -> Expr Double
forall a. Ord a => a -> ProbTree a -> Expr Double
classExpr a
c Tree (Map a Double)
l) (a -> Tree (Map a Double) -> Expr Double
forall a. Ord a => a -> ProbTree a -> Expr Double
classExpr a
c Tree (Map a Double)
r)