Created
November 27, 2013 19:57
-
-
Save ian-ross/7682237 to your computer and use it in GitHub Desktop.
Initial mixed-radix FFT implementation (http://www.skybluetrades.net/blog/posts/2013/11/27/data-analysis-fft-6.html)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{-# LANGUAGE ScopedTypeVariables, TypeSynonymInstances, FlexibleInstances #-} | |
module FFT where | |
import Prelude hiding (length, sum, map, zipWith, (++), foldr, foldr1, or, | |
concat, concatMap, replicate, scanl, scanl1, scanr, null, | |
init, last, tail, head, filter, reverse, product, | |
maximum, zip, dropWhile) | |
import qualified Prelude as P | |
import Data.Complex | |
import Data.Vector | |
-- Packages used only for testing. | |
import Data.Bits | |
import Control.Applicative ((<$>)) | |
import Test.QuickCheck | |
-- Typing Vector (Complex Double) or Vector (Vector (Complex Double)) | |
-- gets old quickly. | |
type CD = Complex Double | |
type VCD = Vector CD | |
type VVCD = Vector (Vector CD) | |
-- Roots of unity. | |
omega :: Int -> CD | |
omega n = cis (2 * pi / fromIntegral n) | |
-- DFT and inverse DFT drivers. | |
dft, idft :: VCD -> VCD | |
dft = dft' 1 1 | |
idft v = dft' (-1) (1.0 / (fromIntegral $ length v)) v | |
-- Naïve DFT. | |
dft' :: Int -> Double -> VCD -> VCD | |
dft' sign scale h = generate n (((scale :+ 0) *) . doone) | |
where n = length h | |
w = omega n | |
doone i = sum $ zipWith (*) h $ generate n (\k -> w^^(sign*i*k)) | |
-- FFT and inverse FFT drivers. | |
fft, ifft :: VCD -> VCD | |
fft = fft' 1 1 | |
ifft v = fft' (-1) (1.0 / (fromIntegral $ length v)) v | |
-- Mixed-radix decimation-in-time Cooley-Tukey FFT. | |
fft' :: Int -> Double -> VCD -> VCD | |
fft' sign scale h = if n == 1 then h | |
else if null fs | |
then dft' sign scale h | |
else map ((scale :+ 0) *) fullfft | |
where | |
-- Full vector length. | |
n = length h | |
-- Factorise input vector length. | |
splitfs@(lastf, fs) = factors n | |
-- Generate sizing information for Danielson-Lanczos step. | |
wfacs = map (n `div`) $ scanl (*) 1 fs | |
dlinfo = zip wfacs fs | |
-- Compose all Danielson-Lanczos steps and "final" DFT. | |
recomb = foldr (.) (mdft lastf) $ map (dl sign) dlinfo | |
-- Apply Danielson-Lanczos steps and "final" DFT to digit reversal | |
-- ordered input vector. | |
fullfft = recomb $ backpermute h (digrev n splitfs) | |
-- Multiple DFT for "bottom" of algorithm. | |
mdft :: Int -> VCD -> VCD | |
mdft factor h = concatMap (dft' sign 1) $ slicevecs factor h | |
-- Single Danielson-Lanczos step: process all duplicates and | |
-- concatenate into a single vector. | |
dl :: Int -> (Int, Int) -> VCD -> VCD | |
dl sign (wfac, split) v = concatMap doone $ slicevecs wfac v | |
where | |
-- Overall vector length. | |
n = length v | |
-- Size of each diagonal sub-matrix. | |
ns = wfac `div` split | |
-- Root of unity for the size of diagonal matrix we need. | |
w = omega $ sign * wfac | |
-- Basic diagonal entries for a given column. | |
ds c = map ((w^^) . (c *)) $ enumFromN 0 ns | |
-- Twiddled diagonal entries in row r, column c (both | |
-- zero-indexed), where each row and column if a wfac x wfac | |
-- matrix. | |
d r c = map (w^(ns*r*c) *) (ds c) | |
-- Process one duplicate by processing all rows and concatenating | |
-- the results into a single vector. | |
doone v = concatMap single $ enumFromN 0 split | |
where vs :: VVCD | |
vs = slicevecs ns v | |
-- Multiply a single block by its appropriate diagonal | |
-- elements. | |
mult :: Int -> Int -> VCD | |
mult r c = zipWith (*) (d r c) (vs!c) | |
-- Multiply all blocks by the corresponding diagonal | |
-- elements in a single row. | |
single :: Int -> VCD | |
single r = foldl1' (zipWith (+)) $ map (mult r) $ enumFromN 0 split | |
-- Generate digit reversal permutation using elementary "modulo" | |
-- permutations: last digit is not permuted to match with using a | |
-- simple DFT at the "bottom" of the overall algorithm. | |
digrev :: Int -> (Int, Vector Int) -> Vector Int | |
digrev n (lastf, fs) = foldl1' (%.%) $ map dupperm subperms | |
where | |
-- Sizes of the individual permutations that we need, one per | |
-- factor. | |
sizes = scanl div n fs | |
-- Partial sub-permutations, one per factor. | |
subperms = reverse $ zipWith perm sizes fs | |
-- Duplicate a sub-permutation to fill the required vector length. | |
dupperm p = | |
let sublen = length p | |
shift di = map (+(sublen * di)) p | |
in concatMap shift $ enumFromN 0 (n `div` sublen) | |
-- Generate a single "modulo" permutation. | |
perm n fac = concatMap doone $ enumFromN 0 fac | |
where n1 = n `div` fac | |
doone i = generate n1 (\j -> j * fac + i) | |
-- Composition of permutations. | |
(%.%) :: Vector Int -> Vector Int -> Vector Int | |
p1 %.% p2 = backpermute p2 p1 | |
-- Slice a vector v into equally sized parts, each of length m. | |
slicevecs :: Int -> VCD -> VVCD | |
slicevecs m v = map (\i -> slice (i * m) m v) $ enumFromN 0 (length v `div` m) | |
-- From Haskell wiki. | |
primes :: [Int] | |
primes = 2 : primes' | |
where primes' = sieve [3, 5 ..] 9 primes' | |
sieve (x:xs) q ps@ ~(p:t) | |
| x < q = x : sieve xs q ps | |
| True = sieve [x | x <- xs, rem x p /= 0] (P.head t^2) t | |
-- Simple prime factorisation: small factors only; largest/last factor | |
-- picked out as "special". | |
factors :: Int -> (Int, Vector Int) | |
factors n = let (lst, rest) = go n primes in (lst, fromList rest) | |
where go cur pss@(p:ps) | |
| cur == p = (p, []) | |
| cur `mod` p == 0 = let (lst, rest) = go (cur `div` p) pss | |
in (lst, p : rest) | |
| otherwise = go cur ps | |
-- Testing and debugging code. | |
-- Clean up number display. | |
defuzz :: VCD -> VCD | |
defuzz = map (\(r :+ i) -> df r :+ df i) | |
where df x = if abs x < 1.0E-6 then 0 else x | |
-- Check FFT against DFT. | |
check :: VCD -> (Double, VCD) | |
check v = let diff = defuzz $ zipWith (-) (fft v) (dft v) | |
in (maximum $ map magnitude diff, diff) | |
-- Check FFT-inverse FFT round-trip. | |
icheck :: VCD -> (Double, VCD) | |
icheck v = let diff = defuzz $ zipWith (-) v (ifft $ fft v) | |
in (maximum $ map magnitude diff, diff) | |
-- QuickCheck property for FFT vs. DFT testing. | |
prop_dft_vs_fft (v :: VCD) = fst (check v) < 1.0E-6 | |
-- QuickCheck property for inverse FFT round-trip testing. | |
prop_ifft (v :: VCD) = maximum (map magnitude diff) < 1.0E-6 | |
where diff = zipWith (-) v (ifft $ fft v) | |
-- Non-zero length arbitrary vectors. | |
instance Arbitrary VCD where | |
arbitrary = fromList <$> listOf1 arbitrary | |
-- Check all 0/1 vectors of a given length. | |
repcheck :: (VCD -> (Double, VCD)) -> Int -> Maybe Int | |
repcheck checker n = if null chks | |
then Nothing | |
else Just $ snd $ head chks | |
where vs = map tobits $ enumFromN (0::Int) (2^n) | |
tobits i = generate n (\j -> if testBit i j then 1:+0 else 0:+0) | |
chks = dropWhile (\(d,_) -> d < 1.0E-6) $ | |
zip (map (fst . checker) vs) (enumFromN 0 (2^n)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment