{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
module Data.Array.Accelerate.Math.FFT (
Mode(..),
Numeric,
fft,
fft1D,
fft2D,
fft3D,
) where
import Data.Array.Accelerate as A
import Data.Array.Accelerate.Data.Complex
import Data.Array.Accelerate.Math.FFT.Type
import Data.Array.Accelerate.Math.FFT.Mode
import qualified Data.Array.Accelerate.Sugar.Shape as A ( rank, shapeR )
import qualified Data.Array.Accelerate.Math.FFT.Adhoc as Adhoc
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
import qualified Data.Array.Accelerate.Math.FFT.LLVM.Native as Native
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
import qualified Data.Array.Accelerate.Math.FFT.LLVM.PTX as PTX
#endif
import Prelude as P
fft :: forall sh e. (Shape sh, Slice sh, Numeric e)
=> Mode
-> Acc (Array (sh:.Int) (Complex e))
-> Acc (Array (sh:.Int) (Complex e))
fft :: Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
fft Mode
mode Acc (Array (sh :. Int) (Complex e))
arr
= let
scale :: Exp (Complex e)
scale = Exp Int -> Exp (Complex 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
indexHead (Acc (Array (sh :. Int) (Complex e)) -> Exp (sh :. Int)
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array (sh :. Int) (Complex e))
arr))
rank :: Int
rank = Shape (sh :. Int) => Int
forall sh. Shape sh => Int
A.rank @(sh:.Int)
shR :: ShapeR (EltR (sh :. Int))
shR = Shape (sh :. Int) => ShapeR (EltR (sh :. Int))
forall sh. Shape sh => ShapeR (EltR sh)
A.shapeR @(sh:.Int)
eR :: NumericR e
eR = Numeric e => NumericR e
forall a. Numeric a => NumericR a
numericR @e
go :: Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
go =
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
(if Int
rank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
P.<= Int
5 then ForeignAcc
(ArraysR (Array (sh :. Int) (Complex e))
-> ArraysR (Array (sh :. Int) (Complex e)))
-> (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> ShapeR (EltR sh, Int)
-> NumericR e
-> ForeignAcc
(Array (EltR sh, Int) (Vec2 e) -> Array (EltR sh, Int) (Vec2 e))
forall sh e.
HasCallStack =>
Mode
-> ShapeR sh
-> NumericR e
-> ForeignAcc (Array sh (Vec2 e) -> Array sh (Vec2 e))
Native.fft Mode
mode ShapeR (EltR sh, Int)
ShapeR (EltR (sh :. Int))
shR NumericR e
eR) else (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall a. a -> a
id) ((Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
(if Int
rank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
P.<= Int
3 then ForeignAcc
(ArraysR (Array (sh :. Int) (Complex e))
-> ArraysR (Array (sh :. Int) (Complex e)))
-> (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> ShapeR (EltR sh, Int)
-> NumericR e
-> ForeignAcc
(Array (EltR sh, Int) (Vec2 e) -> Array (EltR sh, Int) (Vec2 e))
forall sh e.
HasCallStack =>
Mode
-> ShapeR sh
-> NumericR e
-> ForeignAcc (Array sh (Vec2 e) -> Array sh (Vec2 e))
PTX.fft Mode
mode ShapeR (EltR sh, Int)
ShapeR (EltR (sh :. Int))
shR NumericR e
eR) else (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall a. a -> a
id) ((Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e)))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
in
case Mode
mode of
Mode
Inverse -> (Exp (Complex e) -> Exp (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
A.map (Exp (Complex e) -> Exp (Complex e) -> Exp (Complex e)
forall a. Fractional a => a -> a -> a
/Exp (Complex e)
scale) (Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
go Acc (Array (sh :. Int) (Complex e))
arr)
Mode
_ -> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
go Acc (Array (sh :. Int) (Complex e))
arr
fft1D :: forall e. Numeric e
=> Mode
-> Acc (Array DIM1 (Complex e))
-> Acc (Array DIM1 (Complex e))
fft1D :: Mode
-> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
fft1D Mode
mode Acc (Array DIM1 (Complex e))
arr
= let
scale :: Exp (Complex e)
scale = Exp Int -> Exp (Complex e)
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Acc (Array DIM1 (Complex e)) -> Exp Int
forall e. Elt e => Acc (Vector e) -> Exp Int
A.length Acc (Array DIM1 (Complex e))
arr)
eR :: NumericR e
eR = Numeric e => NumericR e
forall a. Numeric a => NumericR a
numericR @e
go :: Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
go =
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
ForeignAcc
(ArraysR (Array DIM1 (Complex e))
-> ArraysR (Array DIM1 (Complex e)))
-> (Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e))
-> Acc (Array DIM1 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM1 (Vec2 e) -> Array DIM1 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM1 (Vec2 e) -> Array DIM1 (Vec2 e))
Native.fft1D Mode
mode NumericR e
eR) ((Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> (Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e))
-> Acc (Array DIM1 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
ForeignAcc
(ArraysR (Array DIM1 (Complex e))
-> ArraysR (Array DIM1 (Complex e)))
-> (Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e))
-> Acc (Array DIM1 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM1 (Vec2 e) -> Array DIM1 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM1 (Vec2 e) -> Array DIM1 (Vec2 e))
PTX.fft1D Mode
mode NumericR e
eR) ((Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> (Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e)))
-> Acc (Array DIM1 (Complex e))
-> Acc (Array DIM1 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
Mode
-> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
in
case Mode
mode of
Mode
Inverse -> (Exp (Complex e) -> Exp (Complex e))
-> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
A.map (Exp (Complex e) -> Exp (Complex e) -> Exp (Complex e)
forall a. Fractional a => a -> a -> a
/Exp (Complex e)
scale) (Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
go Acc (Array DIM1 (Complex e))
arr)
Mode
_ -> Acc (Array DIM1 (Complex e)) -> Acc (Array DIM1 (Complex e))
go Acc (Array DIM1 (Complex e))
arr
fft2D :: forall e. Numeric e
=> Mode
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
fft2D :: Mode
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
fft2D Mode
mode Acc (Array DIM2 (Complex e))
arr
= let
scale :: Exp (Complex e)
scale = Exp Int -> Exp (Complex e)
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Acc (Array DIM2 (Complex e)) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
A.size Acc (Array DIM2 (Complex e))
arr)
eR :: NumericR e
eR = Numeric e => NumericR e
forall a. Numeric a => NumericR a
numericR @e
go :: Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
go =
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
ForeignAcc
(ArraysR (Array DIM2 (Complex e))
-> ArraysR (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM2 (Vec2 e) -> Array DIM2 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM2 (Vec2 e) -> Array DIM2 (Vec2 e))
Native.fft2D Mode
mode NumericR e
eR) ((Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
ForeignAcc
(ArraysR (Array DIM2 (Complex e))
-> ArraysR (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM2 (Vec2 e) -> Array DIM2 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM2 (Vec2 e) -> Array DIM2 (Vec2 e))
PTX.fft2D Mode
mode NumericR e
eR) ((Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
fft'
fft' :: Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
fft' Acc (Array DIM2 (Complex e))
a = Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall e. Elt e => Acc (Array DIM2 e) -> Acc (Array DIM2 e)
A.transpose (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mode
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
(Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall a b c.
(Arrays a, Arrays b, Arrays c) =>
(Acc a -> Acc b) -> (Acc b -> Acc c) -> Acc a -> Acc c
>-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall e. Elt e => Acc (Array DIM2 e) -> Acc (Array DIM2 e)
A.transpose (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e))
-> Acc (Array DIM2 (Complex e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mode
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
(Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e)))
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall a b. (a -> b) -> a -> b
$ Acc (Array DIM2 (Complex e))
a
in
case Mode
mode of
Mode
Inverse -> (Exp (Complex e) -> Exp (Complex e))
-> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
A.map (Exp (Complex e) -> Exp (Complex e) -> Exp (Complex e)
forall a. Fractional a => a -> a -> a
/Exp (Complex e)
scale) (Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
go Acc (Array DIM2 (Complex e))
arr)
Mode
_ -> Acc (Array DIM2 (Complex e)) -> Acc (Array DIM2 (Complex e))
go Acc (Array DIM2 (Complex e))
arr
fft3D :: forall e. Numeric e
=> Mode
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
fft3D :: Mode
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
fft3D Mode
mode Acc (Array DIM3 (Complex e))
arr
= let scale :: Exp (Complex e)
scale = Exp Int -> Exp (Complex e)
forall a b. (FromIntegral a b, Integral a) => Exp a -> Exp b
A.fromIntegral (Acc (Array DIM3 (Complex e)) -> Exp Int
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
A.size Acc (Array DIM3 (Complex e))
arr)
eR :: NumericR e
eR = Numeric e => NumericR e
forall a. Numeric a => NumericR a
numericR @e
go :: Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
go =
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
ForeignAcc
(ArraysR (Array DIM3 (Complex e))
-> ArraysR (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM3 (Vec2 e) -> Array DIM3 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM3 (Vec2 e) -> Array DIM3 (Vec2 e))
Native.fft3D Mode
mode NumericR e
eR) ((Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
ForeignAcc
(ArraysR (Array DIM3 (Complex e))
-> ArraysR (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall as bs (asm :: * -> *).
(Arrays as, Arrays bs, Foreign asm) =>
asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc (Mode
-> NumericR e
-> ForeignAcc (Array DIM3 (Vec2 e) -> Array DIM3 (Vec2 e))
forall e.
Mode
-> NumericR e
-> ForeignAcc (Array DIM3 (Vec2 e) -> Array DIM3 (Vec2 e))
PTX.fft3D Mode
mode NumericR e
eR) ((Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall a b. (a -> b) -> a -> b
$
#endif
Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
fft'
fft' :: Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
fft' Acc (Array DIM3 (Complex e))
a = Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall e. Elt e => Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mode
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
(Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall a b c.
(Arrays a, Arrays b, Arrays c) =>
(Acc a -> Acc b) -> (Acc b -> Acc c) -> Acc a -> Acc c
>-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall e. Elt e => Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mode
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
(Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall a b c.
(Arrays a, Arrays b, Arrays c) =>
(Acc a -> Acc b) -> (Acc b -> Acc c) -> Acc a -> Acc c
>-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall e. Elt e => Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e))
-> Acc (Array DIM3 (Complex e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mode
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall sh e.
(Shape sh, Slice sh, Numeric e) =>
Mode
-> Acc (Array (sh :. Int) (Complex e))
-> Acc (Array (sh :. Int) (Complex e))
Adhoc.fft Mode
mode
(Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e)))
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall a b. (a -> b) -> a -> b
$ Acc (Array DIM3 (Complex e))
a
in
case Mode
mode of
Mode
Inverse -> (Exp (Complex e) -> Exp (Complex e))
-> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
A.map (Exp (Complex e) -> Exp (Complex e) -> Exp (Complex e)
forall a. Fractional a => a -> a -> a
/Exp (Complex e)
scale) (Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
go Acc (Array DIM3 (Complex e))
arr)
Mode
_ -> Acc (Array DIM3 (Complex e)) -> Acc (Array DIM3 (Complex e))
go Acc (Array DIM3 (Complex e))
arr
rotate3D :: Elt e => Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D :: Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D Acc (Array DIM3 e)
arr = Exp DIM3
-> (Exp DIM3 -> Exp DIM3)
-> Acc (Array DIM3 e)
-> Acc (Array DIM3 e)
forall sh sh' a.
(Shape sh, Shape sh', Elt a) =>
Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute Exp DIM3
sh Exp DIM3 -> Exp DIM3
rot Acc (Array DIM3 e)
arr
where
sh :: Exp DIM3
sh :: Exp DIM3
sh =
let Z
Z :. Exp Int
z :. Exp Int
y :. Exp Int
x = Exp (Plain (((Z :. Exp Int) :. Exp Int) :. Exp Int))
-> ((Z :. Exp Int) :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift (Acc (Array DIM3 e) -> Exp DIM3
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape Acc (Array DIM3 e)
arr) :: Z :. Exp Int :. Exp Int :. Exp Int
in Exp Int -> Exp Int -> Exp Int -> Exp DIM3
forall i.
Elt i =>
Exp i -> Exp i -> Exp i -> Exp (((Z :. i) :. i) :. i)
index3 Exp Int
y Exp Int
x Exp Int
z
rot :: Exp DIM3 -> Exp DIM3
rot :: Exp DIM3 -> Exp DIM3
rot Exp DIM3
ix =
let Z
Z :. Exp Int
z :. Exp Int
y :. Exp Int
x = Exp (Plain (((Z :. Exp Int) :. Exp Int) :. Exp Int))
-> ((Z :. Exp Int) :. Exp Int) :. Exp Int
forall (c :: * -> *) e. Unlift c e => c (Plain e) -> e
unlift Exp (Plain (((Z :. Exp Int) :. Exp Int) :. Exp Int))
Exp DIM3
ix :: Z :. Exp Int :. Exp Int :. Exp Int
in Exp Int -> Exp Int -> Exp Int -> Exp DIM3
forall i.
Elt i =>
Exp i -> Exp i -> Exp i -> Exp (((Z :. i) :. i) :. i)
index3 Exp Int
x Exp Int
z Exp Int
y