module Math.FFT.Base where
import Control.Applicative
import Control.Arrow
import Control.Exception
import Control.Concurrent
import Control.Monad
import Data.Array.CArray
import Data.Array.CArray.Base ( shapeToStride, sBounds, mallocForeignPtrArrayAligned
, mapCArrayInPlace)
import Data.Complex
import Data.Bits
import Data.Generics
import Data.List
import Data.Typeable
import Foreign.C.Types
import Foreign.C.String
import Foreign.Marshal.Array
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import Foreign.Storable.Complex
import System.IO.Unsafe (unsafePerformIO)
class (Storable a, RealFloat a) => FFTWReal a where
plan_guru_dft :: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr (Complex a)
-> Ptr (Complex a) -> FFTWSign -> FFTWFlag -> IO Plan
plan_guru_dft_r2c :: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr a
-> Ptr (Complex a) -> FFTWFlag -> IO Plan
plan_guru_dft_c2r :: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr (Complex a)
-> Ptr a -> FFTWFlag -> IO Plan
plan_guru_r2r :: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr a
-> Ptr a -> Ptr FFTWKind -> FFTWFlag -> IO Plan
instance FFTWReal Double where
plan_guru_dft = c_plan_guru_dft
plan_guru_dft_r2c = c_plan_guru_dft_r2c
plan_guru_dft_c2r = c_plan_guru_dft_c2r
plan_guru_r2r = c_plan_guru_r2r
lock = unsafePerformIO $ newMVar ()
withLock :: IO a -> IO a
withLock = withMVar lock . const
type Plan = Ptr FFTWPlan
type FFTWPlan = ()
newtype Flag = Flag { unFlag :: FFTWFlag }
deriving (Eq, Show, Num, Bits)
type FFTWFlag = CUInt
c_measure :: FFTWFlag
c_measure = 0
c_destroy_input :: FFTWFlag
c_destroy_input = 1
c_unaligned :: FFTWFlag
c_unaligned = 2
c_conserve_memory :: FFTWFlag
c_conserve_memory = 4
c_exhaustive :: FFTWFlag
c_exhaustive = 8
c_preserve_input :: FFTWFlag
c_preserve_input = 16
c_patient :: FFTWFlag
c_patient = 32
c_estimate :: FFTWFlag
c_estimate = 64
nullFlag :: Flag
nullFlag = Flag 0
destroyInput :: Flag
destroyInput = Flag c_destroy_input
preserveInput :: Flag
preserveInput = Flag c_preserve_input
unaligned :: Flag
unaligned = Flag c_unaligned
conserveMemory :: Flag
conserveMemory = Flag c_conserve_memory
estimate :: Flag
estimate = Flag c_estimate
measure :: Flag
measure = Flag c_measure
patient :: Flag
patient = Flag c_patient
exhaustive :: Flag
exhaustive = Flag c_exhaustive
data Sign = DFTForward | DFTBackward
deriving (Eq,Show)
type FFTWSign = CInt
c_forward :: FFTWSign
c_forward = (1)
c_backward :: FFTWSign
c_backward = 1
unSign DFTForward = c_forward
unSign DFTBackward = c_backward
data Kind = R2HC | HC2R
| DHT
| REDFT00 | REDFT10 | REDFT01 | REDFT11
| RODFT00 | RODFT01 | RODFT10 | RODFT11
deriving (Eq,Show)
unKind :: Kind -> FFTWKind
unKind k = case k of
R2HC -> c_r2hc
HC2R -> c_hc2r
DHT -> c_dht
REDFT00 -> c_redft00
REDFT10 -> c_redft10
REDFT01 -> c_redft01
REDFT11 -> c_redft11
RODFT00 -> c_rodft00
RODFT01 -> c_rodft01
RODFT10 -> c_rodft10
RODFT11 -> c_rodft11
type FFTWKind = CInt
c_r2hc :: FFTWKind
c_r2hc = 0
c_hc2r :: FFTWKind
c_hc2r = 1
c_dht :: FFTWKind
c_dht = 2
c_redft00 :: FFTWKind
c_redft00 = 3
c_redft10 :: FFTWKind
c_redft10 = 5
c_redft01 :: FFTWKind
c_redft01 = 4
c_redft11 :: FFTWKind
c_redft11 = 6
c_rodft00 :: FFTWKind
c_rodft00 = 7
c_rodft10 :: FFTWKind
c_rodft10 = 9
c_rodft01 :: FFTWKind
c_rodft01 = 8
c_rodft11 :: FFTWKind
c_rodft11 = 10
data IODim = IODim { n :: Int, is :: Int, os :: Int }
deriving (Eq, Show)
instance Storable IODim where
sizeOf _ = (12)
alignment _ = alignment (undefined :: CInt)
peek p = do
n' <- (\hsc_ptr -> peekByteOff hsc_ptr 0) p
is' <- (\hsc_ptr -> peekByteOff hsc_ptr 4) p
os' <- (\hsc_ptr -> peekByteOff hsc_ptr 8) p
return (IODim n' is' os')
poke p (IODim n' is' os') = do
(\hsc_ptr -> pokeByteOff hsc_ptr 0) p n'
(\hsc_ptr -> pokeByteOff hsc_ptr 4) p is'
(\hsc_ptr -> pokeByteOff hsc_ptr 8) p os'
type TSpec = ([IODim],[IODim])
data DFT = CC | RC | CR | CRO | RR
deriving (Eq, Show)
check :: Plan -> IO ()
check p = when (p == nullPtr) . ioError $ userError "invalid plan"
execute :: Plan -> IO ()
execute p = check p >> c_execute p
unsafeNormalize :: (Ix i, Shapable i, Fractional e, Storable e)
=> [Int] -> CArray i e -> CArray i e
unsafeNormalize tdims a = mapCArrayInPlace (* s) a
where s = 1 / fromIntegral (product $ map (shape a !!) tdims)
dftG :: (FFTWReal r, Ix i, Shapable i) => Sign -> Flag -> [Int] -> CArray i (Complex r) -> CArray i (Complex r)
dftG s f tdims ain = case s of
DFTForward -> dftGU s f tdims ain
DFTBackward -> unsafeNormalize tdims (dftGU s f tdims ain)
dftCRG :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r
dftCRG f tdims ain = unsafeNormalize tdims (dftCRGU f tdims ain)
dftCROG :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r
dftCROG f tdims ain = unsafeNormalize tdims (dftCROGU f tdims ain)
dftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r)
dftN = dftG DFTForward estimate
idftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r)
idftN = dftG DFTBackward estimate
dftRCN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i (Complex r)
dftRCN = dftRCG estimate
dftCRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r
dftCRN = dftCRG estimate
dftCRON :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r
dftCRON = dftCROG estimate
fzr = flip zip . repeat
drr = (dftRRN .) . fzr
dftRRN :: (FFTWReal r, Ix i, Shapable i) => [(Int,Kind)] -> CArray i r -> CArray i r
dftRRN = dftRRG estimate
dftRHN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dftRHN = drr R2HC
dftHRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dftHRN = drr HC2R
dhtN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dhtN = drr DHT
dct1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dct1N = drr REDFT00
dct2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dct2N = drr REDFT01
dct3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dct3N = drr REDFT10
dct4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dct4N = drr REDFT11
dst1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dst1N = drr RODFT00
dst2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dst2N = drr RODFT01
dst3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dst3N = drr RODFT10
dst4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r
dst4N = drr RODFT11
dft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r)
dft = dftN [0]
idft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r)
idft = idftN [0]
dftRC :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i (Complex r)
dftRC = dftRCN [0]
dftCR :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r
dftCR = dftCRN [0]
dftCRO :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r
dftCRO = dftCRON [0]
dftRH :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dftRH = dftRHN [0]
dftHR :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dftHR = dftHRN [0]
dht :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dht = dhtN [0]
dct1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dct1 = dct1N [0]
dct2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dct2 = dct2N [0]
dct3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dct3 = dct3N [0]
dct4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dct4 = dct4N [0]
dst1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dst1 = dst1N [0]
dst2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dst2 = dst2N [0]
dst3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dst3 = dst3N [0]
dst4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r
dst4 = dst4N [0]
infix 7 `has`
a `has` b = a .&. b == b
transformCArray :: (Ix i, Storable a, Storable b)
=> Flag -> CArray i a -> (i,i) -> (FFTWFlag -> Ptr a -> Ptr b -> IO Plan) -> CArray i b
transformCArray f a lu planner = if f `has` estimate
&& not (any (f `has`) [measure, patient, exhaustive])
then go else transformCArray' f a lu planner
where go = unsafePerformIO $ do
ofp <- mallocForeignPtrArrayAligned (rangeSize lu)
withCArray a $ \ip ->
withForeignPtr ofp $ \op -> do
p <- withLock $ planner (unFlag f) ip op
execute p
unsafeForeignPtrToCArray ofp lu
transformCArray' :: (Ix i, Storable a, Storable b)
=> Flag -> CArray i a -> (i,i) -> (FFTWFlag -> Ptr a -> Ptr b -> IO Plan) -> CArray i b
transformCArray' f a lu planner = unsafePerformIO $ do
ofp <- mallocForeignPtrArrayAligned (rangeSize lu)
wfp <- mallocForeignPtrArrayAligned sz
withCArray a $ \ip ->
withForeignPtr ofp $ \op ->
withForeignPtr wfp $ \wp -> do
p <- withLock $ planner (unFlag $ f') wp op
copyArray wp ip sz
execute p
unsafeForeignPtrToCArray ofp lu
where sz = size a
f' = f .&. complement preserveInput .|. destroyInput
dftShape :: (Ix i, Shapable i, IArray CArray e)
=> DFT -> [Int] -> CArray i e -> ((i,i),TSpec)
dftShape dft tdims arr = assert valid (oBounds,tspec)
where shp = shape arr
rnk = rank arr
strides = shapeToStride shp
valid = not (null tdims) && 0 <= minimum tdims
&& maximum tdims < rnk && nub tdims == tdims
tspec = (d,d')
where d = zipWith3 IODim (filt lShape) (filt strides) (filt oStrides)
d' = zipWith3 IODim (filt' lShape) (filt' strides) (filt' oStrides)
filt s = map (s !!) tdims
filt' s = map (s !!) ([0 .. rnk 1] \\ tdims)
oShape = adjust f ldim shp
where f = case dft of
RC -> (\n -> n `div` 2 + 1)
CR -> (\n -> (n 1) * 2)
CRO -> (\n -> (n 1) * 2 + 1)
_ -> id
lShape = adjust f ldim shp
where f = case dft of
CR -> (\n -> (n 1) * 2)
CRO -> (\n -> (n 1) * 2 + 1)
_ -> id
oBounds = sBounds oShape
oStrides = shapeToStride oShape
ldim = last tdims
withTSpec :: TSpec -> (CInt -> Ptr IODim -> CInt -> Ptr IODim -> IO a) -> IO a
withTSpec (dims,dims') f = withArrayLen dims $ \r ds ->
withArrayLen dims' $ \hr hds ->
f (fromIntegral r) ds (fromIntegral hr) hds
adjust :: (a -> a) -> Int -> [a] -> [a]
adjust f i = uncurry (++) . second (\(x:xs) -> f x : xs) . splitAt i
dftGU :: (FFTWReal r, Ix i, Shapable i) => Sign -> Flag -> [Int] -> CArray i (Complex r) -> CArray i (Complex r)
dftGU s f tdims ain = transformCArray f ain bds go
where go f ip op = withTSpec tspec $ \r ds hr hds ->
plan_guru_dft r ds hr hds ip op (unSign s) f
(bds,tspec) = dftShape CC tdims ain
dftRCG :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i r -> CArray i (Complex r)
dftRCG f tdims ain = transformCArray f ain bds go
where go f ip op = withTSpec tspec $ \r ds hr hds ->
plan_guru_dft_r2c r ds hr hds ip op f
(bds,tspec) = dftShape RC tdims ain
dftCRG_ :: (FFTWReal r, Ix i, Shapable i) => Bool -> Flag -> [Int] -> CArray i (Complex r) -> CArray i r
dftCRG_ odd f tdims ain = tCArr f ain bds go
where go f ip op = withTSpec tspec $ \r ds hr hds ->
plan_guru_dft_c2r r ds hr hds ip op f
(bds,tspec) = dftShape (if odd then CRO else CR) tdims ain
tCArr = if length tdims == 1 && f `has` preserveInput
then transformCArray
else transformCArray'
dftCRGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r
dftCRGU = dftCRG_ False
dftCROGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r
dftCROGU = dftCRG_ True
dftRRG :: (FFTWReal r, Ix i, Shapable i) => Flag -> [(Int,Kind)] -> CArray i r -> CArray i r
dftRRG f tk ain = tCArr f ain bds go
where go f ip op = withTSpec tspec $ \r ds hr hds ->
withArray (map unKind ks) $ \pk ->
plan_guru_r2r r ds hr hds ip op pk f
(bds,tspec) = dftShape RR tdims ain
(tdims,ks) = unzip tk
tCArr = if any (== HC2R) ks && not (f `has` preserveInput)
then transformCArray'
else transformCArray
exportWisdomString :: IO String
exportWisdomString = do
pc <- c_export_wisdom_string
peekCString pc `finally` c_free pc
importWisdomString :: String -> IO Bool
importWisdomString str =
(==1) <$> withCString str c_import_wisdom_string
importWisdomSystem :: IO Bool
importWisdomSystem = (==1) <$> c_import_wisdom_system
foreign import ccall safe "fftw3.h fftw_plan_guru_dft" c_plan_guru_dft
:: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr (Complex Double)
-> Ptr (Complex Double) -> FFTWSign -> FFTWFlag -> IO Plan
foreign import ccall safe "fftw3.h fftw_plan_guru_dft_r2c" c_plan_guru_dft_r2c
:: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr Double
-> Ptr (Complex Double) -> FFTWFlag -> IO Plan
foreign import ccall safe "fftw3.h fftw_plan_guru_dft_c2r" c_plan_guru_dft_c2r
:: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr (Complex Double)
-> Ptr Double -> FFTWFlag -> IO Plan
foreign import ccall safe "fftw3.h fftw_plan_guru_r2r" c_plan_guru_r2r
:: CInt -> Ptr IODim -> CInt -> Ptr IODim -> Ptr Double
-> Ptr Double -> Ptr FFTWKind -> FFTWFlag -> IO Plan
foreign import ccall safe "fftw3.h fftw_execute" c_execute
:: Plan -> IO ()
foreign import ccall safe "fftw3.h fftw_execute_dft" c_execute_dft
:: Plan -> Ptr (Complex Double) -> Ptr (Complex Double) -> IO ()
foreign import ccall safe "fftw3.h fftw_execute_dft_r2c" c_execute_dft_r2c
:: Plan -> Ptr Double -> Ptr (Complex Double) -> IO ()
foreign import ccall safe "fftw3.h fftw_execute_dft_c2r" c_execute_dft_c2r
:: Plan -> Ptr (Complex Double) -> Ptr Double -> IO ()
foreign import ccall safe "fftw3.h fftw_execute_r2r" c_execute_r2r
:: Plan -> Ptr Double -> Ptr Double -> IO ()
foreign import ccall unsafe "fftw3.h fftw_export_wisdom_to_string"
c_export_wisdom_string :: IO CString
foreign import ccall unsafe "fftw3.h fftw_import_wisdom_from_string"
c_import_wisdom_string :: CString -> IO CInt
foreign import ccall unsafe "fftw3.h fftw_import_system_wisdom"
c_import_wisdom_system :: IO CInt
foreign import ccall unsafe "fftw3.h fftw_free" c_free :: Ptr a -> IO ()