module Examples.Math.Fft where
import qualified Prelude as P
import Feldspar
import Feldspar.Vector
import Feldspar.Matrix
import Feldspar.Compiler
-- | Wrapper to define the size of vectors in 'fft'
fftInstance :: Data [Complex Float] -> Data [Complex Float]
fftInstance = freezeVector . fft . unfreezeVector' 256
-- | Wrapper to define the size of vectors in 'fft'
fft_wrapped ::Data' D256 [Complex Float] -> Data [Complex Float]
fft_wrapped = wrap fft
-- | Wrapper to define the size of vectors in 'ifft'
ifftInstance :: Data [Complex Float] -> Data [Complex Float]
ifftInstance = freezeVector . ifft . unfreezeVector' 256
-- | Wrapper to define the size of vectors in 'ifft'
ifft_wrapped ::Data' D256 [Complex Float] -> Data [Complex Float]
ifft_wrapped = wrap ifft
-- =================== INTERFACE ==================================
-- | Radix-2 Decimation-In-Frequeny Fast Fourier Transformation of the given complex vector
-- The given vector must be power-of-two sized, (for example 2, 4, 8, 16, 32, etc.)
fft :: DVector (Complex Float) -> DVector (Complex Float)
fft v = bitRev (loglen-1) $ fftCore (loglen-1) v
where loglen = f2i $ logBase 2 $ i2f $ length v
-- | Radix-2 Decimation-In-Frequeny Inverse Fast Fourier Transformation of the given complex vector
-- The given vector must be power-of-two sized, (for example 2, 4, 8, 16, 32, etc.)
ifft :: DVector (Complex Float) -> DVector (Complex Float)
ifft v = bitRev (loglen-1) $ ifftCore (loglen-1) v
where loglen = f2i $ logBase 2 $ i2f $ length v
-- ================================================================
-- | fftCore function uses 2^(n+1) input vector.
-- Output from the last stage needs to be bit reversed using bitRev function (if required)
fftCore :: Data Index -> DVector (Complex Float) -> DVector (Complex Float)
fftCore n v = composeOn stage (reverse (0...n)) v
stage k (Indexed l ixf Empty) = (Indexed l ixf' Empty)
where
k2 = 1 << k
ixf' i = condition (testBit i k) (twid * (b-a)) (a+b)
where
a = ixf i
b = ixf (i `xor` k2)
twid = cis (-pi*(i2f (lsbs k i)) / i2f k2)
-- | ifftCore function uses 2^(n+1) input vector.
-- Output from the last stage needs to be bit reversed using bitRev function (if required)
ifftCore :: Data Index -> DVector (Complex Float) -> DVector (Complex Float)
ifftCore n v = map (/ (complex (i2f (2^(n+1))) 0)) $ composeOn istage (reverse (0...n)) v
istage k (Indexed l ixf Empty) = (Indexed l ixf' Empty)
where
k2 = 1 << k
ixf' i = condition (testBit i k) (twid * (b-a)) (a+b)
where
a = ixf i
b = ixf (i `xor` k2)
twid = cis (pi*(i2f (lsbs k i)) / i2f k2)
-- | bitRev function transforms the given vector to bitreversal order
-- parameter n is the size of the input vector
bitRev :: Type a => Data Index -> Vector (Data a) -> Vector (Data a)
bitRev n = pipe riffle (1...n)
-- | Helper functions for fftCore and ifftCore and bitRev
pipe :: (Syntactic a) => (Data Index -> a -> a) -> Vector (Data Index) -> a -> a
pipe = flip.fold.flip
composeOn f is as = fold (flip f) as is
allOnes = complement 0
oneBits n = complement (allOnes << n)
lsbs k i = i .&. oneBits k
par m n f = mat2Vec m n . map f . vec2Mat m n
-- k at least 1
rotBit :: Data Index -> Data Index -> Data Index
rotBit 0 _ = error "k should be at least 1"
rotBit k i = lefts .|. rights
where
ir = i >> 1
rights = ir .&. (oneBits k)
lefts = (((ir >> k) << 1) .|. (i .&. 1)) << k
riffle k (Indexed l ixf Empty) = indexed l (ixf.rotBit k)
vec2Mat :: Data Index -> Data Index -> Vector (Data a) -> Matrix a
vec2Mat m n (Indexed l ixf Empty) = indexedMat (1 << m) (1 << n) ixf'
where
ixf' i j = ixf $ (i << n) `xor` j
mat2Vec :: Type a => Data Index -> Data Index -> Matrix a -> Vector (Data a)
mat2Vec m n matr = Indexed (1 << m << n) ixf Empty
where
ixf i = matr ! y ! x
where
y = i >> n
x = i .&. (oneBits n)
-- Ad-hoc function to generate a power-of-two sequence in order to test fft and ifft
pow2Seq :: Data DefaultWord -> DVector (Complex Float)
pow2Seq n = indexed (2 ^ n) (\i ->complex (i2f i) 0 )