Linear.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

-- | Oblique split candidates from an L1-regularised logistic hyperplane.
module DataFrame.DecisionTree.Linear (
    bestLinearCandidate,
    fitLinearCandidate,
    careRowsFromFeatures,
    careLabels,
    featName,
    imputeMean,
    materializeFeatureForCare,
) where

import DataFrame.DecisionTree.Numeric (NumExpr (..), numExprCols, numericCols)
import DataFrame.DecisionTree.Types (
    CarePoint (..),
    Direction (..),
    TreeConfig (..),
 )
import DataFrame.Expression.Operators ((.*.), (.+.), (.>.))
import qualified DataFrame.Functions as F
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 disabled or too few care points.
bestLinearCandidate ::
    TreeConfig -> T.Text -> DataFrame -> [CarePoint] -> Maybe (Expr Bool)
bestLinearCandidate cfg target df carePoints
    | not (useLinearSolver cfg) = Nothing
    | length carePoints < minCarePointsForLinear cfg = Nothing
    | otherwise = fitLinearCandidate cfg target df carePoints

-- | Fit an L1 logistic regression to the care points, as a condition.
fitLinearCandidate ::
    TreeConfig -> T.Text -> DataFrame -> [CarePoint] -> Maybe (Expr Bool)
fitLinearCandidate cfg target df carePoints =
    case materializedFeatures target df carePoints of
        [] -> Nothing
        feats -> linearFromFeatures cfg carePoints feats

-- | Each feature's materialised care values with the expression behind them.
materializedFeatures ::
    T.Text ->
    DataFrame ->
    [CarePoint] ->
    [(Expr Double, (T.Text, VU.Vector Double))]
materializedFeatures target df carePoints =
    mapMaybe
        ( \ne ->
            (\m -> (doubleExprOf ne (snd m), m))
                <$> materializeFeatureForCare df carePoints ne
        )
        (featureCols target df)

-- | The expression behind a feature; nullable ones take the care-point mean.
doubleExprOf :: NumExpr -> VU.Vector Double -> Expr Double
doubleExprOf (NDouble e) _ = e
doubleExprOf (NMaybeDouble e) vals = F.fromMaybe (VU.sum vals / fromIntegral (max 1 (VU.length vals))) e

featureCols :: T.Text -> DataFrame -> [NumExpr]
featureCols target df = filter (notElem target . numExprCols) (numericCols df)

linearFromFeatures ::
    TreeConfig ->
    [CarePoint] ->
    [(Expr Double, (T.Text, VU.Vector Double))] ->
    Maybe (Expr Bool)
linearFromFeatures cfg carePoints feats
    | VU.all (== 0) weights = Nothing
    | degenerateHyperplane rows weights (LS.lmIntercept model) = Nothing
    | otherwise = Just hyperplane
  where
    mats = map snd feats
    hyperplane =
        foldl
            (\acc (w, e) -> acc .+. (F.lit w .*. e))
            (F.lit (LS.lmIntercept model))
            [(w, e) | (w, (e, _)) <- zip (VU.toList weights) feats, w /= 0]
            .>. F.lit (0 :: Double)
    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, or Nothing for one class.
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 one side.
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.
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.
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, mean-imputing nullable values.
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)