Skip to content

Instantly share code, notes, and snippets.

@ian-ross
Created November 18, 2013 19:22
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ian-ross/7533714 to your computer and use it in GitHub Desktop.
Save ian-ross/7533714 to your computer and use it in GitHub Desktop.
import Prelude hiding (length, sum, map, zipWith, (++), foldr, concat)
import qualified Prelude as P
import Data.Complex
import Data.Vector
import Data.Bits
i :: Complex Double
i = 0 :+ 1
omega :: Int -> Complex Double
omega n = cis (2 * pi / fromIntegral n)
dft, idft :: Vector (Complex Double) -> Vector (Complex Double)
dft = dft' 1 1
idft v = dft' (-1) (1.0 / (fromIntegral $ length v)) v
dft' :: Int -> Double -> Vector (Complex Double) -> Vector (Complex Double)
dft' sign scale h = generate bigN (((scale :+ 0) *) . doone)
where bigN = length h
w = omega bigN
doone n = sum $
zipWith (*) h $ generate bigN (\k -> w^^(sign*n*k))
fft, ifft :: Vector (Complex Double) -> Vector (Complex Double)
fft = fft' 1 1
ifft v = fft' (-1) (1.0 / (fromIntegral $ length v)) v
fft' :: Int -> Double -> Vector (Complex Double) -> Vector (Complex Double)
fft' sign scale h =
if n <= 2
then dft' sign scale h
else map ((scale :+ 0) *) $ recomb $ backpermute h (bitrev n)
where n = length h
recomb = foldr (.) id $ map dl $ iterateN (log2 n) (`div` 2) n
dl m v = let w = omega m
m2 = m `div` 2
ds = map ((w ^^) . (sign *)) $ enumFromN 0 m2
doone v = let v0 = slice 0 m2 v
v1 = zipWith (*) ds $ slice m2 m2 v
in (zipWith (+) v0 v1) ++ (zipWith (-) v0 v1)
in concat $ P.map doone $ slicevecs m v
slicevecs m v = P.map (\i -> slice (i * m) m v) [0..n `div` m - 1]
bitrev :: Int -> Vector Int
bitrev n =
let nbits = log2 n
bs = generate nbits id
onebit i r b = if testBit i b then setBit r (nbits - 1 - b) else r
in generate n (\i -> foldl' (onebit i) 0 bs)
log2 :: Int -> Int
log2 1 = 0
log2 n = 1 + log2 (n `div` 2)
tst bigN = foldr (\l r -> l P.++ " . " P.++ r) "id" $
map ((("dl " P.++ show bigN P.++ " ") P.++) . show) $
iterateN (log2 bigN) (`div` 2) bigN
defuzz :: Vector (Complex Double) -> Vector (Complex Double)
defuzz = map (\(r :+ i) -> df r :+ df i)
where df x = if abs x < 1.0E-6 then 0 else x
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment