Skip to content

Commit 95a2811

Browse files
jiegilletleios
authored andcommitted
Added DFT and re-wrote FFT (#397)
1 parent 2fa064d commit 95a2811

File tree

2 files changed

+27
-17
lines changed

2 files changed

+27
-17
lines changed
+25-15
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,32 @@
11
import Data.Complex
2-
import Data.Array
3-
import Data.Ratio
2+
import Data.List (partition)
3+
import Data.Map ((!))
44
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
514

615
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
1217
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
2029

2130
main = do
22-
print $ fft [0,1,2,3]
31+
print $ dft [0, 1, 2, 3]
32+
print $ fft [0, 1, 2, 3]

contents/cooley_tukey/cooley_tukey.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ For some reason, though, putting code to this transformation really helped me fi
7878
{% sample lang="cpp" %}
7979
[import:23-33, lang:"c_cpp"](code/c++/fft.cpp)
8080
{% sample lang="hs" %}
81-
[import:4-13, lang:"julia"](code/julia/fft.jl)
81+
[import:7-13, lang:"haskell"](code/haskell/fft.hs)
8282
{% sample lang="py" %}
8383
[import:5-11, lang:"python"](code/python/fft.py)
8484
{% sample lang="scratch" %}
@@ -125,7 +125,7 @@ In the end, the code looks like:
125125
{% sample lang="cpp" %}
126126
[import:36-66, lang:"c_cpp"](code/c++/fft.cpp)
127127
{% sample lang="hs" %}
128-
[import:6-19, lang:"haskell"](code/haskell/fft.hs)
128+
[import:15-28, lang:"haskell"](code/haskell/fft.hs)
129129
{% sample lang="py" %}
130130
[import:13-24, lang:"python"](code/python/fft.py)
131131
{% sample lang="scratch" %}

0 commit comments

Comments
 (0)