Statistics.hs

{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

module DataFrame.Operations.Statistics where

import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU

import Prelude as P

import Control.Exception (throw)
import Data.Function ((&))
import Data.Maybe (fromMaybe, isJust)
import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
import DataFrame.Errors (DataFrameException (..))
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    empty,
    getColumn,
 )
import DataFrame.Internal.Expression
import DataFrame.Internal.Interpreter
import DataFrame.Internal.Nullable (BaseType)
import DataFrame.Internal.Row (showValue, toAny)
import DataFrame.Internal.Statistics
import DataFrame.Internal.Types
import DataFrame.Operations.Core
import DataFrame.Operations.Subset (filterJust)
import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore)
import Text.Printf (printf)
import Type.Reflection (typeRep)

{- | Show a frequency table for a categorical feaure.

__Examples:__

@
ghci> df <- D.readCsv ".\/data\/housing.csv"

ghci> D.frequencies "ocean_proximity" df

---------------------------------------------------------------------
   Statistic    | <1H OCEAN | INLAND | ISLAND | NEAR BAY | NEAR OCEAN
----------------|-----------|--------|--------|----------|-----------
      Text      |    Any    |  Any   |  Any   |   Any    |    Any
----------------|-----------|--------|--------|----------|-----------
 Count          | 9136      | 6551   | 5      | 2290     | 2658
 Percentage (%) | 44.26%    | 31.74% | 0.02%  | 11.09%   | 12.88%
@
-}
frequencies :: forall a. (Columnable a) => Expr a -> DataFrame -> DataFrame
frequencies expr df =
    let
        counts = valueCounts expr df
        calculatePercentage cs k = toAny $ toPct2dp (fromIntegral k / fromIntegral (P.sum $ map snd cs))
        initDf =
            empty
                & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"])
        freqs col =
            L.foldl'
                ( \d (col, k) ->
                    insertVector
                        (showValue @a col)
                        (V.fromList [toAny k, calculatePercentage counts k])
                        d
                )
                initDf
                counts
     in
        case columnAsVector expr df of
            Left err -> throw err
            Right column -> freqs column

-- | Calculates the mean of a given column as a standalone value.
mean ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
mean (Col name) df = case _getColumnAsDouble name df of
    Just xs -> meanDouble' xs
    Nothing -> error "[INTERNAL ERROR] Column is non-numeric"
mean expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> mean' xs

meanMaybe ::
    forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double
meanMaybe (Col name) df =
    (mean' . optionalToDoubleVector)
        (either throw id (columnAsVector (Col @(Maybe a) name) df))
meanMaybe expr df = case interpret @(Maybe a) df expr of
    Left e -> throw e
    Right (TColumn col) -> case toVector @(Maybe a) col of
        Left e -> throw e
        Right xs -> (mean' . optionalToDoubleVector) xs

-- | Calculates the median of a given column as a standalone value.
median ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
median (Col name) df = case columnAsUnboxedVector (Col @a name) df of
    Right xs -> median' xs
    Left e -> throw e
median expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> median' xs

-- | Calculates the median of a given column (containing optional values) as a standalone value.
medianMaybe ::
    forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double
medianMaybe (Col name) df =
    (median' . optionalToDoubleVector)
        (either throw id (columnAsVector (Col @(Maybe a) name) df))
medianMaybe expr df = case interpret @(Maybe a) df expr of
    Left e -> throw e
    Right (TColumn col) -> case toVector @(Maybe a) col of
        Left e -> throw e
        Right xs -> (median' . optionalToDoubleVector) xs

-- | Calculates the nth percentile of a given column as a standalone value.
percentile ::
    forall a.
    (Columnable a, Real a, VU.Unbox a) => Int -> Expr a -> DataFrame -> Double
percentile n (Col name) df = case columnAsUnboxedVector (Col @a name) df of
    Right xs -> percentile' n xs
    Left e -> throw e
percentile n expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> percentile' n xs

-- | Calculates the nth percentile of a given column as a standalone value.
genericPercentile ::
    forall a.
    (Columnable a, Ord a) => Int -> Expr a -> DataFrame -> a
genericPercentile n (Col name) df = case columnAsVector (Col @a name) df of
    Right xs -> percentileOrd' n xs
    Left e -> throw e
genericPercentile n expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toVector @a col of
        Left e -> throw e
        Right xs -> percentileOrd' n xs

-- | Calculates the standard deviation of a given column as a standalone value.
standardDeviation ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
standardDeviation (Col name) df = case columnAsUnboxedVector (Col @a name) df of
    Right xs -> (sqrt . variance') xs
    Left e -> throw e
standardDeviation expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> (sqrt . variance') xs

-- | Calculates the skewness of a given column as a standalone value.
skewness ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
skewness (Col name) df = case columnAsUnboxedVector (Col @a name) df of
    Right xs -> skewness' xs
    Left e -> throw e
skewness expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> skewness' xs

-- | Calculates the variance of a given column as a standalone value.
variance ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
variance (Col name) df = case _getColumnAsDouble name df of
    Just xs -> varianceDouble' xs
    Nothing -> error "[INTERNAL ERROR] Column is non-numeric"
variance expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> variance' xs

-- | Calculates the inter-quartile range of a given column as a standalone value.
interQuartileRange ::
    forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double
interQuartileRange (Col name) df = case columnAsUnboxedVector (Col @a name) df of
    Right xs -> interQuartileRange' xs
    Left e -> throw e
interQuartileRange expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn col) -> case toUnboxedVector @a col of
        Left e -> throw e
        Right xs -> interQuartileRange' xs

-- | Calculates the Pearson's correlation coefficient between two given columns as a standalone value.
correlation :: T.Text -> T.Text -> DataFrame -> Maybe Double
correlation first second df = do
    f <- _getColumnAsDouble first df
    s <- _getColumnAsDouble second df
    correlation' f s

_getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double)
_getColumnAsDouble name df = case getColumn name df of
    Just (UnboxedColumn (f :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of
        Just Refl -> Just f
        Nothing -> case sIntegral @a of
            STrue -> Just (VU.map fromIntegral f)
            SFalse -> case sFloating @a of
                STrue -> Just (VU.map realToFrac f)
                SFalse -> Nothing
    Nothing ->
        throw $
            ColumnNotFoundException name "_getColumnAsDouble" (M.keys $ columnIndices df)
    _ -> Nothing -- Return a type mismatch error here.
{-# INLINE _getColumnAsDouble #-}

optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double
optionalToDoubleVector =
    VU.fromList
        . V.foldl'
            (\acc e -> if isJust e then realToFrac (fromMaybe 0 e) : acc else acc)
            []

-- | Calculates the sum of a given column as a standalone value.
sum ::
    forall a. (Columnable a, Num a) => Expr a -> DataFrame -> a
sum (Col name) df = case getColumn name df of
    Nothing -> throw $ ColumnNotFoundException name "sum" (M.keys $ columnIndices df)
    Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of
        Just Refl -> VG.sum column
        Nothing -> 0
    Just ((BoxedColumn (column :: V.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of
        Just Refl -> VG.sum column
        Nothing -> 0
    Just ((OptionalColumn (column :: V.Vector (Maybe a')))) -> case testEquality (typeRep @a') (typeRep @a) of
        Just Refl -> VG.sum (VG.map (fromMaybe 0) column)
        Nothing -> 0
sum expr df = case interpret df expr of
    Left e -> throw e
    Right (TColumn xs) -> case toVector @a @V.Vector xs of
        Left e -> throw e
        Right xs -> VG.sum xs

{- | /O(n)/ Impute missing values in a column using a derived scalar.

Given

* an expression @f :: 'Expr' b -> 'Expr' b@ that, when interpreted over a
  non-nullable column, produces the same value in every row (for example a
  mean, median, or other aggregate), and
* a nullable column @'Expr' ('Maybe' b)@

this function:

1. Drops all @Nothing@ values from the target column.
2. Interprets @f@ on the remaining non-null values.
3. Checks that the resulting column contains a single repeated value.
4. Uses that value to impute all @Nothing@s in the original column.

==== __Throws__

* 'DataFrameException' - if the column does not exist, is empty,

==== __Example__
@
>>> :set -XOverloadedStrings
>>> import qualified DataFrame as D
>>> let df =
...       D.fromNamedColumns
...         [ ("age", D.fromList [Just 10, Nothing, Just 20 :: Maybe Int]) ]
>>>
>>> -- Impute missing ages with the mean of the observed ages
>>> D.imputeWith F.mean "age" df
-- age
-- ----
-- 10
-- 15
-- 20
@
-}
instance {-# OVERLAPPING #-} (Columnable b) => ImputeOp (Maybe b) where
    runImpute = imputeCore

    runImputeWith f col@(Col columnName) df =
        case interpret @b (filterJust columnName df) (f (Col @b columnName)) of
            Left e -> throw e
            Right (TColumn value) -> case headColumn @b value of
                Left e -> throw e
                Right h ->
                    if all (== h) (toList @b value)
                        then imputeCore col h df
                        else error "Impute expression returned more than one value"
    runImputeWith _ _ df = df

imputeWith ::
    forall a.
    (ImputeOp a, Columnable (BaseType a)) =>
    (Expr (BaseType a) -> Expr (BaseType a)) ->
    Expr a ->
    DataFrame ->
    DataFrame
imputeWith = runImputeWith

applyStatistic ::
    (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double
applyStatistic f name df = apply =<< _getColumnAsDouble name (filterJust name df)
  where
    apply col =
        let
            res = f col
         in
            if isNaN res then Nothing else pure res
{-# INLINE applyStatistic #-}

applyStatistics ::
    (VU.Vector Double -> VU.Vector Double) ->
    T.Text ->
    DataFrame ->
    Maybe (VU.Vector Double)
applyStatistics f name df = fmap f (_getColumnAsDouble name (filterJust name df))

-- | Descriptive statistics of the numeric columns.
summarize :: DataFrame -> DataFrame
summarize df =
    fold
        columnStats
        (columnNames df)
        ( fromNamedColumns
            [
                ( "Statistic"
                , fromList
                    [ "Count" :: T.Text
                    , "Mean"
                    , "Minimum"
                    , "25%"
                    , "Median"
                    , "75%"
                    , "Max"
                    , "StdDev"
                    , "IQR"
                    , "Skewness"
                    ]
                )
            ]
        )
  where
    columnStats name d =
        if all isJust (stats name)
            then
                insertUnboxedVector
                    name
                    (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name))
                    d
            else d
    stats name =
        let
            count = fromIntegral . numElements <$> getColumn name df
            quantiles = applyStatistics (quantiles' (VU.fromList [0, 1, 2, 3, 4]) 4) name df
            min' = flip (VG.!) 0 <$> quantiles
            quartile1 = flip (VG.!) 1 <$> quantiles
            median' = flip (VG.!) 2 <$> quantiles
            quartile3 = flip (VG.!) 3 <$> quantiles
            max' = flip (VG.!) 4 <$> quantiles
            iqr = (-) <$> quartile3 <*> quartile1
            doubleColumn col = _getColumnAsDouble col (filterJust col df)
         in
            [ count
            , mean' <$> doubleColumn name
            , min'
            , quartile1
            , median'
            , quartile3
            , max'
            , sqrt . variance' <$> doubleColumn name
            , iqr
            , skewness' <$> doubleColumn name
            ]

-- | Round a @Double@ to Specified Precision
roundTo :: Int -> Double -> Double
roundTo n x = fromInteger (round $ x * 10 ^ n) / 10.0 ^^ n

toPct2dp :: Double -> String
toPct2dp x
    | x < 0.00005 = "<0.01%"
    | otherwise = printf "%.2f%%" (x * 100)