Linear.hs

{-# 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)