{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DataFrame.Internal.Statistics where

import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as VA
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Exception (throw)
import Control.Monad.ST (runST)
import DataFrame.Errors (DataFrameException (..))

mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double
mean' :: forall a. (Real a, Unbox a) => Vector a -> Double
mean' Vector a
samp
    | Vector a -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector a
samp = DataFrameException -> Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Double) -> DataFrameException -> Double
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"mean"
    | Bool
otherwise = a -> Double
forall a. Real a => a -> Double
rtf (Vector a -> a
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector a
samp) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
samp)
{-# INLINE [0] mean' #-}

meanDouble' :: VU.Vector Double -> Double
meanDouble' :: Vector Double -> Double
meanDouble' Vector Double
samp
    | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Double
samp = DataFrameException -> Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Double) -> DataFrameException -> Double
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"mean"
    | Bool
otherwise = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Double
samp Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
samp)
{-# INLINE meanDouble' #-}

meanInt' :: VU.Vector Int -> Double
meanInt' :: Vector Int -> Double
meanInt' Vector Int
samp
    | Vector Int -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Int
samp = DataFrameException -> Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Double) -> DataFrameException -> Double
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"mean"
    | Bool
otherwise = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Int
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Int
samp) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
samp)
{-# INLINE meanInt' #-}

{-# RULES
"mean'/Double" [1] forall (xs :: VU.Vector Double).
    mean' xs =
        meanDouble' xs
"mean'/Int" [1] forall (xs :: VU.Vector Int).
    mean' xs =
        meanInt' xs
    #-}

median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double
median' :: forall a. (Real a, Unbox a) => Vector a -> Double
median' Vector a
samp
    | Vector a -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector a
samp = DataFrameException -> Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Double) -> DataFrameException -> Double
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"median"
    | Bool
otherwise = (forall s. ST s Double) -> Double
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Double) -> Double)
-> (forall s. ST s Double) -> Double
forall a b. (a -> b) -> a -> b
$ do
        MVector s a
mutableSamp <- Vector a -> ST s (MVector (PrimState (ST s)) a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector a
samp
        MVector (PrimState (ST s)) a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VA.sort MVector s a
MVector (PrimState (ST s)) a
mutableSamp
        let len :: Int
len = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
samp
            middleIndex :: Int
middleIndex = Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        a
middleElement <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s a
MVector (PrimState (ST s)) a
mutableSamp Int
middleIndex
        if Int -> Bool
forall a. Integral a => a -> Bool
odd Int
len
            then Double -> ST s Double
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Double
forall a. Real a => a -> Double
rtf a
middleElement)
            else do
                a
prev <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s a
MVector (PrimState (ST s)) a
mutableSamp (Int
middleIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                Double -> ST s Double
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Double
forall a. Real a => a -> Double
rtf (a
middleElement a -> a -> a
forall a. Num a => a -> a -> a
+ a
prev) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
{-# INLINE median' #-}


data VarAcc
    = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double
    deriving (Int -> VarAcc -> ShowS
[VarAcc] -> ShowS
VarAcc -> String
(Int -> VarAcc -> ShowS)
-> (VarAcc -> String) -> ([VarAcc] -> ShowS) -> Show VarAcc
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> VarAcc -> ShowS
showsPrec :: Int -> VarAcc -> ShowS
$cshow :: VarAcc -> String
show :: VarAcc -> String
$cshowList :: [VarAcc] -> ShowS
showList :: [VarAcc] -> ShowS
Show)

varianceStep :: VarAcc -> Double -> VarAcc
varianceStep :: VarAcc -> Double -> VarAcc
varianceStep (VarAcc !Int
n !Double
mean !Double
m2) !Double
x =
    let !n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        !delta :: Double
delta = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mean
        !mean' :: Double
mean' = Double
mean Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
delta Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n'
        !m2' :: Double
m2' = Double
m2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
delta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mean')
     in Int -> Double -> Double -> VarAcc
VarAcc Int
n' Double
mean' Double
m2'
{-# INLINE varianceStep #-}

computeVariance :: VarAcc -> Double
computeVariance :: VarAcc -> Double
computeVariance (VarAcc !Int
n Double
_ !Double
m2)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = Double
0 
    | Bool
otherwise = Double
m2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE computeVariance #-}

variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double
variance' :: forall a. (Real a, Unbox a) => Vector a -> Double
variance' = VarAcc -> Double
computeVariance (VarAcc -> Double) -> (Vector a -> VarAcc) -> Vector a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarAcc -> Double -> VarAcc) -> VarAcc -> Vector Double -> VarAcc
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' VarAcc -> Double -> VarAcc
varianceStep (Int -> Double -> Double -> VarAcc
VarAcc Int
0 Double
0 Double
0) (Vector Double -> VarAcc)
-> (Vector a -> Vector Double) -> Vector a -> VarAcc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Double) -> Vector a -> Vector Double
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map a -> Double
forall a. Real a => a -> Double
rtf
{-# INLINE variance' #-}

varianceDouble' :: VU.Vector Double -> Double
varianceDouble' :: Vector Double -> Double
varianceDouble' = VarAcc -> Double
computeVariance (VarAcc -> Double)
-> (Vector Double -> VarAcc) -> Vector Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarAcc -> Double -> VarAcc) -> VarAcc -> Vector Double -> VarAcc
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' VarAcc -> Double -> VarAcc
varianceStep (Int -> Double -> Double -> VarAcc
VarAcc Int
0 Double
0 Double
0)
{-# INLINE varianceDouble' #-}


data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Int -> SkewAcc -> ShowS
[SkewAcc] -> ShowS
SkewAcc -> String
(Int -> SkewAcc -> ShowS)
-> (SkewAcc -> String) -> ([SkewAcc] -> ShowS) -> Show SkewAcc
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SkewAcc -> ShowS
showsPrec :: Int -> SkewAcc -> ShowS
$cshow :: SkewAcc -> String
show :: SkewAcc -> String
$cshowList :: [SkewAcc] -> ShowS
showList :: [SkewAcc] -> ShowS
Show)

skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc
skewnessStep :: forall a. (Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc
skewnessStep (SkewAcc !Int
n !Double
mean !Double
m2 !Double
m3) !a
x' =
    let !n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        x :: Double
x = a -> Double
forall a. Real a => a -> Double
rtf a
x'
        !k :: Double
k = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n'
        !delta :: Double
delta = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mean
        !mean' :: Double
mean' = Double
mean Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
delta Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k
        !m2' :: Double
m2' = Double
m2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
delta Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k
        !m3' :: Double
m3' = Double
m3 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
delta Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
delta Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
m2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k
     in Int -> Double -> Double -> Double -> SkewAcc
SkewAcc Int
n' Double
mean' Double
m2' Double
m3'
{-# INLINE skewnessStep #-}

computeSkewness :: SkewAcc -> Double
computeSkewness :: SkewAcc -> Double
computeSkewness (SkewAcc Int
n Double
_ Double
m2 Double
m3)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
3 = Double
0 
    | Bool
otherwise = (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
m3) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
m2 Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
3)
{-# INLINE computeSkewness #-}

skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double
skewness' :: forall a. (Unbox a, Real a, Num a) => Vector a -> Double
skewness' = SkewAcc -> Double
computeSkewness (SkewAcc -> Double) -> (Vector a -> SkewAcc) -> Vector a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SkewAcc -> a -> SkewAcc) -> SkewAcc -> Vector a -> SkewAcc
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' SkewAcc -> a -> SkewAcc
forall a. (Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc
skewnessStep (Int -> Double -> Double -> Double -> SkewAcc
SkewAcc Int
0 Double
0 Double
0 Double
0)
{-# INLINE skewness' #-}

data CorrelationStats
    = CorrelationStats
        {-# UNPACK #-} !Double
        {-# UNPACK #-} !Double
        {-# UNPACK #-} !Double
        {-# UNPACK #-} !Double
        {-# UNPACK #-} !Double

correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double
correlation' :: Vector Double -> Vector Double -> Maybe Double
correlation' Vector Double
xs Vector Double
ys
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = Maybe Double
forall a. Maybe a
Nothing
    | Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
ys = Maybe Double
forall a. Maybe a
Nothing
    | Bool
otherwise =
        let nf :: Double
nf = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
            initial :: CorrelationStats
initial = Double -> Double -> Double -> Double -> Double -> CorrelationStats
CorrelationStats Double
0 Double
0 Double
0 Double
0 Double
0
            (CorrelationStats Double
sumX Double
sumY Double
sumXX Double
sumYY Double
sumXY) = (CorrelationStats -> Int -> Double -> CorrelationStats)
-> CorrelationStats -> Vector Double -> CorrelationStats
forall b a. Unbox b => (a -> Int -> b -> a) -> a -> Vector b -> a
VU.ifoldl' CorrelationStats -> Int -> Double -> CorrelationStats
step CorrelationStats
initial Vector Double
xs

            !num :: Double
num = Double
nf Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumXY Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sumX Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumY
            !den :: Double
den = Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
nf Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumXX Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sumX Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumX) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
nf Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumYY Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sumY Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sumY))
         in Double -> Maybe Double
forall a. a -> Maybe a
Just (Double
num Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den)
  where
    n :: Int
n = Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
xs
    step :: CorrelationStats -> Int -> Double -> CorrelationStats
step (CorrelationStats Double
sx Double
sy Double
sxx Double
syy Double
sxy) Int
i Double
x =
        let !y :: Double
y = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Double
ys Int
i
         in Double -> Double -> Double -> Double -> Double -> CorrelationStats
CorrelationStats (Double
sx Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x) (Double
sy Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y) (Double
sxx Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x) (Double
syy Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y) (Double
sxy Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y)
{-# INLINE correlation' #-}

quantiles' ::
    (VU.Unbox a, Num a, Real a) =>
    VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
quantiles' :: forall a.
(Unbox a, Num a, Real a) =>
Vector Int -> Int -> Vector a -> Vector Double
quantiles' Vector Int
qs Int
q Vector a
samp
    | Vector a -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector a
samp = DataFrameException -> Vector Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector Double)
-> DataFrameException -> Vector Double
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"quantiles"
    | Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = DataFrameException -> Vector Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector Double)
-> DataFrameException -> Vector Double
forall a b. (a -> b) -> a -> b
$ Int -> DataFrameException
WrongQuantileNumberException Int
q
    | (Int -> Bool) -> Vector Int -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
q) Vector Int
qs = DataFrameException -> Vector Double
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector Double)
-> DataFrameException -> Vector Double
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> DataFrameException
WrongQuantileIndexException Vector Int
qs Int
q
    | Bool
otherwise = (forall s. ST s (Vector Double)) -> Vector Double
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Double)) -> Vector Double)
-> (forall s. ST s (Vector Double)) -> Vector Double
forall a b. (a -> b) -> a -> b
$ do
        let !n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
samp
        MVector s a
mutableSamp <- Vector a -> ST s (MVector (PrimState (ST s)) a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector a
samp
        MVector (PrimState (ST s)) a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VA.sort MVector s a
MVector (PrimState (ST s)) a
mutableSamp
        (Int -> ST s Double) -> Vector Int -> ST s (Vector Double)
forall (m :: * -> *) a b.
(Monad m, Unbox a, Unbox b) =>
(a -> m b) -> Vector a -> m (Vector b)
VU.mapM
            ( \Int
i -> do
                let !p :: Double
p = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
q
                    !position :: Double
position = Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    !index :: Int
index = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
position
                    !f :: Double
f = Double
position Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
index
                Double
x <- (a -> Double) -> ST s a -> ST s Double
forall a b. (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Double
forall a. Real a => a -> Double
rtf (MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s a
MVector (PrimState (ST s)) a
mutableSamp Int
index)
                if Double
f Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                    then Double -> ST s Double
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
x
                    else do
                        Double
y <- (a -> Double) -> ST s a -> ST s Double
forall a b. (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Double
forall a. Real a => a -> Double
rtf (MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s a
MVector (PrimState (ST s)) a
mutableSamp (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
                        Double -> ST s Double
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> ST s Double) -> Double -> ST s Double
forall a b. (a -> b) -> a -> b
$ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
f) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
f Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y
            )
            Vector Int
qs
{-# INLINE quantiles' #-}

percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double
percentile' :: forall a. (Unbox a, Num a, Real a) => Int -> Vector a -> Double
percentile' Int
n = Vector Double -> Double
forall a. Unbox a => Vector a -> a
VU.head (Vector Double -> Double)
-> (Vector a -> Vector Double) -> Vector a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector a -> Vector Double
forall a.
(Unbox a, Num a, Real a) =>
Vector Int -> Int -> Vector a -> Vector Double
quantiles' ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int
n]) Int
100

quantilesOrd' ::
    (Ord a, Eq a) =>
    VU.Vector Int -> Int -> V.Vector a -> V.Vector a
quantilesOrd' :: forall a.
(Ord a, Eq a) =>
Vector Int -> Int -> Vector a -> Vector a
quantilesOrd' Vector Int
qs Int
q Vector a
samp
    | Vector a -> Bool
forall a. Vector a -> Bool
V.null Vector a
samp = DataFrameException -> Vector a
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector a) -> DataFrameException -> Vector a
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
EmptyDataSetException Text
"quantiles"
    | Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = DataFrameException -> Vector a
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector a) -> DataFrameException -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> DataFrameException
WrongQuantileNumberException Int
q
    | (Int -> Bool) -> Vector Int -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
q) Vector Int
qs = DataFrameException -> Vector a
forall a e. Exception e => e -> a
throw (DataFrameException -> Vector a) -> DataFrameException -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> DataFrameException
WrongQuantileIndexException Vector Int
qs Int
q
    | Bool
otherwise = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
        let !n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
samp
        MVector s a
mutableSamp <- Vector a -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw Vector a
samp
        MVector (PrimState (ST s)) a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VA.sort MVector s a
MVector (PrimState (ST s)) a
mutableSamp
        (Int -> ST s a) -> Vector Int -> ST s (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM
            ( \Int
i -> do
                let !p :: Double
p = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
q
                    !position :: Double
position = Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    !index :: Int
index = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
position
                
                
                MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector s a
MVector (PrimState (ST s)) a
mutableSamp Int
index
            )
            (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector Int
qs)

percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a
percentileOrd' :: forall a. (Ord a, Eq a) => Int -> Vector a -> a
percentileOrd' Int
n = Vector a -> a
forall a. Vector a -> a
V.head (Vector a -> a) -> (Vector a -> Vector a) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector a -> Vector a
forall a.
(Ord a, Eq a) =>
Vector Int -> Int -> Vector a -> Vector a
quantilesOrd' ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int
n]) Int
100

interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double
interQuartileRange' :: forall a. (Unbox a, Num a, Real a) => Vector a -> Double
interQuartileRange' Vector a
samp =
    let quartiles :: Vector Double
quartiles = Vector Int -> Int -> Vector a -> Vector Double
forall a.
(Unbox a, Num a, Real a) =>
Vector Int -> Int -> Vector a -> Vector Double
quantiles' ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int
1, Int
3]) Int
4 Vector a
samp
     in Vector Double
quartiles Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Vector Double
quartiles Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
0
{-# INLINE interQuartileRange' #-}

meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double
meanSquaredError :: Vector Double -> Vector Double -> Maybe Double
meanSquaredError Vector Double
target Vector Double
prediction =
    let
        squareDiff :: Double
squareDiff = (Double -> Int -> Double -> Double)
-> Double -> Vector Double -> Double
forall b a. Unbox b => (a -> Int -> b -> a) -> a -> Vector b -> a
VU.ifoldl' (\Double
sq Int
i Double
e -> (Double
e Double -> Double -> Double
forall a. Num a => a -> a -> a
- Vector Double
target Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i) Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sq) Double
0 Vector Double
prediction
     in
        Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
squareDiff Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
target) (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
prediction))
{-# INLINE meanSquaredError #-}

mutualInformationBinned ::
    Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double
mutualInformationBinned :: Int -> Vector Double -> Vector Double -> Maybe Double
mutualInformationBinned Int
k Vector Double
xs Vector Double
ys
    | Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
ys = Maybe Double
forall a. Maybe a
Nothing
    | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Double
xs = Maybe Double
forall a. Maybe a
Nothing
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = Maybe Double
forall a. Maybe a
Nothing
    | Double
rx Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
ry Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = Double -> Maybe Double
forall a. a -> Maybe a
Just Double
0
    | Bool
otherwise =
        let bx :: Vector Int
bx = (Double -> Int) -> Vector Double -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Double -> Double -> Int -> Double -> Int
binIndex Double
xmin Double
xmax Int
k) Vector Double
xs
            by :: Vector Int
by = (Double -> Int) -> Vector Double -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Double -> Double -> Int -> Double -> Int
binIndex Double
ymin Double
ymax Int
k) Vector Double
ys
            n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
xs) :: Double
            mx :: Vector Int
mx = Int -> Vector Int -> Vector Int
bincount Int
k Vector Int
bx
            my :: Vector Int
my = Int -> Vector Int -> Vector Int
bincount Int
k Vector Int
by
            mxy :: Vector Int
mxy = Int -> Vector Int -> Vector Int -> Vector Int
jointBincount Int
k Vector Int
bx Vector Int
by
         in Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$
                [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                    [ let !cxy :: Double
cxy = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c
                          !pxy :: Double
pxy = Double
cxy Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
                          !px :: Double
px = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int
mx Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
                          !py :: Double
py = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int
my Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
                       in if Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
0 else Double
pxy Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
2 (Double
pxy Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
px Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
py))
                    | Int
i <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    , Int
j <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    , let !c :: Int
c = Vector Int
mxy Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)
                    ]
  where
    (Double
xmin, Double
xmax) = (Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
VU.minimum Vector Double
xs, Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
VU.maximum Vector Double
xs)
    (Double
ymin, Double
ymax) = (Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
VU.minimum Vector Double
ys, Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
VU.maximum Vector Double
ys)
    rx :: Double
rx = Double
xmax Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xmin
    ry :: Double
ry = Double
ymax Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
ymin

binIndex :: Double -> Double -> Int -> Double -> Int
binIndex :: Double -> Double -> Int -> Double -> Int
binIndex Double
lo Double
hi Int
k Double
x
    | Double
hi Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
lo = Int
0
    | Bool
otherwise =
        let !t :: Double
t = (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo)
            !ix :: Int
ix = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t) :: Int
         in Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
ix)
{-# INLINE binIndex #-}

bincount :: Int -> VU.Vector Int -> VU.Vector Int
bincount :: Int -> Vector Int -> Vector Int
bincount Int
k Vector Int
bs = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
mv <- Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
k Int
0)
    Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
bs ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
b -> do
        let i :: Int
i
                | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Int
0
                | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                | Bool
otherwise = Int
b
        Int
x <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s Int
MVector (PrimState (ST s)) Int
mv Int
i
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.write MVector s Int
MVector (PrimState (ST s)) Int
mv Int
i (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
mv
{-# INLINE bincount #-}

jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int
jointBincount :: Int -> Vector Int -> Vector Int -> Vector Int
jointBincount Int
k Vector Int
bx Vector Int
by = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
mv <- Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int
0)
    Vector (Int, Int) -> ((Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
bx Vector Int
by) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, Int
j) -> do
        let ii :: Int
ii = Int -> Int -> Int -> Int
forall {a}. Ord a => a -> a -> a -> a
clamp Int
i Int
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            jj :: Int
jj = Int -> Int -> Int -> Int
forall {a}. Ord a => a -> a -> a -> a
clamp Int
j Int
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            ix :: Int
ix = Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jj
        Int
x <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s Int
MVector (PrimState (ST s)) Int
mv Int
ix
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.write MVector s Int
MVector (PrimState (ST s)) Int
mv Int
ix (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
mv
  where
    clamp :: a -> a -> a -> a
clamp a
z a
a a
b = a -> a -> a
forall a. Ord a => a -> a -> a
max a
a (a -> a -> a
forall a. Ord a => a -> a -> a
min a
b a
z)
{-# INLINE jointBincount #-}

rtf :: (Real a) => a -> Double
rtf :: forall a. Real a => a -> Double
rtf = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac
{-# NOINLINE [1] rtf #-}

{-# RULES
"rtf/Double" [2] forall (x :: Double). rtf x = x
    #-}