HaskellでFFT,STVectorといっしょ。

module Numerical.FFT ( fft1d ) where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (ST, runST)
import Data.Vector (Vector, MVector, (!), generate, map, thaw, freeze)
import qualified Data.Vector as V (length)
import Data.Vector.Mutable (STVector, new, set, read, write)
import qualified Data.Vector.Mutable as MV (length)
import Data.Maybe
import Data.Complex
import Data.Bits

import Prelude hiding ((++),length,read,map)

{- counterつきloop -}
loop :: (PrimMonad m) => Int -> Int -> (Int -> m ()) -> m ()
loop iStart iMax f = loopIter iStart 
  where
    loopIter i = when (i < iMax) ( do f i
                                      loopIter (i+1) )
                 
{- バタフライ演算 -}
butterflyCalc :: (RealFloat a, PrimMonad m) => 
                 Int -> 
                 Int -> 
                 Vector (Complex a)  -> 
                 MVector (PrimState m) (Complex a) -> 
                 MVector (PrimState m) (Complex a) ->
                 m (Vector (Complex a))
                 
butterflyCalc n m twiddlesTable prevDatas nextDatas 
  | n == 1 = do result <- freeze prevDatas
                return result
  | otherwise = do let m' = m * 2
                       n' = (div n 2)
                   loop 0 n' (\j -> loop 0 m (\i -> (do a <- read prevDatas (j*m+i)
                                                        b <- read prevDatas ((j+n')*m+i)
                                                        let bw = b*twiddlesTable!(i*n')
                                                        write nextDatas (j*m'+i)   (a+bw)
                                                        write nextDatas (j*m'+i+m) (a-bw)
                                                        return ())))
                   butterflyCalc n' m' twiddlesTable nextDatas prevDatas
    
fft1dCore' :: (RealFloat a) => Vector (Complex a) -> Vector (Complex a) -> Vector (Complex a)
fft1dCore' twiddlesTable inDatas = 
  runST $ do
    prevDatas <- thaw inDatas
    let dataLen = V.length inDatas
    nextDatas <- new dataLen
    set nextDatas 0
    result <- butterflyCalc dataLen 1 twiddlesTable prevDatas nextDatas
    return result
    
{- 時間間引きのstockham FFT -}
fft1dCore :: (RealFloat a) => Vector (Complex a) -> Bool -> Vector (Complex a)
fft1dCore inDatas bInverse = 
  let numOfSample = V.length inDatas
      signOfRootOfUnity = if bInverse then 1 else -1 
      twiddlesTable = generate numOfSample (\i -> cis (signOfRootOfUnity*2*pi*(fromIntegral i)/(fromIntegral numOfSample))) --回転子のωNのテーブルを作ります
      outDatas = fft1dCore' twiddlesTable inDatas
  in if bInverse 
     then map (\z -> z / (fromIntegral numOfSample)) outDatas 
     else outDatas

fft1d :: (RealFloat a) => Vector (Complex a) -> Bool -> Maybe (Vector (Complex a))
fft1d inDatas bInverse = 
  let n = V.length inDatas
  in if (n .&. (n-1) == 0)
     then Just (fft1dCore inDatas bInverse)
     else Nothing -- Dataの長さが2のべき乗でありません

[プログラミング] HaskellでFFT,UnboxなSTVectorといっしょ。
正格評価とかも使ってみた。

{-# LANGUAGE BangPatterns #-}

module Numerical.FFT ( fft1d ) where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (ST, runST)
import Data.Vector.Unboxed (Unbox, Vector, MVector, (!), generate, map, force, thaw, freeze)
import qualified Data.Vector.Unboxed as VU (length)
import Data.Vector.Unboxed.Mutable (STVector, new, set, unsafeRead, unsafeWrite)
import qualified Data.Vector.Unboxed.Mutable as MUV (length)
import Data.Maybe
import Data.Complex
import Data.Bits

import Prelude hiding ((++),length,read,map)

{- counterつきloop -}
loop :: (PrimMonad m) => Int -> Int -> (Int -> m ()) -> m ()
loop iStart iMax f = loopIter iStart 
  where
    loopIter i = when (i < iMax) ( do f i
                                      loopIter (i+1) )
                 
{- バタフライ演算 -}
butterflyCalc :: (Unbox a, RealFloat a, PrimMonad m) => 
                 Int -> 
                 Int -> 
                 Vector (Complex a)  -> 
                 MVector (PrimState m) (Complex a) -> 
                 MVector (PrimState m) (Complex a) ->
                 m (Vector (Complex a))
                 
butterflyCalc n m !twiddlesTable !prevDatas !nextDatas 
  | n == 1 = do result <- freeze prevDatas
                return result
  | otherwise = do let m' = m * 2
                       n' = (div n 2)
                   loop 0 n' (\j -> loop 0 m (\i -> (do 
                                                        a <- unsafeRead prevDatas (j*m+i)
                                                        b <- unsafeRead prevDatas ((j+n')*m+i)
                                                        let !bw = b*twiddlesTable!(i*n') 
                                                        unsafeWrite nextDatas (j*m'+i)   (a+bw)
                                                        unsafeWrite nextDatas (j*m'+i+m) (a-bw)
                                                        return ())))
                   butterflyCalc n' m' twiddlesTable nextDatas prevDatas
    
fft1dCore' :: (Unbox a, RealFloat a) => Vector (Complex a) -> Vector (Complex a) -> Vector (Complex a)
fft1dCore' twiddlesTable inDatas = 
  runST $ do
    prevDatas <- thaw inDatas
    let dataLen = VU.length inDatas
    nextDatas <- new dataLen
    set nextDatas 0
    result <- butterflyCalc dataLen 1 twiddlesTable prevDatas nextDatas
    return result