{-# LANGUAGE BangPatterns, FlexibleContexts #-}-- |-- Module : Statistics.Transform-- Copyright : (c) 2011 Bryan O'Sullivan-- License : BSD3---- Maintainer : bos@serpentine.com-- Stability : experimental-- Portability : portable---- Fourier-related transformations of mathematical functions.---- These functions are written for simplicity and correctness, not-- speed. If you need a fast FFT implementation for your application,-- you should strongly consider using a library of FFTW bindings-- instead.moduleStatistics.Transform(-- * Type synonymsCD -- * Discrete cosine transform,dct ,dct_ ,idct ,idct_ -- * Fast Fourier transform,fft ,ifft )whereimportControl.Monad(when)importControl.Monad.ST(ST)importData.Bits(shiftL,shiftR)importData.Complex(Complex(..),conjugate,realPart)importNumeric.SpecFunctions(log2)importqualifiedData.Vector.GenericasGimportqualifiedData.Vector.Generic.MutableasMimportqualifiedData.Vector.UnboxedasUimportqualifiedData.VectorasVtypeCD =ComplexDouble-- | Discrete cosine transform (DCT-II).dct ::(G.Vectorv CD ,G.Vectorv Double,G.Vectorv Int)=>v Double->v Doubledct :: forall (v :: * -> *). (Vector v CD, Vector v Double, Vector v Int) => v Double -> v Double dct =v CD -> v Double forall (v :: * -> *). (Vector v CD, Vector v Double, Vector v Int) => v CD -> v Double dctWorker (v CD -> v Double) -> (v Double -> v CD) -> v Double -> v Double forall b c a. (b -> c) -> (a -> b) -> a -> c .(Double -> CD) -> v Double -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map(Double -> Double -> CD forall a. a -> a -> Complex a :+Double 0){-# INLINABLEdct #-}{-# SPECIAlIZEdct ::U.VectorDouble->U.VectorDouble#-}{-# SPECIAlIZEdct ::V.VectorDouble->V.VectorDouble#-}-- | Discrete cosine transform (DCT-II). Only real part of vector is-- transformed, imaginary part is ignored.dct_ ::(G.Vectorv CD ,G.Vectorv Double,G.Vectorv Int)=>v CD ->v Doubledct_ :: forall (v :: * -> *). (Vector v CD, Vector v Double, Vector v Int) => v CD -> v Double dct_ =v CD -> v Double forall (v :: * -> *). (Vector v CD, Vector v Double, Vector v Int) => v CD -> v Double dctWorker (v CD -> v Double) -> (v CD -> v CD) -> v CD -> v Double forall b c a. (b -> c) -> (a -> b) -> a -> c .(CD -> CD) -> v CD -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map(\(Double i :+Double _)->Double i Double -> Double -> CD forall a. a -> a -> Complex a :+Double 0){-# INLINABLEdct_ #-}{-# SPECIAlIZEdct_ ::U.VectorCD ->U.VectorDouble#-}{-# SPECIAlIZEdct_ ::V.VectorCD ->V.VectorDouble#-}dctWorker ::(G.Vectorv CD ,G.Vectorv Double,G.Vectorv Int)=>v CD ->v Double{-# INLINEdctWorker #-}dctWorker :: forall (v :: * -> *). (Vector v CD, Vector v Double, Vector v Int) => v CD -> v Double dctWorker v CD xs -- length 1 is special cased because shuffle algorithms fail for it.|v CD -> Int forall (v :: * -> *) a. Vector v a => v a -> Int G.lengthv CD xs Int -> Int -> Bool forall a. Eq a => a -> a -> Bool ==Int 1=(CD -> Double) -> v CD -> v Double forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map((Double 2Double -> Double -> Double forall a. Num a => a -> a -> a *)(Double -> Double) -> (CD -> Double) -> CD -> Double forall b c a. (b -> c) -> (a -> b) -> a -> c .CD -> Double forall a. Complex a -> a realPart)v CD xs |v CD -> Bool forall (v :: * -> *) a. Vector v a => v a -> Bool vectorOK v CD xs =(CD -> Double) -> v CD -> v Double forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.mapCD -> Double forall a. Complex a -> a realPart(v CD -> v Double) -> v CD -> v Double forall a b. (a -> b) -> a -> b $(CD -> CD -> CD) -> v CD -> v CD -> v CD forall (v :: * -> *) a b c. (Vector v a, Vector v b, Vector v c) => (a -> b -> c) -> v a -> v b -> v c G.zipWithCD -> CD -> CD forall a. Num a => a -> a -> a (*)v CD weights (v CD -> v CD forall (v :: * -> *). Vector v CD => v CD -> v CD fft v CD interleaved )|Bool otherwise=[Char] -> v Double forall a. HasCallStack => [Char] -> a error[Char] "Statistics.Transform.dct: bad vector length"whereinterleaved :: v CD interleaved =v CD -> v Int -> v CD forall (v :: * -> *) a. (HasCallStack, Vector v a, Vector v Int) => v a -> v Int -> v a G.backpermutev CD xs (v Int -> v CD) -> v Int -> v CD forall a b. (a -> b) -> a -> b $Int -> Int -> Int -> v Int forall (v :: * -> *) a. (Vector v a, Enum a) => a -> a -> a -> v a G.enumFromThenToInt 0Int 2(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 2)v Int -> v Int -> v Int forall (v :: * -> *) a. Vector v a => v a -> v a -> v a G.++Int -> Int -> Int -> v Int forall (v :: * -> *) a. (Vector v a, Enum a) => a -> a -> a -> v a G.enumFromThenTo(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 1)(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 3)Int 1weights :: v CD weights =CD -> v CD -> v CD forall (v :: * -> *) a. Vector v a => a -> v a -> v a G.consCD 2(v CD -> v CD) -> ((Int -> CD) -> v CD) -> (Int -> CD) -> v CD forall b c a. (b -> c) -> (a -> b) -> a -> c .Int -> (Int -> CD) -> v CD forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a G.generate(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 1)((Int -> CD) -> v CD) -> (Int -> CD) -> v CD forall a b. (a -> b) -> a -> b $\Int x ->CD 2CD -> CD -> CD forall a. Num a => a -> a -> a *CD -> CD forall a. Floating a => a -> a exp((Double 0Double -> Double -> CD forall a. a -> a -> Complex a :+(-Double 1))CD -> CD -> CD forall a. Num a => a -> a -> a *Int -> CD fi (Int x Int -> Int -> Int forall a. Num a => a -> a -> a +Int 1)CD -> CD -> CD forall a. Num a => a -> a -> a *CD forall a. Floating a => a piCD -> CD -> CD forall a. Fractional a => a -> a -> a /(CD 2CD -> CD -> CD forall a. Num a => a -> a -> a *CD n ))wheren :: CD n =Int -> CD fi Int len len :: Int len =v CD -> Int forall (v :: * -> *) a. Vector v a => v a -> Int G.lengthv CD xs -- | Inverse discrete cosine transform (DCT-III). It's inverse of-- 'dct' only up to scale parameter:---- > (idct . dct) x = (* length x)idct ::(G.Vectorv CD ,G.Vectorv Double)=>v Double->v Doubleidct :: forall (v :: * -> *). (Vector v CD, Vector v Double) => v Double -> v Double idct =v CD -> v Double forall (v :: * -> *). (Vector v CD, Vector v Double) => v CD -> v Double idctWorker (v CD -> v Double) -> (v Double -> v CD) -> v Double -> v Double forall b c a. (b -> c) -> (a -> b) -> a -> c .(Double -> CD) -> v Double -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map(Double -> Double -> CD forall a. a -> a -> Complex a :+Double 0){-# INLINABLEidct #-}{-# SPECIAlIZEidct ::U.VectorDouble->U.VectorDouble#-}{-# SPECIAlIZEidct ::V.VectorDouble->V.VectorDouble#-}-- | Inverse discrete cosine transform (DCT-III). Only real part of vector is-- transformed, imaginary part is ignored.idct_ ::(G.Vectorv CD ,G.Vectorv Double)=>v CD ->v Doubleidct_ :: forall (v :: * -> *). (Vector v CD, Vector v Double) => v CD -> v Double idct_ =v CD -> v Double forall (v :: * -> *). (Vector v CD, Vector v Double) => v CD -> v Double idctWorker (v CD -> v Double) -> (v CD -> v CD) -> v CD -> v Double forall b c a. (b -> c) -> (a -> b) -> a -> c .(CD -> CD) -> v CD -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map(\(Double i :+Double _)->Double i Double -> Double -> CD forall a. a -> a -> Complex a :+Double 0){-# INLINABLEidct_ #-}{-# SPECIAlIZEidct_ ::U.VectorCD ->U.VectorDouble#-}{-# SPECIAlIZEidct_ ::V.VectorCD ->V.VectorDouble#-}idctWorker ::(G.Vectorv CD ,G.Vectorv Double)=>v CD ->v Double{-# INLINEidctWorker #-}idctWorker :: forall (v :: * -> *). (Vector v CD, Vector v Double) => v CD -> v Double idctWorker v CD xs |v CD -> Bool forall (v :: * -> *) a. Vector v a => v a -> Bool vectorOK v CD xs =Int -> (Int -> Double) -> v Double forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a G.generateInt len Int -> Double interleave |Bool otherwise=[Char] -> v Double forall a. HasCallStack => [Char] -> a error[Char] "Statistics.Transform.dct: bad vector length"whereinterleave :: Int -> Double interleave Int z |Int -> Bool forall a. Integral a => a -> Bool evenInt z =v Double vals v Double -> Int -> Double forall (v :: * -> *) a. Vector v a => v a -> Int -> a `G.unsafeIndex`Int -> Int halve Int z |Bool otherwise=v Double vals v Double -> Int -> Double forall (v :: * -> *) a. Vector v a => v a -> Int -> a `G.unsafeIndex`(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int -> Int halve Int z Int -> Int -> Int forall a. Num a => a -> a -> a -Int 1)vals :: v Double vals =(CD -> Double) -> v CD -> v Double forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.mapCD -> Double forall a. Complex a -> a realPart(v CD -> v Double) -> (v CD -> v CD) -> v CD -> v Double forall b c a. (b -> c) -> (a -> b) -> a -> c .v CD -> v CD forall (v :: * -> *). Vector v CD => v CD -> v CD ifft (v CD -> v Double) -> v CD -> v Double forall a b. (a -> b) -> a -> b $(CD -> CD -> CD) -> v CD -> v CD -> v CD forall (v :: * -> *) a b c. (Vector v a, Vector v b, Vector v c) => (a -> b -> c) -> v a -> v b -> v c G.zipWithCD -> CD -> CD forall a. Num a => a -> a -> a (*)v CD weights v CD xs weights :: v CD weights =CD -> v CD -> v CD forall (v :: * -> *) a. Vector v a => a -> v a -> v a G.consCD n (v CD -> v CD) -> v CD -> v CD forall a b. (a -> b) -> a -> b $Int -> (Int -> CD) -> v CD forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a G.generate(Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 1)((Int -> CD) -> v CD) -> (Int -> CD) -> v CD forall a b. (a -> b) -> a -> b $\Int x ->CD 2CD -> CD -> CD forall a. Num a => a -> a -> a *CD n CD -> CD -> CD forall a. Num a => a -> a -> a *CD -> CD forall a. Floating a => a -> a exp((Double 0Double -> Double -> CD forall a. a -> a -> Complex a :+Double 1)CD -> CD -> CD forall a. Num a => a -> a -> a *Int -> CD fi (Int x Int -> Int -> Int forall a. Num a => a -> a -> a +Int 1)CD -> CD -> CD forall a. Num a => a -> a -> a *CD forall a. Floating a => a piCD -> CD -> CD forall a. Fractional a => a -> a -> a /(CD 2CD -> CD -> CD forall a. Num a => a -> a -> a *CD n ))wheren :: CD n =Int -> CD fi Int len len :: Int len =v CD -> Int forall (v :: * -> *) a. Vector v a => v a -> Int G.lengthv CD xs -- | Inverse fast Fourier transform.ifft ::G.Vectorv CD =>v CD ->v CD ifft :: forall (v :: * -> *). Vector v CD => v CD -> v CD ifft v CD xs |v CD -> Bool forall (v :: * -> *) a. Vector v a => v a -> Bool vectorOK v CD xs =(CD -> CD) -> v CD -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.map((CD -> CD -> CD forall a. Fractional a => a -> a -> a /Int -> CD fi (v CD -> Int forall (v :: * -> *) a. Vector v a => v a -> Int G.lengthv CD xs ))(CD -> CD) -> (CD -> CD) -> CD -> CD forall b c a. (b -> c) -> (a -> b) -> a -> c .CD -> CD forall a. Num a => Complex a -> Complex a conjugate)(v CD -> v CD) -> (v CD -> v CD) -> v CD -> v CD forall b c a. (b -> c) -> (a -> b) -> a -> c .v CD -> v CD forall (v :: * -> *). Vector v CD => v CD -> v CD fft (v CD -> v CD) -> (v CD -> v CD) -> v CD -> v CD forall b c a. (b -> c) -> (a -> b) -> a -> c .(CD -> CD) -> v CD -> v CD forall (v :: * -> *) a b. (Vector v a, Vector v b) => (a -> b) -> v a -> v b G.mapCD -> CD forall a. Num a => Complex a -> Complex a conjugate(v CD -> v CD) -> v CD -> v CD forall a b. (a -> b) -> a -> b $v CD xs |Bool otherwise=[Char] -> v CD forall a. HasCallStack => [Char] -> a error[Char] "Statistics.Transform.ifft: bad vector length"{-# INLINABLEifft #-}{-# SPECIAlIZEifft ::U.VectorCD ->U.VectorCD #-}{-# SPECIAlIZEifft ::V.VectorCD ->V.VectorCD #-}-- | Radix-2 decimation-in-time fast Fourier transform.fft ::G.Vectorv CD =>v CD ->v CD fft :: forall (v :: * -> *). Vector v CD => v CD -> v CD fft v CD v |v CD -> Bool forall (v :: * -> *) a. Vector v a => v a -> Bool vectorOK v CD v =(forall s. ST s (Mutable v s CD)) -> v CD forall (v :: * -> *) a. Vector v a => (forall s. ST s (Mutable v s a)) -> v a G.create((forall s. ST s (Mutable v s CD)) -> v CD) -> (forall s. ST s (Mutable v s CD)) -> v CD forall a b. (a -> b) -> a -> b $doMutable v s CD mv <-v CD -> ST s (Mutable v (PrimState (ST s)) CD) forall (m :: * -> *) (v :: * -> *) a. (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a) G.thawv CD v Mutable v s CD -> ST s () forall (v :: * -> * -> *) s. MVector v CD => v s CD -> ST s () mfft Mutable v s CD mv Mutable v s CD -> ST s (Mutable v s CD) forall a. a -> ST s a forall (m :: * -> *) a. Monad m => a -> m a returnMutable v s CD mv |Bool otherwise=[Char] -> v CD forall a. HasCallStack => [Char] -> a error[Char] "Statistics.Transform.fft: bad vector length"{-# INLINABLEfft #-}{-# SPECIAlIZEfft ::U.VectorCD ->U.VectorCD #-}{-# SPECIAlIZEfft ::V.VectorCD ->V.VectorCD #-}-- Vector length must be power of two. It's not checkedmfft ::(M.MVectorv CD )=>v s CD ->STs (){-# INLINEmfft #-}mfft :: forall (v :: * -> * -> *) s. MVector v CD => v s CD -> ST s () mfft v s CD vec =Int -> Int -> ST s () forall {m :: * -> *}. (PrimState m ~ s, PrimMonad m) => Int -> Int -> m () bitReverse Int 0Int 0wherebitReverse :: Int -> Int -> m () bitReverse Int i Int j |Int i Int -> Int -> Bool forall a. Eq a => a -> a -> Bool ==Int len Int -> Int -> Int forall a. Num a => a -> a -> a -Int 1=Int -> Int -> m () forall {m :: * -> *}. (PrimState m ~ s, PrimMonad m) => Int -> Int -> m () stage Int 0Int 1|Bool otherwise=doBool -> m () -> m () forall (f :: * -> *). Applicative f => Bool -> f () -> f () when(Int i Int -> Int -> Bool forall a. Ord a => a -> a -> Bool <Int j )(m () -> m ()) -> m () -> m () forall a b. (a -> b) -> a -> b $v (PrimState m) CD -> Int -> Int -> m () forall (m :: * -> *) (v :: * -> * -> *) a. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m () M.swapv s CD v (PrimState m) CD vec Int i Int j letinner :: Int -> Int -> m () inner Int k Int l |Int k Int -> Int -> Bool forall a. Ord a => a -> a -> Bool <=Int l =Int -> Int -> m () inner (Int k Int -> Int -> Int forall a. Bits a => a -> Int -> a `shiftR`Int 1)(Int l Int -> Int -> Int forall a. Num a => a -> a -> a -Int k )|Bool otherwise=Int -> Int -> m () bitReverse (Int i Int -> Int -> Int forall a. Num a => a -> a -> a +Int 1)(Int l Int -> Int -> Int forall a. Num a => a -> a -> a +Int k )Int -> Int -> m () inner (Int len Int -> Int -> Int forall a. Bits a => a -> Int -> a `shiftR`Int 1)Int j stage :: Int -> Int -> m () stage Int l !Int l1 |Int l Int -> Int -> Bool forall a. Eq a => a -> a -> Bool ==Int m =() -> m () forall a. a -> m a forall (m :: * -> *) a. Monad m => a -> m a return()|Bool otherwise=dolet!l2 :: Int l2 =Int l1 Int -> Int -> Int forall a. Bits a => a -> Int -> a `shiftL`Int 1!e :: Double e =-Double 6.283185307179586Double -> Double -> Double forall a. Fractional a => a -> a -> a /Int -> Double forall a b. (Integral a, Num b) => a -> b fromIntegralInt l2 flight :: Int -> Double -> m () flight Int j !Double a |Int j Int -> Int -> Bool forall a. Eq a => a -> a -> Bool ==Int l1 =Int -> Int -> m () stage (Int l Int -> Int -> Int forall a. Num a => a -> a -> a +Int 1)Int l2 |Bool otherwise=doletbutterfly :: Int -> m () butterfly Int i |Int i Int -> Int -> Bool forall a. Ord a => a -> a -> Bool >=Int len =Int -> Double -> m () flight (Int j Int -> Int -> Int forall a. Num a => a -> a -> a +Int 1)(Double a Double -> Double -> Double forall a. Num a => a -> a -> a +Double e )|Bool otherwise=doleti1 :: Int i1 =Int i Int -> Int -> Int forall a. Num a => a -> a -> a +Int l1 Double xi1 :+Double yi1 <-v (PrimState m) CD -> Int -> m CD forall (m :: * -> *) (v :: * -> * -> *) a. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a M.readv s CD v (PrimState m) CD vec Int i1 let!c :: Double c =Double -> Double forall a. Floating a => a -> a cosDouble a !s :: Double s =Double -> Double forall a. Floating a => a -> a sinDouble a d :: CD d =(Double c Double -> Double -> Double forall a. Num a => a -> a -> a *Double xi1 Double -> Double -> Double forall a. Num a => a -> a -> a -Double s Double -> Double -> Double forall a. Num a => a -> a -> a *Double yi1 )Double -> Double -> CD forall a. a -> a -> Complex a :+(Double s Double -> Double -> Double forall a. Num a => a -> a -> a *Double xi1 Double -> Double -> Double forall a. Num a => a -> a -> a +Double c Double -> Double -> Double forall a. Num a => a -> a -> a *Double yi1 )CD ci <-v (PrimState m) CD -> Int -> m CD forall (m :: * -> *) (v :: * -> * -> *) a. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a M.readv s CD v (PrimState m) CD vec Int i v (PrimState m) CD -> Int -> CD -> m () forall (m :: * -> *) (v :: * -> * -> *) a. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () M.writev s CD v (PrimState m) CD vec Int i1 (CD ci CD -> CD -> CD forall a. Num a => a -> a -> a -CD d )v (PrimState m) CD -> Int -> CD -> m () forall (m :: * -> *) (v :: * -> * -> *) a. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () M.writev s CD v (PrimState m) CD vec Int i (CD ci CD -> CD -> CD forall a. Num a => a -> a -> a +CD d )Int -> m () butterfly (Int i Int -> Int -> Int forall a. Num a => a -> a -> a +Int l2 )Int -> m () butterfly Int j Int -> Double -> m () flight Int 0Double 0len :: Int len =v s CD -> Int forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int M.lengthv s CD vec m :: Int m =Int -> Int log2Int len ------------------------------------------------------------------ Helpers----------------------------------------------------------------fi ::Int->CD fi :: Int -> CD fi =Int -> CD forall a b. (Integral a, Num b) => a -> b fromIntegralhalve ::Int->Inthalve :: Int -> Int halve =(Int -> Int -> Int forall a. Bits a => a -> Int -> a `shiftR`Int 1)vectorOK ::G.Vectorv a =>v a ->Bool{-# INLINEvectorOK #-}vectorOK :: forall (v :: * -> *) a. Vector v a => v a -> Bool vectorOK v a v =(Int 1Int -> Int -> Int forall a. Bits a => a -> Int -> a `shiftL`Int -> Int log2Int n )Int -> Int -> Bool forall a. Eq a => a -> a -> Bool ==Int n wheren :: Int n =v a -> Int forall (v :: * -> *) a. Vector v a => v a -> Int G.lengthv a v