{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
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
rootsOfUnity
:: (Shape sh, Slice sh, A.Floating e, A.FromIntegral Int e)
=> Exp (sh :. Int)
-> Acc (Array (sh:.Int) (Complex e))
rootsOfUnity sh =
let n = A.fromIntegral (A.indexHead sh)
in
A.generate sh (\ix -> let i = A.fromIntegral (A.indexHead ix)
k = 2 * pi * i / n
in
A.lift ( cos k :+ (-sin k) ))
inverseRootsOfUnity
:: (Shape sh, Slice sh, A.Floating e, A.FromIntegral Int e)
=> Exp (sh :. Int)
-> Acc (Array (sh:.Int) (Complex e))
inverseRootsOfUnity sh =
let n = A.fromIntegral (A.indexHead sh)
in
A.generate sh (\ix -> let i = A.fromIntegral (A.indexHead ix)
k = 2 * pi * i / n
in
A.lift ( cos k :+ sin k ))