{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
-- |
-- Module      : Data.Array.Accelerate.Math.DFT.Roots
-- Copyright   : [2012..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.Math.DFT.Roots (

  rootsOfUnity, inverseRootsOfUnity,

) where

import Prelude                                  as P
import Data.Array.Accelerate                    as A
import Data.Array.Accelerate.Data.Complex


-- | Calculate the roots of unity for the forward transform
--
rootsOfUnity
    :: (Shape sh, Slice sh, A.Floating e, A.FromIntegral Int e)
    => Exp (sh :. Int)
    -> Acc (Array (sh:.Int) (Complex e))
rootsOfUnity :: Exp (sh :. Int) -> Acc (Array (sh :. Int) (Complex e))
rootsOfUnity Exp (sh :. Int)
sh =
  let n :: Exp e
n = Exp Int -> Exp e
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Exp (sh :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
A.indexHead Exp (sh :. Int)
sh)
  in
  Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
A.generate Exp (sh :. Int)
sh (\Exp (sh :. Int)
ix -> let i :: Exp e
i = Exp Int -> Exp e
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Exp (sh :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
A.indexHead Exp (sh :. Int)
ix)
                            k :: Exp e
k = Exp e
2 Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
* Exp e
forall a. Floating a => a
pi Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
* Exp e
i Exp e -> Exp e -> Exp e
forall a. Fractional a => a -> a -> a
/ Exp e
n
                        in
                        Complex (Exp e) -> Exp (Plain (Complex (Exp e)))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
A.lift ( Exp e -> Exp e
forall a. Floating a => a -> a
cos Exp e
k Exp e -> Exp e -> Complex (Exp e)
forall a. a -> a -> Complex a
:+ (-Exp e -> Exp e
forall a. Floating a => a -> a
sin Exp e
k) ))


-- | Calculate the roots of unity for an inverse transform
--
inverseRootsOfUnity
    :: (Shape sh, Slice sh, A.Floating e, A.FromIntegral Int e)
    => Exp (sh :. Int)
    -> Acc (Array (sh:.Int) (Complex e))
inverseRootsOfUnity :: Exp (sh :. Int) -> Acc (Array (sh :. Int) (Complex e))
inverseRootsOfUnity Exp (sh :. Int)
sh =
  let n :: Exp e
n = Exp Int -> Exp e
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Exp (sh :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
A.indexHead Exp (sh :. Int)
sh)
  in
  Exp (sh :. Int)
-> (Exp (sh :. Int) -> Exp (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall sh a.
(Shape sh, Elt a) =>
Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
A.generate Exp (sh :. Int)
sh (\Exp (sh :. Int)
ix -> let i :: Exp e
i = Exp Int -> Exp e
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Exp (sh :. Int) -> Exp Int
forall sh a. (Elt sh, Elt a) => Exp (sh :. a) -> Exp a
A.indexHead Exp (sh :. Int)
ix)
                            k :: Exp e
k = Exp e
2 Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
* Exp e
forall a. Floating a => a
pi Exp e -> Exp e -> Exp e
forall a. Num a => a -> a -> a
* Exp e
i Exp e -> Exp e -> Exp e
forall a. Fractional a => a -> a -> a
/ Exp e
n
                        in
                        Complex (Exp e) -> Exp (Plain (Complex (Exp e)))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
A.lift ( Exp e -> Exp e
forall a. Floating a => a -> a
cos Exp e
k Exp e -> Exp e -> Complex (Exp e)
forall a. a -> a -> Complex a
:+ Exp e -> Exp e
forall a. Floating a => a -> a
sin Exp e
k ))