    {-# 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 :: forall (v :: * -> *).
    (Vector v (Double, Double), Vector v Double) =>
    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 :: forall a b (v :: * -> *).
    (Ord a, Ord b, Vector v a, Vector v b, Vector v (a, b),
     Vector v Int, Vector v Double, Vector v (Double, Double),
     Vector v (Int, a), Vector v (Int, b)) =>
    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 a b. (a -> b) -> [a] -> [b]
    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 #-}