-- | Applying these transforms to the input of a DFT causes the output 
--   to be centered so that the zero frequency is in the middle. 
module Data.Array.Repa.Algorithms.DFT.Center
        ( center1d
        , center2d
        , center3d)
where
import Data.Array.Repa                          as R
import Data.Array.Repa.Algorithms.Complex       as R

-- | Apply the centering transform to a vector.
center1d
        :: Source r Complex
        => Array  r DIM1 Complex -> Array D DIM1 Complex
{-# INLINE center1d #-}
center1d :: Array r DIM1 Complex -> Array D DIM1 Complex
center1d Array r DIM1 Complex
arr
 = Array r DIM1 Complex
-> (DIM1 -> DIM1)
-> ((DIM1 -> Complex) -> DIM1 -> Complex)
-> Array D DIM1 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
R.traverse Array r DIM1 Complex
arr DIM1 -> DIM1
forall a. a -> a
id
        (\DIM1 -> Complex
get ix :: DIM1
ix@(DIM0
_ :. Int
x) -> ((-Complex
1) Complex -> Int -> Complex
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
x) Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
* DIM1 -> Complex
get DIM1
ix)


-- | Apply the centering transform to a matrix.
center2d
        :: Source r Complex
        => Array  r DIM2 Complex -> Array D DIM2 Complex
{-# INLINE center2d #-}
center2d :: Array r DIM2 Complex -> Array D DIM2 Complex
center2d Array r DIM2 Complex
arr
 = Array r DIM2 Complex
-> (DIM2 -> DIM2)
-> ((DIM2 -> Complex) -> DIM2 -> Complex)
-> Array D DIM2 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
R.traverse Array r DIM2 Complex
arr DIM2 -> DIM2
forall a. a -> a
id
        (\DIM2 -> Complex
get ix :: DIM2
ix@(DIM0
_ :. Int
y :. Int
x) -> ((-Complex
1) Complex -> Int -> Complex
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)) Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
* DIM2 -> Complex
get DIM2
ix)


-- | Apply the centering transform to a 3d array.
center3d 
        :: Source r Complex
        => Array  r DIM3 Complex -> Array D DIM3 Complex
{-# INLINE center3d #-}
center3d :: Array r DIM3 Complex -> Array D DIM3 Complex
center3d Array r DIM3 Complex
arr
 = Array r DIM3 Complex
-> (DIM3 -> DIM3)
-> ((DIM3 -> Complex) -> DIM3 -> Complex)
-> Array D DIM3 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
R.traverse Array r DIM3 Complex
arr DIM3 -> DIM3
forall a. a -> a
id
        (\DIM3 -> Complex
get ix :: DIM3
ix@(DIM0
_ :. Int
z :. Int
y :. Int
x) -> ((-Complex
1) Complex -> Int -> Complex
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
z Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)) Complex -> Complex -> Complex
forall a. Num a => a -> a -> a
* DIM3 -> Complex
get DIM3
ix)