|
1 | 1 | import Data.Complex
|
2 |
| -import Data.Array |
3 |
| -import Data.Ratio |
| 2 | +import Data.List (partition) |
| 3 | +import Data.Map ((!)) |
4 | 4 | import qualified Data.Map as M
|
| 5 | +import Data.Ratio |
| 6 | + |
| 7 | +dft :: [Complex Double] -> [Complex Double] |
| 8 | +dft x = matMult dftMat x |
| 9 | + where |
| 10 | + n = length x |
| 11 | + w = exp $ (-2) * pi * (0 :+ 1) / fromIntegral n |
| 12 | + dftMat = [[w ^ (j * k) | j <- [0 .. n - 1]] | k <- [0 .. n - 1]] |
| 13 | + matMult m x = map (sum . zipWith (*) x) m |
5 | 14 |
|
6 | 15 | fft :: [Complex Double] -> [Complex Double]
|
7 |
| -fft x = let n = length x |
8 |
| - i = 0 :+ 1 |
9 |
| - w = M.fromList [(k%n, exp ((-2)*pi*i*(fromIntegral k)/(fromIntegral n)) ) | k<-[0..n-1]] |
10 |
| - arr = fft' n w (listArray (0,n-1) x) |
11 |
| - in [arr!k | k<-[0..n-1]] |
| 16 | +fft x = fft' x |
12 | 17 | where
|
13 |
| - fft' 1 _ x = x |
14 |
| - fft' n w x = let n2 = div n 2 |
15 |
| - e = fft' n2 w (listArray (0, n2-1) [x!k | k<-[0,2..n-1]]) |
16 |
| - o = fft' n2 w (listArray (0, n2-1) [x!k | k<-[1,3..n-1]]) |
17 |
| - in array (0, n-1) $ concat [[(k, e!k + o!k * w M.!(k%n)), |
18 |
| - (k + n2, e!k - o!k * w M.!(k%n))] |
19 |
| - | k <- [0..n2-1]] |
| 18 | + n = length x |
| 19 | + w0 = exp ((-2) * pi * (0 :+ 1) / fromIntegral n) |
| 20 | + w = M.fromList [(k % n, w0 ^ k) | k <- [0 .. n - 1]] |
| 21 | + fft' [x] = [x] |
| 22 | + fft' x = |
| 23 | + let (evens, odds) = partition (even . fst) $ zip [0 ..] x |
| 24 | + e = fft' $ map snd evens |
| 25 | + o = fft' $ map snd odds |
| 26 | + x1 = zipWith3 (\e o k -> e + o * w ! (k %n)) e o [0 ..] |
| 27 | + x2 = zipWith3 (\e o k -> e - o * w ! (k %n)) e o [0 ..] |
| 28 | + in x1 ++ x2 |
20 | 29 |
|
21 | 30 | main = do
|
22 |
| - print $ fft [0,1,2,3] |
| 31 | + print $ dft [0, 1, 2, 3] |
| 32 | + print $ fft [0, 1, 2, 3] |
0 commit comments