-----------------------------------------------------------------------------
-- |
-- Module      :  Matrix.Cholesky
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- This module contains an implementation of Levinson-Durbin recursion.
--
-----------------------------------------------------------------------------

module Matrix.Levinson (levinson) where

import Data.Array
import Data.Complex

-- * Functions

-- Section 6.3.3 in Kay, formulas 6.46--6.48

-- TODO: rho is typing as complex, but it is real
-- TODO: add stepdown function
-- TODO: some applications may want all model estimations from [1..p]

-- | levinson takes an array, r, of autocorrelation values, and a
-- model order, p, and returns an array, a, of the model estimate and
-- rho, the noise power.

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

levinson :: 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 = (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] ]
	  rho :: Array a (Complex b)
rho = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
1,a
p) [ (a
k, a -> Complex b
rhok a
k) | a
k <- [a
1..a
p] ]
	  ak :: a -> a -> Complex b
ak a
1 a
1             = -Array a (Complex b)
rforall i e. Ix i => Array i e -> i -> e
!a
1 forall a. Fractional a => a -> a -> a
/ Array a (Complex b)
rforall i e. Ix i => Array i e -> i -> e
!a
0
	  ak a
k a
i | a
kforall a. Eq a => a -> a -> Bool
==a
i      = -(Array a (Complex b)
rforall i e. Ix i => Array i e -> i -> e
!a
k forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ 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
l) forall a. Num a => a -> a -> a
* Array a (Complex b)
rforall i e. Ix i => Array i e -> i -> e
!(a
kforall a. Num a => a -> a -> a
-a
l) | a
l <- [a
1..(a
kforall a. Num a => a -> a -> a
-a
1)] ]) forall a. Fractional 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)
		 | 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, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
k,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)))
	  rhok :: a -> Complex b
rhok a
1 = (Complex b
1 forall a. Num a => a -> a -> a
- (forall a. Num a => a -> a
abs (Array (a, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
1,a
1)))forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)) forall a. Num a => a -> a -> a
* Array a (Complex b)
rforall i e. Ix i => Array i e -> i -> e
!a
0
	  rhok a
k = (Complex b
1 forall a. Num a => a -> a -> a
- (forall a. Num a => a -> a
abs (Array (a, a) (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!(a
k,a
k)))forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)) 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)

-- r = array (0,2) [ (0, (2.0 :+ 0.0)), (1, ((-1.0) :+ 1.0)), (2, (0.0 :+ 0.0)) ]
-- a = fst (levinson r 2)

-- verify = a == array (1,2) [(1,1.0 :+ (-1.0)),(2,0.0 :+ (-1.0))]