Initial mixed-radix FFT implementation (http://www.skybluetrades.net/blog/posts/2013/11/27/data-analysis-fft-6.html)

FFT-v2.hs

{-# 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))