-----------------------------------------------------------------------------
-- |
-- Module      :  DSP.Estimation.Spectral.AR
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- This module contains a few algorithms for AR parameter estimation.
-- Algorithms are taken from Steven M. Kay, /Modern Spectral Estimation:
-- Theory and Application/, which is one of the standard texts on the
-- subject.  When possible, variable conventions are the same in the code
-- as they are found in the text.
--
-----------------------------------------------------------------------------

module DSP.Estimation.Spectral.AR where

import DSP.Correlation
import Matrix.Levinson
import Matrix.Cholesky

import DSP.Basic((^!))
import Data.Array
import Data.Complex

-- * Functions

-------------------------------------------------------------------------------
-- ar_yw x p
-------------------------------------------------------------------------------

-- Section 7.3 in Kay

-- | Computes an AR(p) model estimate from x using the Yule-Walker method

ar_yw :: (Ix a, Integral a, RealFloat b) => Array a (Complex b)      -- ^ x
      -> a                        -- ^ p
      -> (Array a (Complex b), b) -- ^ (a,rho)

ar_yw :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> (Array a (Complex b), b)
ar_yw Array a (Complex b)
x a
p = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> (Array a (Complex b), b)
levinson Array a (Complex b)
r a
p
    where r :: Array a (Complex b)
r = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
0,a
p) [ (a
k, 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) | a
k <- [a
0..a
p] ]

-------------------------------------------------------------------------------
-- ar_cov x p
-------------------------------------------------------------------------------

-- Section 7.4 in Kay, but I factored out the 1/(N-p) term, and only
-- generate the lower triangle of cxx

-- TODO: use modified Prony method instead of matrix solver

-- | Computes an AR(p) model estimate from x using the covariance method

ar_cov :: (Ix a, Integral a, RealFloat b) => Array a (Complex b)      -- ^ x
       -> a                        -- ^ p
       -> (Array a (Complex b), b) -- ^ (a,rho)

ar_cov :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> (Array a (Complex b), b)
ar_cov Array a (Complex b)
x a
p = (Array a (Complex b)
a, b
sig2  forall a. Fractional a => a -> a -> a
/ (forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
nforall a. Num a => a -> a -> a
-a
p)))
    where a :: Array a (Complex b)
a = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b)
-> Array a (Complex b) -> Array a (Complex b)
cholesky Array (a, a) (Complex b)
m Array a (Complex b)
v
 	  sig2 :: b
sig2 = forall a. Complex a -> a
realPart ((a -> a -> Complex b
cxx a
0 a
0) forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Array a (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!a
k forall a. Num a => a -> a -> a
* (a -> a -> Complex b
cxx a
0 a
k) | a
k <- [a
1..a
p] ])
	  m :: Array (a, a) (Complex b)
m = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((a
1,a
1),(a
p,a
p)) [ ((a
j,a
k), a -> a -> Complex b
cxx a
j a
k) | a
j <- [a
1..a
p], a
k <- [a
1..a
j] ]
	  v :: Array a (Complex b)
v = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
1,a
p) [ (a
j, -(a -> a -> Complex b
cxx a
j a
0)) | a
j <- [a
1..a
p] ]
	  cxx :: a -> a -> Complex b
cxx a
j a
k = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ (forall a. Num a => Complex a -> Complex a
conjugate (Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
-a
j))) forall a. Num a => a -> a -> a
* Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
-a
k) | a
i <- [a
p..(a
nforall a. Num a => a -> a -> a
-a
1)] ]
	  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

-------------------------------------------------------------------------------
-- ar_mcov x p
-------------------------------------------------------------------------------

-- Section 7.5 in Kay, but I factored out the 1/(2(N-p)) term, and only
-- generate the lower triangle of cxx

-- | Computes an AR(p) model estimate from x using the modified covariance method

ar_mcov :: (Ix a, Integral a, RealFloat b) => Array a (Complex b)      -- ^ x
        -> a                        -- ^ p
        -> (Array a (Complex b), b) -- ^ (a,rho)

ar_mcov :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> (Array a (Complex b), b)
ar_mcov Array a (Complex b)
x a
p = (Array a (Complex b)
a, b
sig2  forall a. Fractional a => a -> a -> a
/ (forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
2forall a. Num a => a -> a -> a
*(a
nforall a. Num a => a -> a -> a
-a
p))))
    where a :: Array a (Complex b)
a = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b)
-> Array a (Complex b) -> Array a (Complex b)
cholesky Array (a, a) (Complex b)
m Array a (Complex b)
v
	  sig2 :: b
sig2 = forall a. Complex a -> a
realPart ((a -> a -> Complex b
cxx a
0 a
0) forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Array a (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!a
k forall a. Num a => a -> a -> a
* (a -> a -> Complex b
cxx a
0 a
k) | a
k <- [a
1..a
p] ])
	  m :: Array (a, a) (Complex b)
m = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((a
1,a
1),(a
p,a
p)) [ ((a
j,a
k), a -> a -> Complex b
cxx a
j a
k) | a
j <- [a
1..a
p], a
k <- [a
1..a
j] ]
	  v :: Array a (Complex b)
v = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
1,a
p) [ (a
j, -(a -> a -> Complex b
cxx a
j a
0)) | a
j <- [a
1..a
p] ]
	  cxx :: a -> a -> Complex b
cxx a
j a
k = (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ (forall a. Num a => Complex a -> Complex a
conjugate (Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
-a
j))) forall a. Num a => a -> a -> a
* Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
-a
k) | a
i <- [a
p..(a
nforall a. Num a => a -> a -> a
-a
1)] ] forall a. Num a => a -> a -> a
+ 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
j) forall a. Num a => a -> a -> a
* (forall a. Num a => Complex a -> Complex a
conjugate (Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
+a
k))) | a
i <- [a
0..(a
nforall a. Num a => a -> a -> a
-a
1forall a. Num a => a -> a -> a
-a
p)] ])
	  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

-------------------------------------------------------------------------------
-- ar_burg x p
-------------------------------------------------------------------------------

-- Section 7.6 in Kay

-- TODO: rho doesn't need to be an array
-- TODO: kk doesn't need to be an array
-- TODO: ef and eb don't need to be 2-D arrays

-- | Computes an AR(p) model estimate from x using the Burg' method

ar_burg :: (Ix a, Integral a, RealFloat b) => Array a (Complex b)      -- ^ x
        -> a                        -- ^ p
        -> (Array a (Complex b), b) -- ^ (a,rho)

ar_burg :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> (Array a (Complex b), b)
ar_burg Array a (Complex b)
x a
p = (forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
1,a
p) [ (a
k, Array (a, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
p,a
k)) | a
k <- [a
1..a
p] ], forall a. Complex a -> a
realPart (Array a (Complex b)
rhoforall i e. Ix i => Array i e -> i -> e
!a
p))
    where a :: Array (a, a) (Complex b)
a = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((a
1,a
1),(a
p,a
p)) [ ((a
k,a
i), a -> a -> Complex b
ak a
k a
i) | a
k <- [a
1..a
p], a
i <- [a
1..a
k] ]
	  ak :: a -> a -> Complex b
ak a
k a
i | a
iforall a. Eq a => a -> a -> Bool
==a
k      = Array a (Complex b)
kkforall i e. Ix i => Array i e -> i -> e
!a
k
		 | Bool
otherwise = Array (a, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
i) forall a. Num a => a -> a -> a
+ Array a (Complex b)
kkforall i e. Ix i => Array i e -> i -> e
!a
k forall a. Num a => a -> a -> a
* (forall a. Num a => Complex a -> Complex a
conjugate (Array (a, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
kforall a. Num a => a -> a -> a
-a
i)))
	  kk :: Array a (Complex b)
kk = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
1,a
p) [ (a
k, -Complex b
2 forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Array (a, a) (Complex b)
efforall i e. Ix i => Array i e -> i -> e
!((a
kforall a. Num a => a -> a -> a
-a
1),a
i) forall a. Num a => a -> a -> a
* (forall a. Num a => Complex a -> Complex a
conjugate (Array (a, a) (Complex b)
ebforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
iforall a. Num a => a -> a -> a
-a
1))) | a
i <- [a
k..(a
nforall a. Num a => a -> a -> a
-a
1)] ] forall a. Fractional a => a -> a -> a
/ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ forall a. Num a => a -> a
abs (Array (a, a) (Complex b)
efforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
i)) forall a. Num a => a -> Int -> a
^! Int
2 forall a. Num a => a -> a -> a
+ forall a. Num a => a -> a
abs (Array (a, a) (Complex b)
ebforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
iforall a. Num a => a -> a -> a
-a
1)) forall a. Num a => a -> Int -> a
^! Int
2 | a
i <- [a
k..(a
nforall a. Num a => a -> a -> a
-a
1)] ]) | a
k <- [a
1..a
p] ]
	  rho :: Array a (Complex b)
rho = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
0,a
p) ((a
0, forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> a -> Complex b
rxx_b Array a (Complex b)
x a
0) forall a. a -> [a] -> [a]
: [ (a
k, (Complex b
1 forall a. Num a => a -> a -> a
- forall a. Num a => a -> a
abs (Array a (Complex b)
kkforall i e. Ix i => Array i e -> i -> e
!a
k) forall a. Num a => a -> Int -> a
^! Int
2) forall a. Num a => a -> a -> a
* Array a (Complex b)
rhoforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1)) | a
k <- [a
1..a
p] ])
	  ef :: Array (a, a) (Complex b)
ef = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((a
0,a
1),(a
p,a
nforall a. Num a => a -> a -> a
-a
1)) [ ((a
k,a
i), a -> a -> Complex b
efki a
k a
i) | a
k <- [a
0..a
p], a
i <- [(a
kforall a. Num a => a -> a -> a
+a
1)..(a
nforall a. Num a => a -> a -> a
-a
1)] ]
	  eb :: Array (a, a) (Complex b)
eb = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((a
0,a
0),(a
p,a
nforall a. Num a => a -> a -> a
-a
2)) [ ((a
k,a
i), a -> a -> Complex b
ebki a
k a
i) | a
k <- [a
0..a
p], a
i <- [a
k..(a
nforall a. Num a => a -> a -> a
-a
2)] ]
	  efki :: a -> a -> Complex b
efki a
0 a
i = Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!a
i
	  efki a
k a
i = Array (a, a) (Complex b)
efforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
i) forall a. Num a => a -> a -> a
+ Array a (Complex b)
kkforall i e. Ix i => Array i e -> i -> e
!a
k forall a. Num a => a -> a -> a
* Array (a, a) (Complex b)
ebforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
iforall a. Num a => a -> a -> a
-a
1)
	  ebki :: a -> a -> Complex b
ebki a
0 a
i = Array a (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!a
i
	  ebki a
k a
i = Array (a, a) (Complex b)
ebforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
iforall a. Num a => a -> a -> a
-a
1) forall a. Num a => a -> a -> a
+ (forall a. Num a => Complex a -> Complex a
conjugate (Array a (Complex b)
kkforall i e. Ix i => Array i e -> i -> e
!a
k)) forall a. Num a => a -> a -> a
* Array (a, a) (Complex b)
efforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
1,a
i)
	  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

-------------------------------------------------------------------------------
-- ar_rmle x p
-------------------------------------------------------------------------------