-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Transform.Fourier.DFT
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- Not so naive implementation of a Discrete Fourier Transform.
--
-----------------------------------------------------------------------------

{-
We cheat in three ways from a direct translation of the DFT equation:

     X(k) = sum(n=0..N-1) x(n) * e^(-2*j*pi*n*k/N)

1.  We precompute all values of W_N, and exploit the periodicity.
This is just to cut down on the number of sin/cos calls.

2.  We calculate X(0) seperately to prevent multiplication by 1

3.  We factor out x(0) to prevent multiplication by 1
-}

module Numeric.Transform.Fourier.DFT (dft) where

import Data.Array
import Data.Complex

-- We use a helper function here because we may want to have special
-- cases for small DFT's and we want to precompute the suspension all of
-- the twiddle factors.

{-# specialize dft :: Array Int (Complex Float) -> Array Int (Complex Float) #-}
{-# specialize dft :: Array Int (Complex Double) -> Array Int (Complex Double) #-}

dft :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x[n]
    -> Array a (Complex b) -- ^ X[k]
dft :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
dft Array a (Complex b)
a = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b)
-> Array a (Complex b) -> a -> Array a (Complex b)
dft' Array a (Complex b)
a Array a (Complex b)
w a
n
    where w :: Array a (Complex b)
w = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (a
0,a
nforall a. Num a => a -> a -> a
-a
1) [ forall a. Floating a => a -> Complex a
cis (-b
2 forall a. Num a => a -> a -> a
* forall a. Floating a => a
pi forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n) | a
i <- [a
0..(a
nforall a. Num a => a -> a -> a
-a
0)] ]
	  n :: a
n = forall a b. (a, b) -> b
snd (forall i e. Array i e -> (i, i)
bounds Array a (Complex b)
a) forall a. Num a => a -> a -> a
+ a
1

{-# specialize dft' :: Array Int (Complex Float) -> Array Int (Complex Float) -> Int -> Array Int (Complex Float) #-}
{-# specialize dft' :: Array Int (Complex Double) -> Array Int (Complex Double) -> Int -> Array Int (Complex Double) #-}

dft' :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -> Array a (Complex b) -> a -> Array a (Complex b)
dft' :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b)
-> Array a (Complex b) -> a -> Array a (Complex b)
dft' Array a (Complex b)
a Array a (Complex b)
_ a
1 = Array a (Complex b)
a
dft' Array a (Complex b)
a Array a (Complex b)
w a
n = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (a
0,a
nforall a. Num a => a -> a -> a
-a
1) (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 | a
k <- [a
0..(a
nforall a. Num a => a -> a -> a
-a
1)] ] forall a. a -> [a] -> [a]
: [ Array a (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!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
wik a
i a
k | a
k <- [a
1..(a
nforall a. Num a => a -> a -> a
-a
1)] ] | a
i <- [a
1..(a
nforall a. Num a => a -> a -> a
-a
1)] ])
    where wik :: a -> a -> Complex b
wik a
0 a
_ = Complex b
1
          wik a
_ a
0 = Complex b
1
          wik a
i a
k = Array a (Complex b)
wforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
*a
k forall a. Integral a => a -> a -> a
`mod` a
n)