-----------------------------------------------------------------------------
-- |
-- Module  :  ForSyDe.DFT
-- Copyright   :  (c) ForSyDe Group, KTH 2007-2008
-- License     :  BSD-style (see the file LICENSE)
-- 
-- Maintainer  :  forsyde-dev@ict.kth.se
-- Stability   :  experimental
-- Portability :  portable
--
-- This module includes the standard Discrete Fourier Transform (DFT)
-- function, and a fast Fourier transform (FFT) algorithm, for
-- computing the DFT, when the input vectors' length is a power of 2.
-----------------------------------------------------------------------------
module ForSyDe.Shallow.Utility.DFT(dft, fft) where

import ForSyDe.Shallow.Core.Vector
import Data.Complex

-- | The function 'dft' performs a standard Discrete Fourier Transformation
dft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
dft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
dft Int
bigN Vector (Complex Double)
x | Int
bigN Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Vector (Complex Double) -> Int
forall a. Vector a -> Int
lengthV Vector (Complex Double)
x) = (Complex Double -> Complex Double)
-> Vector (Complex Double) -> Vector (Complex Double)
forall a b. (a -> b) -> Vector a -> Vector b
mapV (Int -> Vector (Complex Double) -> Complex Double -> Complex Double
bigX_k Int
bigN Vector (Complex Double)
x) (Vector (Complex Double) -> Vector (Complex Double)
forall b a. Num b => Vector a -> Vector b
nVector Vector (Complex Double)
x)
           | Bool
otherwise = [Char] -> Vector (Complex Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"DFT: Vector has not the right size!"   
  where
    nVector :: Vector a -> Vector b
nVector Vector a
x'        = Int -> (b -> b) -> b -> Vector b
forall a b. (Num a, Eq a) => a -> (b -> b) -> b -> Vector b
iterateV (Vector a -> Int
forall a. Vector a -> Int
lengthV Vector a
x') (b -> b -> b
forall a. Num a => a -> a -> a
+b
1) b
0
    bigX_k :: Int -> Vector (Complex Double) -> Complex Double -> Complex Double
bigX_k Int
bigN' Vector (Complex Double)
x' Complex Double
k = Vector (Complex Double) -> Complex Double
sumV ((Complex Double -> Complex Double -> Complex Double)
-> Vector (Complex Double)
-> Vector (Complex Double)
-> Vector (Complex Double)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWithV Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
(*) Vector (Complex Double)
x' (Complex Double -> Int -> Vector (Complex Double)
bigW' Complex Double
k Int
bigN'))
    bigW' :: Complex Double -> Int -> Vector (Complex Double)
bigW' Complex Double
k' Int
bigN'    = (Complex Double -> Complex Double)
-> Vector (Complex Double) -> Vector (Complex Double)
forall a b. (a -> b) -> Vector a -> Vector b
mapV (Complex Double -> Complex Double -> Complex Double
forall a. Floating a => a -> a -> a
** Complex Double
k') ((Double -> Complex Double)
-> Vector Double -> Vector (Complex Double)
forall a b. (a -> b) -> Vector a -> Vector b
mapV Double -> Complex Double
forall a. Floating a => a -> Complex a
cis (Int -> Vector Double
fullcircle Int
bigN'))
    sumV :: Vector (Complex Double) -> Complex Double
sumV              = (Complex Double -> Complex Double -> Complex Double)
-> Complex Double -> Vector (Complex Double) -> Complex Double
forall a b. (a -> b -> a) -> a -> Vector b -> a
foldlV Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
(+) (Double
0Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0)

fullcircle :: Int -> Vector Double 
fullcircle :: Int -> Vector Double
fullcircle Int
n = Double -> Double -> Int -> Vector Double
forall a t.
(Floating a, Integral t, Eq a) =>
a -> a -> t -> Vector a
fullcircle1 Double
0 (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Int
n
  where
    fullcircle1 :: a -> a -> t -> Vector a
fullcircle1 a
l a
m t
n' 
      | a
l a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
m    = Vector a
forall a. Vector a
NullV
      | Bool
otherwise = -a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
forall a. Floating a => a
pia -> a -> a
forall a. Num a => a -> a -> a
*a
la -> a -> a
forall a. Fractional a => a -> a -> a
/(t -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
n') 
                    a -> Vector a -> Vector a
forall a. a -> Vector a -> Vector a
:> a -> a -> t -> Vector a
fullcircle1 (a
la -> a -> a
forall a. Num a => a -> a -> a
+a
1) a
m t
n' 

-- | The function 'fft' implements a fast Fourier transform (FFT)
-- algorithm, for computing the DFT, when the size N is a power of 2.
fft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
fft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
fft Int
bigN Vector (Complex Double)
xv | Int
bigN Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Vector (Complex Double) -> Int
forall a. Vector a -> Int
lengthV Vector (Complex Double)
xv) = (Int -> Complex Double) -> Vector Int -> Vector (Complex Double)
forall a b. (a -> b) -> Vector a -> Vector b
mapV (Vector (Complex Double) -> Int -> Complex Double
bigX Vector (Complex Double)
xv) (Int -> Vector Int
forall b a. (Num b, Num a, Eq a) => a -> Vector b
kVector Int
bigN)
            | Bool
otherwise = [Char] -> Vector (Complex Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"FFT: Vector has not the right size!"

kVector :: (Num b, Num a, Eq a) => a -> Vector b  
kVector :: a -> Vector b
kVector a
bigN = a -> (b -> b) -> b -> Vector b
forall a b. (Num a, Eq a) => a -> (b -> b) -> b -> Vector b
iterateV a
bigN (b -> b -> b
forall a. Num a => a -> a -> a
+b
1) b
0 

bigX :: Vector (Complex Double) -> Int -> Complex Double
bigX :: Vector (Complex Double) -> Int -> Complex Double
bigX (Complex Double
x0:>Complex Double
x1:>Vector (Complex Double)
NullV) Int
k | Int -> Bool
forall a. Integral a => a -> Bool
even Int
k = Complex Double
x0 Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
+ Complex Double
x1 Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Int -> Int -> Complex Double
bigW Int
2 Int
0
                       | Int -> Bool
forall a. Integral a => a -> Bool
odd Int
k  = Complex Double
x0 Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
- Complex Double
x1 Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Int -> Int -> Complex Double
bigW Int
2 Int
0
bigX Vector (Complex Double)
xv Int
k = Int -> Complex Double
bigF_even Int
k Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
+ Int -> Complex Double
bigF_odd Int
k Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Int -> Int -> Complex Double
bigW Int
bigN (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k)
  where bigF_even :: Int -> Complex Double
bigF_even Int
k' = Vector (Complex Double) -> Int -> Complex Double
bigX (Vector (Complex Double) -> Vector (Complex Double)
forall a. Vector a -> Vector a
evens Vector (Complex Double)
xv) Int
k'
        bigF_odd :: Int -> Complex Double
bigF_odd Int
k' = Vector (Complex Double) -> Int -> Complex Double
bigX (Vector (Complex Double) -> Vector (Complex Double)
forall a. Vector a -> Vector a
odds Vector (Complex Double)
xv) Int
k'
        bigN :: Int
bigN = Vector (Complex Double) -> Int
forall a. Vector a -> Int
lengthV Vector (Complex Double)
xv

bigW :: Int -> Int -> Complex Double
bigW :: Int -> Int -> Complex Double
bigW Int
bigN Int
k = Double -> Complex Double
forall a. Floating a => a -> Complex a
cis (-Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigN))

evens :: Vector a -> Vector a
evens :: Vector a -> Vector a
evens Vector a
NullV   = Vector a
forall a. Vector a
NullV
evens (a
v1:>Vector a
NullV) = a
v1 a -> Vector a -> Vector a
forall a. a -> Vector a -> Vector a
:> Vector a
forall a. Vector a
NullV
evens (a
v1:>a
_:>Vector a
v)  = a
v1 a -> Vector a -> Vector a
forall a. a -> Vector a -> Vector a
:> Vector a -> Vector a
forall a. Vector a -> Vector a
evens Vector a
v

odds :: Vector a -> Vector a
odds :: Vector a -> Vector a
odds Vector a
NullV    = Vector a
forall a. Vector a
NullV
odds (a
_:>Vector a
NullV)   = Vector a
forall a. Vector a
NullV
odds (a
_:>a
v2:>Vector a
v)   = a
v2 a -> Vector a -> Vector a
forall a. a -> Vector a -> Vector a
:> Vector a -> Vector a
forall a. Vector a -> Vector a
odds Vector a
v