Simple audio filtering application of powers-of-two FFT (http://www.skybluetrades.net/blog/posts/2013/11/21/data-analysis-fft-4/index.html)

audio-filtering.hs

module Main where

import Data.WAVE
import Prelude hiding (length, sum, map, zipWith, (++), foldr, concat, zip3,
                       replicate, concatMap)
import qualified Prelude as P
import Data.Complex
import Data.Vector
import Data.Bits
import System.Environment (getArgs)
import System.IO


-- A filter specifies a window size, a window overlap (both in terms
-- of numbers of samples) and a filter function (which should be a
-- window size length vector).
data Filter = Filter { fWindowSize :: Int
                     , fWindowOverlap :: Int
                     , fFilterFunc :: Vector (Complex Double) }

-- Helper for building filter functions.
build :: [(Int, Int)] -> Vector (Complex Double)
build segs = map realToFrac $ concat $ P.map (\(n, v) -> replicate n v) segs

-- Smooth filter function.
smooth :: Vector (Complex Double)
smooth = map realToFrac $ generate 1024 $ \i ->
  cos (2 * pi * fromIntegral i / 1023)

-- Some filters: at a sampling rate of 44100 Hz, 1024 samples is about
-- 0.02 seconds.
filt_noop, filt_hard, filt_hard_olap :: Filter
filt_smooth, filt_smooth_olap :: Filter
filt_noop = Filter 1024 0 (build [(1024,1)])
filt_hard = Filter 1024 0 (build [(128,1), (768,0), (128,1)])
filt_hard_olap = Filter 1024 256 (build [(128,1), (768,0), (128,1)])
filt_smooth = Filter 1024 0  smooth
filt_smooth_olap = Filter 1024 256 smooth
filters :: [Filter]
filters = [filt_noop, filt_hard, filt_hard_olap, filt_smooth, filt_smooth_olap]


-- Main processing driver.
doit :: Filter -> String -> String -> IO ()
doit filt@(Filter win o _) fin fout = do
  -- Read WAV file.
  w <- getWAVEFile fin

  -- Powers of two only.
  let h = waveHeader w
      Just nsamp = waveFrames h
      nwins = (nsamp + o) `div` (win - o)
      nsampout = (win - o) * nwins

  -- Convert all channels to Vector (Complex Double) and extract
  -- individual channels as vectors.
  let d = fromList $ P.map (map sampleToDouble . fromList) $
          P.take (nsampout + o) $ waveSamples w
      chs = map (\i -> map (!i) d) $ enumFromN 0 (waveNumChannels h)

  -- Run filter.
  let chout = map (runFilter filt) chs

  -- Convert data back to frames and write out.
  let dout = map (\i -> map (!i) chout) $ enumFromN 0 nsampout
      sampout = toList $ map (toList . map doubleToSample) dout
      wave = WAVE ((waveHeader w) { waveFrames = Just nsampout }) sampout
  putWAVEFile fout wave


-- Handy pipeline combinator.
(=>=) :: (a -> b) -> (b -> c) -> a -> c
(=>=) = flip (.)


-- Apply a filter to a channel.
runFilter :: Filter -> Vector Double -> Vector Double
runFilter (Filter w o ffn) din = proc ss
  where
    -- Data to Complex Double for our complex-to-complex FFT.
    cdin = map realToFrac din

    -- Window start positions.
    ss = enumFromStepN 0 (w - o) ((length din + o) `div` (w - o))

    -- Processing pipeline.
    proc =     map (\s -> slice s w cdin)      -- Extract windows.
           =>= map fft
           =>= map (zipWith (*) ffn)           -- Filter.
           =>= map ifft
           =>= map (slice (o `div` 2) (w - o)) -- Slice off overlaps.
           =>= toList =>= concat               -- Concatenate windows.
           =>= map realPart                    -- Convert back to real data.


main :: IO ()
main = do
  as <- getArgs
  if P.length as /= 3
    then putStrLn "Usage: filter in.wav out.wav filter-index"
    else do
    let [inwav, outwav, filts] = as
        filtidx = read filts :: Int
    if filtidx < 0 || filtidx >= P.length filters
      then putStrLn "Invalid filter index!"
      else doit (filters !! filtidx) inwav outwav


-- FFT code follows...

omega :: Int -> Complex Double
omega n = cis (2 * pi / fromIntegral n)

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 == 1
  then 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 $ sign * m
                     m2 = m `div` 2
                     ds = map (w ^) $ 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)