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

module DataFrame.Internal.Statistics where

import qualified Data.Vector.Algorithms.Intro as VA
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 = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum ((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 b. (Real a, Fractional b) => a -> b
realToFrac 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 mean' #-}

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 b. (Real a, Fractional b) => a -> b
realToFrac 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 b. (Real a, Fractional b) => a -> b
realToFrac (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' #-}

-- accumulator: count, mean, m2
data VarAcc = VarAcc !Int !Double !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 :: (Real a) => VarAcc -> a -> VarAcc
varianceStep :: forall a. Real a => VarAcc -> a -> VarAcc
varianceStep (VarAcc !Int
n !Double
mean !Double
m2) !a
x =
    let !n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        !delta :: Double
delta = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
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
* (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
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 -- or error "variance of <2 samples"
    | 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 -> a -> VarAcc) -> VarAcc -> Vector a -> VarAcc
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' VarAcc -> a -> VarAcc
forall a. Real a => VarAcc -> a -> VarAcc
varianceStep (Int -> Double -> Double -> VarAcc
VarAcc Int
0 Double
0 Double
0)
{-# INLINE variance' #-}

-- accumulator: count, mean, m2, m3
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 :: SkewAcc -> Double -> SkewAcc
skewnessStep :: SkewAcc -> Double -> SkewAcc
skewnessStep (SkewAcc !Int
n !Double
mean !Double
m2 !Double
m3) !Double
x =
    let !n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        !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 -- or error "skewness of <3 samples"
    | 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.Vector Double -> Double
skewness' :: Vector Double -> Double
skewness' = SkewAcc -> Double
computeSkewness (SkewAcc -> Double)
-> (Vector Double -> SkewAcc) -> Vector Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SkewAcc -> Double -> SkewAcc)
-> SkewAcc -> Vector Double -> SkewAcc
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' SkewAcc -> Double -> SkewAcc
skewnessStep (Int -> Double -> Double -> Double -> SkewAcc
SkewAcc Int
0 Double
0 Double
0 Double
0)
{-# INLINE skewness' #-}

correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double
correlation' :: Vector Double -> Vector Double -> Maybe Double
correlation' 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
    | Int
nI Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = 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
nI
            (!Double
sumX, !Double
sumY, !Double
sumSquaredX, !Double
sumSquaredY, !Double
sumXY) = Int
-> Double
-> Double
-> Double
-> Double
-> Double
-> (Double, Double, Double, Double, Double)
go Int
0 Double
0 Double
0 Double
0 Double
0 Double
0
            !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
sumSquaredX 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
sumSquaredY 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
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
num Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den)
  where
    !nI :: Int
nI = Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
xs
    go :: Int
-> Double
-> Double
-> Double
-> Double
-> Double
-> (Double, Double, Double, Double, Double)
go !Int
i !Double
sumX !Double
sumY !Double
sumSquaredX !Double
sumSquaredY !Double
sumXY
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nI =
            let !x :: Double
x = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Double
xs Int
i
                !y :: Double
y = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Double
ys Int
i
                !sumX' :: Double
sumX' = Double
sumX Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x
                !sumY' :: Double
sumY' = Double
sumY Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y
                !sumSquaredX' :: Double
sumSquaredX' = Double
sumSquaredX Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
                !sumSquaredY' :: Double
sumSquaredY' = Double
sumSquaredY Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y
                !sumXY' :: Double
sumXY' = Double
sumXY Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y
             in Int
-> Double
-> Double
-> Double
-> Double
-> Double
-> (Double, Double, Double, Double, Double)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double
sumX' Double
sumY' Double
sumSquaredX' Double
sumSquaredY' Double
sumXY'
        | Bool
otherwise = (Double
sumX, Double
sumY, Double
sumSquaredX, Double
sumSquaredY, Double
sumXY)
{-# INLINE correlation' #-}

quantiles' :: VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double
quantiles' :: Vector Int -> Int -> Vector Double -> Vector Double
quantiles' Vector Int
qs Int
q Vector Double
samp
    | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Double
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 Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
samp
        MVector s Double
mutableSamp <- Vector Double -> ST s (MVector (PrimState (ST s)) Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector Double
samp
        MVector (PrimState (ST s)) Double -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VA.sort MVector s Double
MVector (PrimState (ST s)) Double
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 <- MVector (PrimState (ST s)) Double -> Int -> ST s Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s Double
MVector (PrimState (ST s)) Double
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 <- MVector (PrimState (ST s)) Double -> Int -> ST s Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.read MVector s Double
MVector (PrimState (ST s)) Double
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' #-}

interQuartileRange' :: VU.Vector Double -> Double
interQuartileRange' :: Vector Double -> Double
interQuartileRange' Vector Double
samp =
    let quartiles :: Vector Double
quartiles = Vector Int -> Int -> Vector Double -> Vector Double
quantiles' ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int
1, Int
3]) Int
4 Vector Double
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 = if Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Int
0 else if Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k then Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else 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 #-}