-----------------------------------------------------------------------------
-- |
-- Module      :  DSP.Correlation
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- This module contains routines to perform cross- and auto-correlation.
-- These formulas can be found in most DSP textbooks.
--
-- In the following routines, x and y are assumed to be of the same
-- length.
--
-----------------------------------------------------------------------------

module DSP.Correlation (rxy, rxy_b, rxy_u, rxx, rxx_b, rxx_u, test) where

import Data.Array
import Data.Complex

-- * Functions

-- TODO: fix these routines to handle the case were x and y are different
-- lengths.

-- | raw cross-correllation

rxy :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                 -> Array a (Complex b) -- ^ y
                                 -> a                   -- ^ k
                                 -> Complex b           -- ^ R_xy[k]

rxy :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy Array a (Complex b)
x Array a (Complex b)
y a
k =
   if a
k forall a. Ord a => a -> a -> Bool
>= a
0
     then forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
+a
k) forall a. Num a => a -> a -> a
* forall a. Num a => Complex a -> Complex a
conjugate (Array a (Complex b)
yforall i e. Ix i => Array i e -> i -> e
!a
i) | a
i <- [a
0..(a
nforall a. Num a => a -> a -> a
-a
1forall a. Num a => a -> a -> a
-a
k)] ]
     else forall a. Num a => Complex a -> Complex a
conjugate (forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy Array a (Complex b)
y Array a (Complex b)
x (-a
k))
    where n :: a
n = forall a b. (a, b) -> b
snd (forall i e. Array i e -> (i, i)
bounds Array a (Complex b)
x) forall a. Num a => a -> a -> a
+ a
1

-- | biased cross-correllation

rxy_b :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                   -> Array a (Complex b) -- ^ y
                                   -> a                   -- ^ k
                                   -> Complex b           -- ^ R_xy[k] \/ N

rxy_b :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy_b Array a (Complex b)
x Array a (Complex b)
y a
k = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy Array a (Complex b)
x Array a (Complex b)
y a
k forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n
    where n :: a
n = forall a b. (a, b) -> b
snd (forall i e. Array i e -> (i, i)
bounds Array a (Complex b)
x) forall a. Num a => a -> a -> a
+ a
1

-- | unbiased cross-correllation

rxy_u :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                   -> Array a (Complex b) -- ^ y
                                   -> a                   -- ^ k
                                   -> Complex b           -- ^ R_xy[k] \/ (N-k)

rxy_u :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy_u Array a (Complex b)
x Array a (Complex b)
y a
k = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy Array a (Complex b)
x Array a (Complex b)
y a
k forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
n forall a. Num a => a -> a -> a
- forall a. Num a => a -> a
abs a
k)
    where n :: a
n = forall a b. (a, b) -> b
snd (forall i e. Array i e -> (i, i)
bounds Array a (Complex b)
x) forall a. Num a => a -> a -> a
+ a
1

-- autocorrellation

-- We define autocorrelation in terms of the cross correlation routines.

-- | raw auto-correllation

rxx :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                 -> a                   -- ^ k
                                 -> Complex b           -- ^ R_xx[k]

rxx :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> Complex b
rxx   Array a (Complex b)
x a
k = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy   Array a (Complex b)
x Array a (Complex b)
x a
k

-- | biased auto-correllation

rxx_b :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                   -> a                   -- ^ k
                                   -> Complex b           -- ^ R_xx[k] \/ N

rxx_b :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> Complex b
rxx_b Array a (Complex b)
x a
k = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy_b Array a (Complex b)
x Array a (Complex b)
x a
k

-- | unbiased auto-correllation

rxx_u :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x
                                   -> a                   -- ^ k
                                   -> Complex b           -- ^ R_xx[k] \/ (N-k)

rxx_u :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> Complex b
rxx_u Array a (Complex b)
x a
k = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy_u Array a (Complex b)
x Array a (Complex b)
x a
k

----------------------------------------------------------------------------
-- test routines
----------------------------------------------------------------------------

xt, yt :: Array Int (Complex Double)
xt :: Array Int (Complex Double)
xt = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
0,Int
4)
   [ (Int
0, Double
1 forall a. a -> a -> Complex a
:+ Double
0),
     (Int
1, Double
0 forall a. a -> a -> Complex a
:+ Double
1),
     (Int
2, (-Double
1) forall a. a -> a -> Complex a
:+  Double
0),
     (Int
3, Double
0 forall a. a -> a -> Complex a
:+ (-Double
1)),
     (Int
4, Double
1 forall a. a -> a -> Complex a
:+ Double
0) ]

yt :: Array Int (Complex Double)
yt = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
0,Int
4)
   [ (Int
0, Double
1 forall a. a -> a -> Complex a
:+ Double
0),
     (Int
1, (-Double
1) forall a. a -> a -> Complex a
:+ Double
0),
     (Int
2, Double
1 forall a. a -> a -> Complex a
:+ Double
0),
     (Int
3, (-Double
1) forall a. a -> a -> Complex a
:+ Double
0),
     (Int
4, Double
1 forall a. a -> a -> Complex a
:+ Double
0) ]

rt :: [Complex Double]
rt :: [Complex Double]
rt = forall a b. (a -> b) -> [a] -> [b]
map (forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b) -> a -> Complex b
rxy_b Array Int (Complex Double)
xt Array Int (Complex Double)
yt) [ Int
0, Int
1, Int
2 ]

test :: Bool
test :: Bool
test  =  [Complex Double]
rt forall a. Eq a => a -> a -> Bool
== [ (Double
0.2 forall a. a -> a -> Complex a
:+ Double
0.0), (Double
0.0 forall a. a -> a -> Complex a
:+ Double
0.0), (Double
0.0 forall a. a -> a -> Complex a
:+ Double
0.2) ]