{-# LANGUAGE TypeOperators, PatternGuards, RankNTypes, ScopedTypeVariables, BangPatterns, FlexibleContexts #-}
{-# OPTIONS -fno-warn-incomplete-patterns #-}

-- | Fast computation of Discrete Fourier Transforms using the Cooley-Tuckey algorithm.
--   Time complexity is O(n log n) in the size of the input.
--
--   This uses a naive divide-and-conquer algorithm, the absolute performance is about
--   50x slower than FFTW in estimate mode.
--
module Data.Array.Repa.Algorithms.FFT
        ( Mode(..)
        , isPowerOfTwo
        , fft3dP
        , fft2dP
        , fft1dP)
where

import Data.Array.Repa.Algorithms.Complex
import Data.Array.Repa                          as R
import Data.Array.Repa.Eval                     as R
import Data.Array.Repa.Unsafe                   as R

import Data.Bits                                ((.&.))
import Prelude                                  as P


data Mode
        = Forward
        | Reverse
        | Inverse
        deriving (Int -> Mode -> ShowS
[Mode] -> ShowS
Mode -> String
(Int -> Mode -> ShowS)
-> (Mode -> String) -> ([Mode] -> ShowS) -> Show Mode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Mode] -> ShowS
$cshowList :: [Mode] -> ShowS
show :: Mode -> String
$cshow :: Mode -> String
showsPrec :: Int -> Mode -> ShowS
$cshowsPrec :: Int -> Mode -> ShowS
Show, Mode -> Mode -> Bool
(Mode -> Mode -> Bool) -> (Mode -> Mode -> Bool) -> Eq Mode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Mode -> Mode -> Bool
$c/= :: Mode -> Mode -> Bool
== :: Mode -> Mode -> Bool
$c== :: Mode -> Mode -> Bool
Eq)


signOfMode :: Mode -> Double
signOfMode :: Mode -> Double
signOfMode Mode
mode
 = case Mode
mode of
        Mode
Forward         -> (-Double
1)
        Mode
Reverse         ->   Double
1
        Mode
Inverse         ->   Double
1
{-# INLINE signOfMode #-}


-- | Check if an `Int` is a power of two. Assumes `n` is a natural number.

-- The implementation can be found in Henry S. Warren, Jr.'s book
-- Hacker's delight, Chapter 2.
isPowerOfTwo :: Int -> Bool
isPowerOfTwo :: Int -> Bool
isPowerOfTwo Int
0 = Bool
True
isPowerOfTwo Int
1 = Bool
False
isPowerOfTwo Int
n = (Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE isPowerOfTwo #-}


-- 3D Transform -----------------------------------------------------------------------------------
-- | Compute the DFT of a 3d array. Array dimensions must be powers of two else `error`.
fft3dP  :: (Source r Complex, Monad m)
        => Mode
        -> Array r DIM3 Complex
        -> m (Array U DIM3 Complex)
fft3dP :: Mode -> Array r DIM3 Complex -> m (Array U DIM3 Complex)
fft3dP Mode
mode Array r DIM3 Complex
arr
 = let  DIM0
_ :. Int
depth :. Int
height :. Int
width   = Array r DIM3 Complex -> DIM3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r DIM3 Complex
arr
        !sign :: Double
sign   = Mode -> Double
signOfMode Mode
mode
        !scale :: Complex
scale  = Int -> Complex
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
height)

   in   if Bool -> Bool
not (Int -> Bool
isPowerOfTwo Int
depth Bool -> Bool -> Bool
&& Int -> Bool
isPowerOfTwo Int
height Bool -> Bool -> Bool
&& Int -> Bool
isPowerOfTwo Int
width)
         then String -> m (Array U DIM3 Complex)
forall a. HasCallStack => String -> a
error (String -> m (Array U DIM3 Complex))
-> String -> m (Array U DIM3 Complex)
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
                [ String
"Data.Array.Repa.Algorithms.FFT: fft3d"
                , String
"  Array dimensions must be powers of two,"
                , String
"  but the provided array is "
                        String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
height String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
width String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
depth ]

         else Array r DIM3 Complex
arr Array r DIM3 Complex
-> m (Array U DIM3 Complex) -> m (Array U DIM3 Complex)
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray`
                case Mode
mode of
                        Mode
Forward -> Array U DIM3 Complex -> m (Array U DIM3 Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U DIM3 Complex -> m (Array U DIM3 Complex))
-> Array U DIM3 Complex -> m (Array U DIM3 Complex)
forall a b. (a -> b) -> a -> b
$ Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array r DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign Array r DIM3 Complex
arr
                        Mode
Reverse -> Array U DIM3 Complex -> m (Array U DIM3 Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U DIM3 Complex -> m (Array U DIM3 Complex))
-> Array U DIM3 Complex -> m (Array U DIM3 Complex)
forall a b. (a -> b) -> a -> b
$ Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array r DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign Array r DIM3 Complex
arr
                        Mode
Inverse -> Array D DIM3 Complex -> m (Array U DIM3 Complex)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
                                (Array D DIM3 Complex -> m (Array U DIM3 Complex))
-> Array D DIM3 Complex -> m (Array U DIM3 Complex)
forall a b. (a -> b) -> a -> b
$  (Complex -> Complex)
-> Array U DIM3 Complex -> Array D DIM3 Complex
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
R.map (Complex -> Complex -> Complex
forall a. Fractional a => a -> a -> a
/ Complex
scale)
                                (Array U DIM3 Complex -> Array D DIM3 Complex)
-> Array U DIM3 Complex -> Array D DIM3 Complex
forall a b. (a -> b) -> a -> b
$  Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array U DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign (Array U DIM3 Complex -> Array U DIM3 Complex)
-> Array U DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double -> Array r DIM3 Complex -> Array U DIM3 Complex
forall r.
Source r Complex =>
Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign Array r DIM3 Complex
arr
{-# INLINE fft3dP #-}


fftTrans3d
        :: Source r Complex
        => Double
        -> Array r DIM3 Complex
        -> Array U DIM3 Complex

fftTrans3d :: Double -> Array r DIM3 Complex -> Array U DIM3 Complex
fftTrans3d Double
sign Array r DIM3 Complex
arr
 = let  ((DIM0 :. Int) :. Int
sh :. Int
len)     = Array r DIM3 Complex -> DIM3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r DIM3 Complex
arr
   in   Array D DIM3 Complex -> Array U DIM3 Complex
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> Array r2 sh e
suspendedComputeP (Array D DIM3 Complex -> Array U DIM3 Complex)
-> Array D DIM3 Complex -> Array U DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Array U DIM3 Complex -> Array D DIM3 Complex
forall r.
Source r Complex =>
Array r DIM3 Complex -> Array D DIM3 Complex
rotate3d (Array U DIM3 Complex -> Array D DIM3 Complex)
-> Array U DIM3 Complex -> Array D DIM3 Complex
forall a b. (a -> b) -> a -> b
$ Double
-> ((DIM0 :. Int) :. Int)
-> Int
-> Array r DIM3 Complex
-> Array U DIM3 Complex
forall sh r.
(Shape sh, Source r Complex) =>
Double
-> sh
-> Int
-> Array r (sh :. Int) Complex
-> Array U (sh :. Int) Complex
fft Double
sign (DIM0 :. Int) :. Int
sh Int
len Array r DIM3 Complex
arr
{-# INLINE fftTrans3d #-}


rotate3d
        :: Source r Complex
        => Array r DIM3 Complex -> Array D DIM3 Complex
rotate3d :: Array r DIM3 Complex -> Array D DIM3 Complex
rotate3d Array r DIM3 Complex
arr
 = DIM3
-> (DIM3 -> DIM3) -> Array r DIM3 Complex -> Array D DIM3 Complex
forall r sh1 sh2 e.
(Shape sh1, Source r e) =>
sh2 -> (sh2 -> sh1) -> Array r sh1 e -> Array D sh2 e
backpermute (DIM0
sh DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
m (DIM0 :. Int) -> Int -> (DIM0 :. Int) :. Int
forall tail head. tail -> head -> tail :. head
:. Int
k ((DIM0 :. Int) :. Int) -> Int -> DIM3
forall tail head. tail -> head -> tail :. head
:. Int
l) DIM3 -> DIM3
forall tail head head head.
(((tail :. head) :. head) :. head)
-> ((tail :. head) :. head) :. head
f Array r DIM3 Complex
arr
 where  (DIM0
sh :. Int
k :. Int
l :. Int
m)             = Array r DIM3 Complex -> DIM3
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r DIM3 Complex
arr
        f :: (((tail :. head) :. head) :. head)
-> ((tail :. head) :. head) :. head
f (tail
sh' :. head
m' :. head
k' :. head
l')       = tail
sh' tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
k' (tail :. head) -> head -> (tail :. head) :. head
forall tail head. tail -> head -> tail :. head
:. head
l' ((tail :. head) :. head)
-> head -> ((tail :. head) :. head) :. head
forall tail head. tail -> head -> tail :. head
:. head
m'
{-# INLINE rotate3d #-}



-- Matrix Transform -------------------------------------------------------------------------------
-- | Compute the DFT of a matrix. Array dimensions must be powers of two else `error`.
fft2dP  :: (Source r Complex, Monad m)
        => Mode
        -> Array r DIM2 Complex
        -> m (Array U DIM2 Complex)
fft2dP :: Mode
-> Array r ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
fft2dP Mode
mode Array r ((DIM0 :. Int) :. Int) Complex
arr
 = let  DIM0
_ :. Int
height :. Int
width    = Array r ((DIM0 :. Int) :. Int) Complex -> (DIM0 :. Int) :. Int
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r ((DIM0 :. Int) :. Int) Complex
arr
        sign :: Double
sign    = Mode -> Double
signOfMode Mode
mode
        scale :: Complex
scale   = Int -> Complex
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
height)

   in   if Bool -> Bool
not (Int -> Bool
isPowerOfTwo Int
height Bool -> Bool -> Bool
&& Int -> Bool
isPowerOfTwo Int
width)
         then String -> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall a. HasCallStack => String -> a
error (String -> m (Array U ((DIM0 :. Int) :. Int) Complex))
-> String -> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
                [ String
"Data.Array.Repa.Algorithms.FFT: fft2d"
                , String
"  Array dimensions must be powers of two,"
                , String
"  but the provided array is " String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
height String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
width ]

         else Array r ((DIM0 :. Int) :. Int) Complex
arr Array r ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray`
                case Mode
mode of
                        Mode
Forward -> Array U ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U ((DIM0 :. Int) :. Int) Complex
 -> m (Array U ((DIM0 :. Int) :. Int) Complex))
-> Array U ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ Double
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign (Array U ((DIM0 :. Int) :. Int) Complex
 -> Array U ((DIM0 :. Int) :. Int) Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign Array r ((DIM0 :. Int) :. Int) Complex
arr
                        Mode
Reverse -> Array U ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U ((DIM0 :. Int) :. Int) Complex
 -> m (Array U ((DIM0 :. Int) :. Int) Complex))
-> Array U ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ Double
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign (Array U ((DIM0 :. Int) :. Int) Complex
 -> Array U ((DIM0 :. Int) :. Int) Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign Array r ((DIM0 :. Int) :. Int) Complex
arr
                        Mode
Inverse -> Array D ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP (Array D ((DIM0 :. Int) :. Int) Complex
 -> m (Array U ((DIM0 :. Int) :. Int) Complex))
-> Array D ((DIM0 :. Int) :. Int) Complex
-> m (Array U ((DIM0 :. Int) :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ (Complex -> Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array D ((DIM0 :. Int) :. Int) Complex
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
R.map (Complex -> Complex -> Complex
forall a. Fractional a => a -> a -> a
/ Complex
scale) (Array U ((DIM0 :. Int) :. Int) Complex
 -> Array D ((DIM0 :. Int) :. Int) Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array D ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign (Array U ((DIM0 :. Int) :. Int) Complex
 -> Array U ((DIM0 :. Int) :. Int) Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign Array r ((DIM0 :. Int) :. Int) Complex
arr
{-# INLINE fft2dP #-}


fftTrans2d
        :: Source r Complex
        => Double
        -> Array r DIM2 Complex
        -> Array U DIM2 Complex

fftTrans2d :: Double
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
fftTrans2d Double
sign Array r ((DIM0 :. Int) :. Int) Complex
arr
 = let  (DIM0 :. Int
sh :. Int
len)     = Array r ((DIM0 :. Int) :. Int) Complex -> (DIM0 :. Int) :. Int
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r ((DIM0 :. Int) :. Int) Complex
arr
   in   Array D ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> Array r2 sh e
suspendedComputeP (Array D ((DIM0 :. Int) :. Int) Complex
 -> Array U ((DIM0 :. Int) :. Int) Complex)
-> Array D ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Array U ((DIM0 :. Int) :. Int) Complex
-> Array D ((DIM0 :. Int) :. Int) Complex
forall sh r e.
(Shape sh, Source r e) =>
Array r ((sh :. Int) :. Int) e -> Array D ((sh :. Int) :. Int) e
transpose (Array U ((DIM0 :. Int) :. Int) Complex
 -> Array D ((DIM0 :. Int) :. Int) Complex)
-> Array U ((DIM0 :. Int) :. Int) Complex
-> Array D ((DIM0 :. Int) :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> (DIM0 :. Int)
-> Int
-> Array r ((DIM0 :. Int) :. Int) Complex
-> Array U ((DIM0 :. Int) :. Int) Complex
forall sh r.
(Shape sh, Source r Complex) =>
Double
-> sh
-> Int
-> Array r (sh :. Int) Complex
-> Array U (sh :. Int) Complex
fft Double
sign DIM0 :. Int
sh Int
len Array r ((DIM0 :. Int) :. Int) Complex
arr
{-# INLINE fftTrans2d #-}


-- Vector Transform -------------------------------------------------------------------------------
-- | Compute the DFT of a vector. Array dimensions must be powers of two else `error`.
fft1dP  :: (Source r Complex, Monad m)
        => Mode
        -> Array r DIM1 Complex
        -> m (Array U DIM1 Complex)
fft1dP :: Mode
-> Array r (DIM0 :. Int) Complex
-> m (Array U (DIM0 :. Int) Complex)
fft1dP Mode
mode Array r (DIM0 :. Int) Complex
arr
 = let  DIM0
_ :. Int
len        = Array r (DIM0 :. Int) Complex -> DIM0 :. Int
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r (DIM0 :. Int) Complex
arr
        sign :: Double
sign    = Mode -> Double
signOfMode Mode
mode
        scale :: Complex
scale   = Int -> Complex
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len

   in   if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Bool
isPowerOfTwo Int
len
         then String -> m (Array U (DIM0 :. Int) Complex)
forall a. HasCallStack => String -> a
error (String -> m (Array U (DIM0 :. Int) Complex))
-> String -> m (Array U (DIM0 :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
                [ String
"Data.Array.Repa.Algorithms.FFT: fft1d"
                , String
"  Array dimensions must be powers of two, "
                , String
"  but the provided array is " String -> ShowS
forall a. [a] -> [a] -> [a]
P.++ Int -> String
forall a. Show a => a -> String
show Int
len ]

         else Array r (DIM0 :. Int) Complex
arr Array r (DIM0 :. Int) Complex
-> m (Array U (DIM0 :. Int) Complex)
-> m (Array U (DIM0 :. Int) Complex)
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray`
                case Mode
mode of
                        Mode
Forward -> Array U (DIM0 :. Int) Complex -> m (Array U (DIM0 :. Int) Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U (DIM0 :. Int) Complex
 -> m (Array U (DIM0 :. Int) Complex))
-> Array U (DIM0 :. Int) Complex
-> m (Array U (DIM0 :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
fftTrans1d Double
sign Array r (DIM0 :. Int) Complex
arr
                        Mode
Reverse -> Array U (DIM0 :. Int) Complex -> m (Array U (DIM0 :. Int) Complex)
forall sh r e (m :: * -> *).
(Shape sh, Source r e, Monad m) =>
Array r sh e -> m (Array r sh e)
now (Array U (DIM0 :. Int) Complex
 -> m (Array U (DIM0 :. Int) Complex))
-> Array U (DIM0 :. Int) Complex
-> m (Array U (DIM0 :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
fftTrans1d Double
sign Array r (DIM0 :. Int) Complex
arr
                        Mode
Inverse -> Array D (DIM0 :. Int) Complex -> m (Array U (DIM0 :. Int) Complex)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP (Array D (DIM0 :. Int) Complex
 -> m (Array U (DIM0 :. Int) Complex))
-> Array D (DIM0 :. Int) Complex
-> m (Array U (DIM0 :. Int) Complex)
forall a b. (a -> b) -> a -> b
$ (Complex -> Complex)
-> Array U (DIM0 :. Int) Complex -> Array D (DIM0 :. Int) Complex
forall sh r a b.
(Shape sh, Source r a) =>
(a -> b) -> Array r sh a -> Array D sh b
R.map (Complex -> Complex -> Complex
forall a. Fractional a => a -> a -> a
/ Complex
scale) (Array U (DIM0 :. Int) Complex -> Array D (DIM0 :. Int) Complex)
-> Array U (DIM0 :. Int) Complex -> Array D (DIM0 :. Int) Complex
forall a b. (a -> b) -> a -> b
$ Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
forall r.
Source r Complex =>
Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
fftTrans1d Double
sign Array r (DIM0 :. Int) Complex
arr
{-# INLINE fft1dP #-}


fftTrans1d
        :: Source r Complex
        => Double
        -> Array r DIM1 Complex
        -> Array U DIM1 Complex

fftTrans1d :: Double
-> Array r (DIM0 :. Int) Complex -> Array U (DIM0 :. Int) Complex
fftTrans1d Double
sign Array r (DIM0 :. Int) Complex
arr
 = let  (DIM0
sh :. Int
len)     = Array r (DIM0 :. Int) Complex -> DIM0 :. Int
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r (DIM0 :. Int) Complex
arr
   in   Double
-> DIM0
-> Int
-> Array r (DIM0 :. Int) Complex
-> Array U (DIM0 :. Int) Complex
forall sh r.
(Shape sh, Source r Complex) =>
Double
-> sh
-> Int
-> Array r (sh :. Int) Complex
-> Array U (sh :. Int) Complex
fft Double
sign DIM0
sh Int
len Array r (DIM0 :. Int) Complex
arr
{-# INLINE fftTrans1d #-}


-- Rank Generalised Worker ------------------------------------------------------------------------
fft     :: (Shape sh, Source r Complex)
        => Double -> sh -> Int
        -> Array r (sh :. Int) Complex
        -> Array U (sh :. Int) Complex

fft :: Double
-> sh
-> Int
-> Array r (sh :. Int) Complex
-> Array U (sh :. Int) Complex
fft !Double
sign !sh
sh !Int
lenVec !Array r (sh :. Int) Complex
vec
 = Int -> Int -> Int -> Array U (sh :. Int) Complex
go Int
lenVec Int
0 Int
1
 where  go :: Int -> Int -> Int -> Array U (sh :. Int) Complex
go !Int
len !Int
offset !Int
stride
         | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
         = Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> Array r2 sh e
suspendedComputeP (Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex)
-> Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall a b. (a -> b) -> a -> b
$ (sh :. Int)
-> ((sh :. Int) -> Complex) -> Array D (sh :. Int) Complex
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction (sh
sh sh -> Int -> sh :. Int
forall tail head. tail -> head -> tail :. head
:. Int
2) (sh :. Int) -> Complex
swivel

         | Bool
otherwise
         = Int
-> Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex
combine Int
len
                (Int -> Int -> Int -> Array U (sh :. Int) Complex
go (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int
offset            (Int
stride Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2))
                (Int -> Int -> Int -> Array U (sh :. Int) Complex
go (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) (Int
stride Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2))

         where  swivel :: (sh :. Int) -> Complex
swivel (sh
sh' :. Int
ix)
                 = case Int
ix of
                        Int
0       -> (Array r (sh :. Int) Complex
vec Array r (sh :. Int) Complex -> (sh :. Int) -> Complex
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` (sh
sh' sh -> Int -> sh :. Int
forall tail head. tail -> head -> tail :. head
:. Int
offset)) Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
+ (Array r (sh :. Int) Complex
vec Array r (sh :. Int) Complex -> (sh :. Int) -> Complex
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` (sh
sh' sh -> Int -> sh :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride)))
                        Int
1       -> (Array r (sh :. Int) Complex
vec Array r (sh :. Int) Complex -> (sh :. Int) -> Complex
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` (sh
sh' sh -> Int -> sh :. Int
forall tail head. tail -> head -> tail :. head
:. Int
offset)) Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
- (Array r (sh :. Int) Complex
vec Array r (sh :. Int) Complex -> (sh :. Int) -> Complex
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh -> e
`unsafeIndex` (sh
sh' sh -> Int -> sh :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride)))

                {-# INLINE combine #-}
                combine :: Int
-> Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex
combine !Int
len'   Array U (sh :. Int) Complex
evens Array U (sh :. Int) Complex
odds
                 = Array U (sh :. Int) Complex
evens Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray` Array U (sh :. Int) Complex
odds Array U (sh :. Int) Complex
-> Array U (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall r e sh b. (Source r e, Shape sh) => Array r sh e -> b -> b
`deepSeqArray`
                   let  odds' :: Array D (sh :. Int) Complex
odds'   = Array U (sh :. Int) Complex
-> ((sh :. Int) -> sh :. Int)
-> (((sh :. Int) -> Complex) -> (sh :. Int) -> Complex)
-> Array D (sh :. Int) Complex
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array U (sh :. Int) Complex
odds (sh :. Int) -> sh :. Int
forall a. a -> a
id (\(sh :. Int) -> Complex
get ix :: sh :. Int
ix@(sh
_ :. Int
k) -> Double -> Int -> Int -> Complex
twiddle Double
sign Int
k Int
len' Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
* (sh :. Int) -> Complex
get sh :. Int
ix)
                   in   Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> Array r2 sh e
suspendedComputeP (Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex)
-> Array D (sh :. Int) Complex -> Array U (sh :. Int) Complex
forall a b. (a -> b) -> a -> b
$ (Array U (sh :. Int) Complex
evens Array U (sh :. Int) Complex
-> Array D (sh :. Int) Complex -> Array D (sh :. Int) Complex
forall sh r1 c r2.
(Shape sh, Source r1 c, Source r2 c, Num c) =>
Array r1 sh c -> Array r2 sh c -> Array D sh c
+^ Array D (sh :. Int) Complex
odds') Array D (sh :. Int) Complex
-> Array D (sh :. Int) Complex -> Array D (sh :. Int) Complex
forall sh r1 e r2.
(Shape sh, Source r1 e, Source r2 e) =>
Array r1 (sh :. Int) e
-> Array r2 (sh :. Int) e -> Array D (sh :. Int) e
R.++ (Array U (sh :. Int) Complex
evens Array U (sh :. Int) Complex
-> Array D (sh :. Int) Complex -> Array D (sh :. Int) Complex
forall sh r1 c r2.
(Shape sh, Source r1 c, Source r2 c, Num c) =>
Array r1 sh c -> Array r2 sh c -> Array D sh c
-^ Array D (sh :. Int) Complex
odds')
{-# INLINE fft #-}


-- Compute a twiddle factor.
twiddle :: Double
        -> Int                  -- index
        -> Int                  -- length
        -> Complex

twiddle :: Double -> Int -> Int -> Complex
twiddle Double
sign Int
k' Int
n'
        =  (Double -> Double
forall a. Floating a => a -> a
cos (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n), Double
sign Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin  (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n))
        where   k :: Double
k       = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k'
                n :: Double
n       = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n'
{-# INLINE twiddle #-}