Source

signals / src / Math / Signals / Transforms / Discrete / DFT.hs

{-# LANGUAGE BangPatterns #-}

module Math.Signals.Transforms.Discrete.DFT where

import qualified Data.Vector.Unboxed as V 
import qualified Data.Vector.Unboxed.Mutable as MV 
import  Data.Complex
import Control.Monad.Identity
--import 


negUnityRoot:: (Integral a1, Integral a2, RealFloat a) => a1 -> a2 -> Complex a
negUnityRoot k n  = exp $!  (:+) 0 $!  -2 * pi * fromIntegral k / fromIntegral n
{-# INLINE negUnityRoot #-}

unityRoot:: (Integral a1, Integral a2, RealFloat a) => a1 -> a2 -> Complex a
unityRoot k n  = exp $!  (:+) 0 $!  2 * pi * fromIntegral k / fromIntegral n
{-# INLINE unityRoot #-}


{- for when you want to multiply a real number by a complex! -}
scale :: Num a => a -> Complex a -> Complex a
scale !a !(rl:+im) = (a* rl) :+ (a*im)
{-#  INLINE scale #-}

realDFT  :: (RealFloat b, MV.Unbox b) => V.Vector b -> V.Vector (Complex b)
realDFT  vec=  case   V.length vec of 
            len ->   V.generate len  (\ixk -> V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                            scale  vEntry  $!  negUnityRoot (ixn * k)  len
{-# INLINE realDFT #-}                        

normalizedRealDFT   :: (RealFloat b, MV.Unbox b) => V.Vector b -> V.Vector (Complex b)
normalizedRealDFT   vec =  case   V.length vec of 
            len ->   V.generate len  (\ixk ->  invSqrtN * V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                            scale  vEntry  $!  negUnityRoot (ixn * k)  len                     
                    invSqrtN = 1 / (sqrt $! fromIntegral len)
{-# INLINE normalizedRealDFT #-}


invDft :: (RealFloat a, MV.Unbox a) => V.Vector (Complex a) -> V.Vector (Complex a)
invDft vec=  case   V.length vec of 
            len ->   V.generate len  (\ixk -> V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                              vEntry  *  unityRoot (ixn * k)  len
{-# INLINE invDft #-}                              

dft :: (RealFloat a, MV.Unbox a) => V.Vector (Complex a) -> V.Vector (Complex a)
dft vec=  case   V.length vec of 
            len ->   V.generate len  (\ixk -> V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                              vEntry  *  negUnityRoot (ixn * k)  len
{-# INLINE dft #-}                              



invNormalizedDFT :: (RealFloat a, MV.Unbox a) => V.Vector (Complex a) -> V.Vector (Complex a)
invNormalizedDFT vec=  case   V.length vec of 
            len ->   V.generate len  (\ixk -> invSqrtN *  V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                              vEntry  *  unityRoot (ixn * k)  len
                    invSqrtN = 1 / (sqrt $! fromIntegral len)
{-# INLINE invNormalizedDFT#-}                    

normalizedDFT :: (RealFloat a, MV.Unbox a) => V.Vector (Complex a) -> V.Vector (Complex a)
normalizedDFT vec=  case   V.length vec of 
            len ->   V.generate len  (\ixk -> invSqrtN *  V.ifoldl' (convolve ixk) 0 vec)  
                where 
                    convolve !k !zSum !ixn !vEntry  = (+) zSum $!  
                              vEntry  *  negUnityRoot (ixn * k)  len
                    invSqrtN = 1 / (sqrt $! fromIntegral len)
{-# INLINE normalizedDFT #-}

idDoubleRealDFT
  :: (RealFloat a, MV.Unbox a) => V.Vector a -> V.Vector (Complex a)
idDoubleRealDFT a = invNormalizedDFT $!  normalizedRealDFT a 


idDoubleDFT
  :: (RealFloat a, MV.Unbox a) =>
     V.Vector (Complex a) -> V.Vector (Complex a)
idDoubleDFT  a  = invNormalizedDFT . normalizedDFT $! a