{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
-- |
-- Module      : Statistics.Correlation.Pearson
--
module Statistics.Correlation
    ( -- * Pearson correlation
      pearson
    , pearsonMatByRow
      -- * Spearman correlation
    , spearman
    , spearmanMatByRow
    ) where

import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import Statistics.Matrix
import Statistics.Sample
import Statistics.Test.Internal (rankUnsorted)


----------------------------------------------------------------
-- Pearson
----------------------------------------------------------------

-- | Pearson correlation for sample of pairs. Exactly same as
-- 'Statistics.Sample.correlation'
pearson :: (G.Vector v (Double, Double), G.Vector v Double)
        => v (Double, Double) -> Double
pearson :: v (Double, Double) -> Double
pearson = v (Double, Double) -> Double
forall (v :: * -> *).
(Vector v (Double, Double), Vector v Double) =>
v (Double, Double) -> Double
correlation
{-# INLINE pearson #-}

-- | Compute pairwise pearson correlation between rows of a matrix
pearsonMatByRow :: Matrix -> Matrix
pearsonMatByRow :: Matrix -> Matrix
pearsonMatByRow Matrix
m
  = Int -> (Int -> Int -> Double) -> Matrix
generateSym (Matrix -> Int
rows Matrix
m)
      (\Int
i Int
j -> Vector (Double, Double) -> Double
forall (v :: * -> *).
(Vector v (Double, Double), Vector v Double) =>
v (Double, Double) -> Double
pearson (Vector (Double, Double) -> Double)
-> Vector (Double, Double) -> Double
forall a b. (a -> b) -> a -> b
$ Matrix -> Int -> Vector
row Matrix
m Int
i Vector -> Vector -> Vector (Double, Double)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
`U.zip` Matrix -> Int -> Vector
row Matrix
m Int
j)
{-# INLINE pearsonMatByRow #-}



----------------------------------------------------------------
-- Spearman
----------------------------------------------------------------

-- | compute spearman correlation between two samples
spearman :: ( Ord a
            , Ord b
            , G.Vector v a
            , G.Vector v b
            , G.Vector v (a, b)
            , G.Vector v Int
            , G.Vector v Double
            , G.Vector v (Double, Double)
            , G.Vector v (Int, a)
            , G.Vector v (Int, b)
            )
         => v (a, b)
         -> Double
spearman :: v (a, b) -> Double
spearman v (a, b)
xy
  = v (Double, Double) -> Double
forall (v :: * -> *).
(Vector v (Double, Double), Vector v Double) =>
v (Double, Double) -> Double
pearson
  (v (Double, Double) -> Double) -> v (Double, Double) -> Double
forall a b. (a -> b) -> a -> b
$ v Double -> v Double -> v (Double, Double)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v a -> v b -> v (a, b)
G.zip (v a -> v Double
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int, Vector v Double,
 Vector v (Int, a)) =>
v a -> v Double
rankUnsorted v a
x) (v b -> v Double
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int, Vector v Double,
 Vector v (Int, a)) =>
v a -> v Double
rankUnsorted v b
y)
  where
    (v a
x, v b
y) = v (a, b) -> (v a, v b)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v (a, b) -> (v a, v b)
G.unzip v (a, b)
xy
{-# INLINE spearman #-}

-- | compute pairwise spearman correlation between rows of a matrix
spearmanMatByRow :: Matrix -> Matrix
spearmanMatByRow :: Matrix -> Matrix
spearmanMatByRow
  = Matrix -> Matrix
pearsonMatByRow (Matrix -> Matrix) -> (Matrix -> Matrix) -> Matrix -> Matrix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vector] -> Matrix
fromRows ([Vector] -> Matrix) -> (Matrix -> [Vector]) -> Matrix -> Matrix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector -> Vector) -> [Vector] -> [Vector]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Vector -> Vector
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int, Vector v Double,
 Vector v (Int, a)) =>
v a -> v Double
rankUnsorted ([Vector] -> [Vector])
-> (Matrix -> [Vector]) -> Matrix -> [Vector]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix -> [Vector]
toRows
{-# INLINE spearmanMatByRow #-}