{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Math.DFT -- Copyright : [2012..2020] The Accelerate Team -- License : BSD3 -- -- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- -- Compute the Discrete Fourier Transform (DFT) along the lower order dimension -- of an array. -- -- This uses a naïve algorithm which takes O(n^2) time. However, you can -- transform an array with an arbitrary extent, unlike with FFT which requires -- each dimension to be a power of two. -- -- The `dft` and `idft` functions compute the roots of unity as needed. If you -- need to transform several arrays with the same extent than it is faster to -- compute the roots once using `rootsOfUnity` or `inverseRootsOfUnity` -- respectively, then call `dftG` directly. -- -- You can also compute single values of the transform using `dftGS` -- module Data.Array.Accelerate.Math.DFT ( dft, idft, dftG, dftGS, ) where import Prelude as P hiding ((!!)) import Data.Array.Accelerate as A import Data.Array.Accelerate.Math.DFT.Roots import Data.Array.Accelerate.Data.Complex -- | Compute the DFT along the low order dimension of an array -- dft :: (Shape sh, Slice sh, A.RealFloat e, A.FromIntegral Int e) => Acc (Array (sh:.Int) (Complex e)) -> Acc (Array (sh:.Int) (Complex e)) dft v = dftG (rootsOfUnity (shape v)) v -- | Compute the inverse DFT along the low order dimension of an array -- idft :: (Shape sh, Slice sh, A.RealFloat e, A.FromIntegral Int e) => Acc (Array (sh:.Int) (Complex e)) -> Acc (Array (sh:.Int) (Complex e)) idft v = let sh = shape v n = indexHead sh roots = inverseRootsOfUnity sh scale = lift (A.fromIntegral n :+ 0) in A.map (/scale) \$ dftG roots v -- | Generic function for computation of forward and inverse DFT. This function -- is also useful if you transform many arrays of the same extent, and don't -- want to recompute the roots for each one. -- -- The extent of the input and roots must match. -- dftG :: forall sh e. (Shape sh, Slice sh, A.RealFloat e) => Acc (Array (sh:.Int) (Complex e)) -- ^ roots of unity -> Acc (Array (sh:.Int) (Complex e)) -- ^ input array -> Acc (Array (sh:.Int) (Complex e)) dftG roots arr = A.fold (+) 0 \$ A.zipWith (*) arr' roots' where base = shape arr l = indexHead base extend = lift (base :. shapeSize base) -- Extend the entirety of the input arrays into a higher dimension, reading -- roots from the appropriate places and then reduce along this axis. -- -- In the calculation for 'roots'', 'i' is the index into the extended -- dimension, with corresponding base index 'ix' which we are attempting to -- calculate the single DFT value of. The rest proceeds as per 'dftGS'. -- arr' = A.generate extend (\ix' -> let i = indexHead ix' in arr !! i) roots' = A.generate extend (\ix' -> let ix :. i = unlift ix' sh :. n = unlift (fromIndex base i) :: Exp sh :. Exp Int k = indexHead ix in roots ! lift (sh :. (k*n) `mod` l)) -- | Compute a single value of the DFT. -- dftGS :: forall sh e. (Shape sh, Slice sh, A.RealFloat e) => Exp (sh :. Int) -- ^ index of the value we want -> Acc (Array (sh:.Int) (Complex e)) -- ^ roots of unity -> Acc (Array (sh:.Int) (Complex e)) -- ^ input array -> Acc (Scalar (Complex e)) dftGS ix roots arr = let k = indexHead ix l = indexHead (shape arr) -- all the roots we need to multiply with roots' = A.generate (shape arr) (\ix' -> let sh :. n = unlift ix' :: Exp sh :. Exp Int in roots ! lift (sh :. (k*n) `mod` l)) in A.foldAll (+) 0 \$ A.zipWith (*) arr roots'