-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Transform.Fourier.CT
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- Cooley-Tukey algorithm for computing the FFT
--
-----------------------------------------------------------------------------

module Numeric.Transform.Fourier.CT (fft_ct1, fft_ct2) where

import Data.List (transpose)
import Data.Array
import Data.Complex

-- | Cooley-Tukey algorithm doing row FFT's then column FFT's

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

fft_ct1 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x[n]
	-> a -- ^ nrows
	-> a -- ^ ncols
	-> (Array a (Complex b) -> Array a (Complex b)) -- ^ FFT function
	-> Array a (Complex b) -- ^ X[k]

fft_ct1 :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b)
-> a
-> a
-> (Array a (Complex b) -> Array a (Complex b))
-> Array a (Complex b)
fft_ct1 Array a (Complex b)
a a
l a
m Array a (Complex b) -> Array a (Complex b)
fft = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
0,a
nforall a. Num a => a -> a -> a
-a
1) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a]
ks (forall i e. Array i e -> [e]
elems Array (a, a) (Complex b)
x')
    where x :: Array (a, a) (Complex b)
x = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) [ Array a (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!a
i | a
i <- [a]
xs ]
	  f :: Array (a, a) (Complex b)
f = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) (forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_rows forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
rows Array (a, a) (Complex b)
x)
	  g :: Array (a, a) (Complex b)
g = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) [ Array (a, a) (Complex b)
fforall i e. Ix i => Array i e -> i -> e
!(a
i,a
j) forall a. Num a => a -> a -> a
* Array a (Complex b)
wforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
*a
j) | a
i <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
j <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]
	  x' :: Array (a, a) (Complex b)
x' = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) (forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_cols forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
cols Array (a, a) (Complex b)
g)
	  wn :: Complex b
wn = 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. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n)
	  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 b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall a. Num a => a -> a -> a
* Complex b
wn) Complex b
1
	  ([a]
xs,[a]
ks) = forall a. Integral a => a -> a -> ([a], [a])
ct_index_map1 a
l a
m
	  n :: a
n = a
l forall a. Num a => a -> a -> a
* a
m

-- | Cooley-Tukey algorithm doing column FFT's then row FFT's

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

fft_ct2 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x[n]
	-> a -- ^ nrows
	-> a -- ^ ncols
	-> (Array a (Complex b) -> Array a (Complex b)) -- ^ fft function
	-> Array a (Complex b) -- ^ X[k]

fft_ct2 :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b)
-> a
-> a
-> (Array a (Complex b) -> Array a (Complex b))
-> Array a (Complex b)
fft_ct2 Array a (Complex b)
a a
l a
m Array a (Complex b) -> Array a (Complex b)
fft = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (a
0,a
nforall a. Num a => a -> a -> a
-a
1) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a]
ks (forall i e. Array i e -> [e]
elems Array (a, a) (Complex b)
x')
    where x :: Array (a, a) (Complex b)
x = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) [ Array a (Complex b)
aforall i e. Ix i => Array i e -> i -> e
!a
i | a
i <- [a]
xs ]
	  f :: Array (a, a) (Complex b)
f = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) (forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_cols forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
cols Array (a, a) (Complex b)
x)
	  g :: Array (a, a) (Complex b)
g = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) [ Array (a, a) (Complex b)
fforall i e. Ix i => Array i e -> i -> e
!(a
i,a
j) forall a. Num a => a -> a -> a
* Array a (Complex b)
wforall i e. Ix i => Array i e -> i -> e
!(a
iforall a. Num a => a -> a -> a
*a
j) | a
i <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
j <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]
	  x' :: Array (a, a) (Complex b)
x' = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray ((a
0,a
0),(a
lforall a. Num a => a -> a -> a
-a
1,a
mforall a. Num a => a -> a -> a
-a
1)) (forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_rows forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
rows Array (a, a) (Complex b)
g)
	  wn :: Complex b
wn = 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. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n)
	  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 b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall a. Num a => a -> a -> a
* Complex b
wn) Complex b
1
	  ([a]
xs,[a]
ks) = forall a. Integral a => a -> a -> ([a], [a])
ct_index_map2 a
l a
m
	  n :: a
n = a
l forall a. Num a => a -> a -> a
* a
m

-- Index maps

{-# specialize ct_index_map1 :: Int -> Int -> ([Int],[Int]) #-}

ct_index_map1 :: (Integral a) => a -> a -> ([a],[a])
ct_index_map1 :: forall a. Integral a => a -> a -> ([a], [a])
ct_index_map1 a
l a
m = ([a]
n,[a]
k)
    where n :: [a]
n = [ a
n1 forall a. Num a => a -> a -> a
+ a
l forall a. Num a => a -> a -> a
* a
n2 | a
n1 <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
n2 <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]
          k :: [a]
k = [ a
m forall a. Num a => a -> a -> a
* a
k1 forall a. Num a => a -> a -> a
+ a
k2 | a
k1 <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
k2 <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]

{-# specialize ct_index_map2 :: Int -> Int -> ([Int],[Int]) #-}

ct_index_map2 :: (Integral a) => a -> a -> ([a],[a])
ct_index_map2 :: forall a. Integral a => a -> a -> ([a], [a])
ct_index_map2 a
l a
m = ([a]
n,[a]
k)
    where n :: [a]
n = [ a
m forall a. Num a => a -> a -> a
* a
n1 forall a. Num a => a -> a -> a
+ a
n2 | a
n1 <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
n2 <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]
          k :: [a]
k = [ a
k1 forall a. Num a => a -> a -> a
+ a
l forall a. Num a => a -> a -> a
* a
k2 | a
k1 <- [a
0..(a
lforall a. Num a => a -> a -> a
-a
1)], a
k2 <- [a
0..(a
mforall a. Num a => a -> a -> a
-a
1)] ]

-- Auxilary functions (also used for PFA)

{-# specialize rows :: Array (Int,Int) (Complex Float) -> [Array Int (Complex Float)] #-}
{-# specialize rows :: Array (Int,Int) (Complex Double) -> [Array Int (Complex Double)] #-}

rows :: (Ix a, Integral a, RealFloat b) => Array (a,a) (Complex b) -> [Array a (Complex b)]
rows :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
rows Array (a, a) (Complex b)
x = [ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (a
0,a
m) [ Array (a, a) (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
i,a
j) | a
j <- [a
0..a
m] ] | a
i <- [a
0..a
l] ]
    where ((a
_,a
_),(a
l,a
m)) = forall i e. Array i e -> (i, i)
bounds Array (a, a) (Complex b)
x

{-# specialize cols :: Array (Int,Int) (Complex Float) -> [Array Int (Complex Float)] #-}
{-# specialize cols :: Array (Int,Int) (Complex Double) -> [Array Int (Complex Double)] #-}

cols :: (Ix a, Integral a, RealFloat b) => Array (a,a) (Complex b) -> [Array a (Complex b)]
cols :: forall a b.
(Ix a, Integral a, RealFloat b) =>
Array (a, a) (Complex b) -> [Array a (Complex b)]
cols Array (a, a) (Complex b)
x = [ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (a
0,a
l) [ Array (a, a) (Complex b)
xforall i e. Ix i => Array i e -> i -> e
!(a
i,a
j) | a
i <- [a
0..a
l] ] | a
j <- [a
0..a
m] ]
    where ((a
_,a
_),(a
l,a
m)) = forall i e. Array i e -> (i, i)
bounds Array (a, a) (Complex b)
x

{-# specialize flatten_rows :: [Array Int (Complex Float)] -> [(Complex Float)] #-}
{-# specialize flatten_rows :: [Array Int (Complex Double)] -> [(Complex Double)] #-}

flatten_rows :: (Ix a, Integral a, RealFloat b) => [Array a (Complex b)] -> [(Complex b)]
flatten_rows :: forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_rows [Array a (Complex b)]
a = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. [a] -> [a] -> [a]
(++) [] forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall i e. Array i e -> [e]
elems [Array a (Complex b)]
a

{-# specialize flatten_cols :: [Array Int (Complex Float)] -> [(Complex Float)] #-}
{-# specialize flatten_cols :: [Array Int (Complex Double)] -> [(Complex Double)] #-}

flatten_cols :: (Ix a, Integral a, RealFloat b) => [Array a (Complex b)] -> [(Complex b)]
flatten_cols :: forall a b.
(Ix a, Integral a, RealFloat b) =>
[Array a (Complex b)] -> [Complex b]
flatten_cols [Array a (Complex b)]
a = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. [a] -> [a] -> [a]
(++) [] forall a b. (a -> b) -> a -> b
$ forall a. [[a]] -> [[a]]
transpose forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall i e. Array i e -> [e]
elems [Array a (Complex b)]
a