Private.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
module Data.Array.Accelerate.FFTW.Private where

import qualified Numeric.FFTW.Batch as Batch
import qualified Numeric.FFTW.Shape as Spectrum
import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Accelerate.CArray.Conversion as Conv
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (Array, (:.)((:.)))

import qualified Data.Array.Comfort.Shape as CShape

import Foreign.Storable.Complex ()

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


type ShapeCyclic = CShape.Cyclic Int

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

newtype ComplexTransform sh a =
   ComplexTransform {runComplexTransform :: Transform sh (Complex a)}

type Element a ac = (Class.Real a, Conv.ComplexElement a ac)

complexTransform ::
   (A.Shape sh, CShape.C loop, Spectrum.MultiCyclic csh, Element a ac) =>
   (sh -> (loop, csh)) -> Batch.Sign -> ComplexTransform sh a
complexTransform bnds sign =
   ComplexTransform $ Conv.withCArrayComplex bnds $ Batch.fourier sign


dft ::
   (A.Shape sh, Element a ac) =>
   ComplexTransform (sh:.Int) a
dft = complexTransform bnds1 Batch.Forward

idft ::
   (A.Shape sh, Element a ac) =>
   ComplexTransform (sh:.Int) a
idft = complexTransform bnds1 Batch.Backward

bnds1 :: (A.Shape sh) => sh:.Int -> (Conv.Cube sh, ShapeCyclic)
bnds1 (sh:.n) = (Conv.Cube sh, CShape.Cyclic n)


dft2d ::
   (A.Shape sh, Element a ac) =>
   ComplexTransform (sh:.Int:.Int) a
dft2d = complexTransform bnds2 Batch.Forward

idft2d ::
   (A.Shape sh, Element a ac) =>
   ComplexTransform (sh:.Int:.Int) a
idft2d = complexTransform bnds2 Batch.Backward

bnds2 ::
   (A.Shape sh) => sh:.Int:.Int -> (Conv.Cube sh, (ShapeCyclic,ShapeCyclic))
bnds2 (sh:.n:.m) = (Conv.Cube sh, (CShape.Cyclic n, CShape.Cyclic m))