-----------------------------------------------------------------------------
-- |
-- Module      :  DSP.Covariance
-- 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-covariance
-- These formulas can be found in most DSP textbooks.
-- 
-- In the following routines, x and y are assumed to be of the same
-- length.
--
-----------------------------------------------------------------------------


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

-- TODO: Cxx(X) = Var(X), but I'm not sure how the lag works into that

module DSP.Covariance (cxy, cxy_b, cxy_u, cxx, cxx_b, cxx_u) where

import Data.Array
import Data.Complex

import DSP.Correlation
import Numeric.Statistics.Moment

-- | raw cross-covariance
--
-- We define covariance in terms of correlation.
--
-- Cxy(X,Y) = E[(X - E[X])(Y - E[Y])] 
--          = E[XY] - E[X]E[Y]
--          = Rxy(X,Y) - E[X]E[Y]

-- cxy x y k | k >= 0 = sum [ (x!(i+k) - xm) * ((conjugate (y!i)) - ym) | i <- [0..(n-1-k)] ]

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

cxy x y k | k >= 0 = rxy x y k - xm * ym
 	  | k < 0  = conjugate (cxy y x (-k))
    where xm = mean (elems x)
          ym = mean (map conjugate (elems y))
	  n = snd (bounds x) + 1

-- | raw auto-covariance
--
-- Cxx(X,X) = E[(X - E[X])(X - E[X])] 
--          = E[XX] - E[X]E[X]
--          = Rxy(X,X) - E[X]^2

-- We define this explicitly to prevent the mean from being calculated
-- twice.

-- cxx x k | k >= 0 = sum [ (x!(i+k) - xm) * (conjugate (x!i - xm)) | i <- [0..(n-1-k)] ]

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

cxx x k | k >= 0 = rxx x k - xm^2
	| k < 0  = conjugate (cxx x (-k))
    where xm = mean (elems x)
	  n = snd (bounds x) + 1

-- Define the biased and unbiased versions in terms of the above.

-- | biased cross-covariance

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

cxy_b x y k = (cxy x y k) / (fromIntegral n)
    where n = snd (bounds x) + 1

-- | unbiased cross-covariance

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

cxy_u x y k = (cxy x y k) / (fromIntegral (n-(abs k)))
    where n = snd (bounds x) + 1

-- | biased auto-covariance

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

cxx_b x k = (cxx x k) / (fromIntegral n)
    where n = snd (bounds x) + 1

-- | unbiased auto-covariance

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

cxx_u x k = (cxx x k) / (fromIntegral (n-(abs k)))
    where n = snd (bounds x) + 1