Manifest.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ConstraintKinds #-}
module Data.Array.Accelerate.FFTW.Manifest (
   Transform,
   Element,
   dft,
   idft,
   dft2d,
   idft2d,
   ) where

import qualified Math.FFT as FFT
import Math.FFT.Base (FFTWReal)

import qualified Data.Array.Accelerate.CArray.Conversion as Conv

import qualified Data.Array.Accelerate.Sugar.Shape as Shape
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (Array, (:.)((:.)))

import Foreign.Storable.Complex ()

import Data.Array.Accelerate.Data.Complex (Complex)


{- $setup
>>> import qualified Data.Array.Accelerate.Utility.Lift.Exp as Exp
>>> import qualified Data.Array.Accelerate.Interpreter as AI
>>> import qualified Data.Array.Accelerate as A
>>> import Data.Array.Accelerate.Utility.Lift.Exp (expr)
>>> import Data.Array.Accelerate.Data.Complex (Complex((:+)))
>>> import Data.Array.Accelerate ((:.)((:.)))
>>>
>>> roundComplex :: A.Exp (Complex Float) -> A.Exp (Complex Int)
>>> roundComplex =
>>>    Exp.modify (expr:+expr) $ \(r:+i) -> A.round r :+ (A.round i :: A.Exp Int)
-}


type Bounds i = (i,i)

type Transform sh a = Array sh a -> Array sh a

type Element a = (FFTWReal a, Conv.ComplexElement a (Complex a))

{- |
>>> AI.run1 (A.map roundComplex) $ dft (A.fromList (A.Z:.3:.3) [1000,0,0, 0,1000,0, 0,0,1000])
Matrix (Z :. 3 :. 3)
  [ 1000 :+ 0,        1000 :+ 0,        1000 :+ 0,
    1000 :+ 0, (-500) :+ (-866),    (-500) :+ 866,
    1000 :+ 0,    (-500) :+ 866, (-500) :+ (-866)]
-}
dft ::
   (A.Shape sh, Element a) =>
   Transform (sh:.Int) (Complex a)
dft = Conv.withCArrayComplex bnds1 $ FFT.dftN [1]

{- |
>>> AI.run1 (A.map roundComplex) $ idft $ dft (A.fromList (A.Z:.3:.3) [1000,0,0, 0,1000,0, 0,0,1000])
Matrix (Z :. 3 :. 3)
  [ 1000 :+ 0,    0 :+ 0,    0 :+ 0,
       0 :+ 0, 1000 :+ 0,    0 :+ 0,
       0 :+ 0,    0 :+ 0, 1000 :+ 0]
-}
idft ::
   (A.Shape sh, Element a) =>
   Transform (sh:.Int) (Complex a)
idft = Conv.withCArrayComplex bnds1 $ FFT.idftN [1]

bnds1 :: (A.Shape sh) => sh:.Int -> Bounds (Int,Int)
bnds1 (sh:.n) = ((0,0), (Shape.size sh-1, n-1))


{- |
>>> AI.run1 (A.map roundComplex) $ dft2d (A.fromList (A.Z:.3:.3) [1000,0,0, 0,0,0, 0,0,0])
Matrix (Z :. 3 :. 3)
  [ 1000 :+ 0, 1000 :+ 0, 1000 :+ 0,
    1000 :+ 0, 1000 :+ 0, 1000 :+ 0,
    1000 :+ 0, 1000 :+ 0, 1000 :+ 0]

>>> AI.run1 (A.map roundComplex) $ dft2d (A.fromList (A.Z:.3:.3) [0,1000,0, 0,0,0, 0,0,0])
Matrix (Z :. 3 :. 3)
  [ 1000 :+ 0, (-500) :+ (-866), (-500) :+ 866,
    1000 :+ 0, (-500) :+ (-866), (-500) :+ 866,
    1000 :+ 0, (-500) :+ (-866), (-500) :+ 866]

>>> AI.run1 (A.map roundComplex) $ dft2d (A.fromList (A.Z:.3:.3) [1000,0,0, 0,1000,0, 0,0,1000])
Matrix (Z :. 3 :. 3)
  [ 3000 :+ 0,    0 :+ 0,    0 :+ 0,
       0 :+ 0,    0 :+ 0, 3000 :+ 0,
       0 :+ 0, 3000 :+ 0,    0 :+ 0]
-}
dft2d ::
   (A.Shape sh, Element a) =>
   Transform (sh:.Int:.Int) (Complex a)
dft2d =
   Conv.withCArrayComplex bnds2 $ FFT.dftN [1,2]

{- |
>>> AI.run1 (A.map roundComplex) $ idft2d $ dft2d (A.fromList (A.Z:.3:.3) [1000,0,0, 0,1000,0, 0,0,1000])
Matrix (Z :. 3 :. 3)
  [ 1000 :+ 0,    0 :+ 0,    0 :+ 0,
       0 :+ 0, 1000 :+ 0,    0 :+ 0,
       0 :+ 0,    0 :+ 0, 1000 :+ 0]
-}
idft2d ::
   (A.Shape sh, Element a) =>
   Transform (sh:.Int:.Int) (Complex a)
idft2d =
   Conv.withCArrayComplex bnds2 $ FFT.idftN [1,2]

bnds2 :: (A.Shape sh) => sh:.Int:.Int -> Bounds (Int,Int,Int)
bnds2 (sh:.n:.m) = ((0,0,0), (Shape.size sh-1, n-1, m-1))