1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
{-# 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.
module Statistics.Transform
(
-- * Type synonyms
CD
-- * Discrete cosine transform
, dct
, dct_
, idct
, idct_
-- * Fast Fourier transform
, fft
, ifft
) where
import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.Bits (shiftL, shiftR)
import Data.Complex (Complex(..), conjugate, realPart)
import Numeric.SpecFunctions (log2)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed as U
type CD = Complex Double
-- | Discrete cosine transform (DCT-II).
dct :: U.Vector Double -> U.Vector Double
dct = dctWorker . G.map (:+0)
-- | Discrete cosine transform (DCT-II). Only real part of vector is
-- transformed, imaginary part is ignored.
dct_ :: U.Vector CD -> U.Vector Double
dct_ = dctWorker . G.map (\(i :+ _) -> i :+ 0)
dctWorker :: U.Vector CD -> U.Vector Double
dctWorker xs
= G.map realPart $ G.zipWith (*) weights (fft interleaved)
where
interleaved = G.backpermute xs $ G.enumFromThenTo 0 2 (len-2) G.++
G.enumFromThenTo (len-1) (len-3) 1
weights = G.cons 2 . G.generate (len-1) $ \x ->
2 * exp ((0:+(-1))*fi (x+1)*pi/(2*n))
where n = fi len
len = G.length xs
-- | Inverse discrete cosine transform (DCT-III). It's inverse of
-- 'dct' only up to scale parameter:
--
-- > (idct . dct) x = (* lenngth x)
idct :: U.Vector Double -> U.Vector Double
idct = idctWorker . G.map (:+0)
-- | Inverse discrete cosine transform (DCT-III). Only real part of vector is
-- transformed, imaginary part is ignored.
idct_ :: U.Vector CD -> U.Vector Double
idct_ = idctWorker . G.map (\(i :+ _) -> i :+ 0)
idctWorker :: U.Vector CD -> U.Vector Double
idctWorker xs = G.generate len interleave
where
interleave z | even z = vals `G.unsafeIndex` halve z
| otherwise = vals `G.unsafeIndex` (len - halve z - 1)
vals = G.map realPart . ifft $ G.zipWith (*) weights xs
weights
= G.cons n
$ G.generate (len - 1) $ \x -> 2 * n * exp ((0:+1) * fi (x+1) * pi/(2*n))
where n = fi len
len = G.length xs
-- | Inverse fast Fourier transform.
ifft :: U.Vector CD -> U.Vector CD
ifft xs = G.map ((/fi (G.length xs)) . conjugate) . fft . G.map conjugate $ xs
-- | Radix-2 decimation-in-time fast Fourier transform.
fft :: U.Vector CD -> U.Vector CD
fft v = G.create $ do
mv <- G.thaw v
mfft mv
return mv
mfft :: (M.MVector v CD) => v s CD -> ST s ()
mfft vec
| 1 `shiftL` m /= len = error "Statistics.Transform.fft: bad vector size"
| otherwise = bitReverse 0 0
where
bitReverse i j | i == len-1 = stage 0 1
| otherwise = do
when (i < j) $ M.swap vec i j
let inner k l | k <= l = inner (k `shiftR` 1) (l-k)
| otherwise = bitReverse (i+1) (l+k)
inner (len `shiftR` 1) j
stage l !l1 | l == m = return ()
| otherwise = do
let !l2 = l1 `shiftL` 1
!e = -6.283185307179586/fromIntegral l2
flight j !a | j == l1 = stage (l+1) l2
| otherwise = do
let butterfly i | i >= len = flight (j+1) (a+e)
| otherwise = do
let i1 = i + l1
xi1 :+ yi1 <- M.read vec i1
let !c = cos a
!s = sin a
d = (c*xi1 - s*yi1) :+ (s*xi1 + c*yi1)
ci <- M.read vec i
M.write vec i1 (ci - d)
M.write vec i (ci + d)
butterfly (i+l2)
butterfly j
flight 0 0
len = M.length vec
m = log2 len
fi :: Int -> CD
fi = fromIntegral
halve :: Int -> Int
halve = (`shiftR` 1)
|
