diff --git a/Math/FFT.hs b/Math/FFT.hs deleted file mode 100644 --- a/Math/FFT.hs +++ /dev/null @@ -1,170 +0,0 @@ --- | --- Module : Math.FFT --- Copyright : (c) 2008 Jed Brown --- License : BSD-style --- --- Maintainer : jed@59A2.org --- Stability : experimental --- Portability : non-portable --- --- This module exposes an interface to FFTW, the Fastest Fourier Transform in --- the West. --- --- These bindings present several levels of interface. All the higher level --- functions ('dft', 'idft', 'dftN', ...) are easily derived from the general --- functions ('dftG', 'dftRCG', ...). Only the general functions let you --- specify planner flags. The higher levels all set 'estimate' so you should --- not have to wait through time consuming planning (see below for more). --- --- The simplest interface is the one-dimensional transforms. If you supply a --- multi-dimensional array, these will only transform the first dimension. --- These functions only take one argument, the array to be transformed. --- --- At the next level, we have multi-dimensional transforms where you specify --- which dimensions to transform in and the array to transform. For instance --- --- > b = dftRCN [0,2] a --- --- is the real to complex transform in dimensions 0 and 2 of the array @a@ which --- must be at least rank 3. The array @b@ will be complex valued with the same --- extent as @a@ in every dimension except @2@. If @a@ had extent @n@ in --- dimension @2@ then the @b@ will have extent @a `div` 2 + 1@ which consists of --- all non-negative frequency components in this dimension (the negative --- frequencies are conjugate to the positive frequencies because of symmetry --- since @a@ is real valued). --- --- The real to real transforms allow different transform kinds in each --- transformed dimension. For example, --- --- > b = dftRRN [(0,DHT), (1,REDFT10), (2,RODFT11)] a --- --- is a Discrete Hartley Transform in dimension 0, a discrete cosine transform --- (DCT-2) in dimension 1, and distrete sine transform (DST-4) in dimension 2 --- where the array @a@ must have rank at least 3. --- --- The general interface is similar to the multi-dimensional interface, takes as --- its first argument, a bitwise '.|.' of planning 'Flag's. (In the complex --- version, the sign of the transform is first.) For example, --- --- > b = dftG DFTBackward (patient .|. destroy_input) [1,2] a --- --- is an inverse DFT in dimensions 1 and 2 of the complex array @a@ which has --- rank at least 3. It will use the patient planner to generate a (near) --- optimal transform. If you compute the same type of transform again, it --- should be very fast since the plan is cached. --- --- Inverse transforms are typically normalized. The un-normalized inverse --- transforms are 'dftGU', 'dftCRGU' and 'dftCROGU'. For example --- --- > b = dftCROGU measure [0,1] a --- --- is an un-normalized inverse DFT in dimensions 0 and 1 of the complex array --- @a@ (representing the non-negative frequencies, where the negative --- frequencies are conjugate) which has rank at least 2. Here, dimension 1 is --- logically odd so if @a@ has extent @n@ in dimension 1, then @b@ will have --- extent @(n - 1) * 2 + 1@ in dimension 1. It is more common that the logical --- dimension is even, in which case we would use 'dftCRGU' in which case @b@ --- would have extent @(n - 1) * 2@ in dimension @1@. --- --- --- The FFTW library separates transforms into two steps. First you compute a --- plan for a given transform, then you execute it. Often the planning stage is --- quite time-consuming, but subsequent transforms of the same size and type --- will be extremely fast. The planning phase actually computes transforms, so --- it overwrites its input array. For many C codes, it is reasonable to re-use --- the same arrays to compute a given transform on different data. This is not --- a very useful paradigm from Haskell. Fortunately, FFTW caches its plans so --- if try to generate a new plan for a transform size which has already been --- planned, the planner will return immediately. Unfortunately, it is not --- possible to consult the cache, so if a plan is cached, we may use more memory --- than is strictly necessary since we must allocate a work array which we --- expect to be overwritten during planning. FFTW can export its cached plans --- to a string. This is known as wisdom. For high performance work, it is a --- good idea to compute plans of the sizes you are interested in, using --- aggressive options (i.e. 'patient'), use 'exportWisdomString' to get a string --- representing these plans, and write this to a file. Then for production --- runs, you can read this in, then add it to the cache with --- 'importWisdomString'. Now you can use the 'estimate' planner so the Haskell --- bindings know that FFTW will not overwrite the input array, and you will --- still get a high quality transform (because it has wisdom). - -module Math.FFT ( - -- * Data types - Sign(..), - Kind(..), - -- * Planner flags - -- ** Algorithm restriction flags - destroyInput, - preserveInput, - -- ** Planning rigor flags - estimate, - measure, - patient, - exhaustive, - - -- * DFT of complex data - -- ** DFT in first dimension only - dft, - idft, - -- ** Multi-dimensional transforms - dftN, - idftN, - -- ** General transform - dftG, - -- ** Un-normalized general transform - dftGU, - - -- * DFT of real data - -- ** DFT in first dimension only - dftRC, - dftCR, - dftCRO, - -- ** Multi-dimensional transforms - dftRCN, - dftCRN, - dftCRON, - -- ** General transform - dftRCG, - dftCRG, - dftCROG, - -- ** Un-normalized general transform - dftCRGU, - dftCROGU, - - -- * Real to real transforms (all un-normalized) - -- ** Transforms in first dimension only - dftRH, - dftHR, - dht, - dct1, - dct2, - dct3, - dct4, - dst1, - dst2, - dst3, - dst4, - -- ** Multi-dimensional transforms with the same transform type in each dimension - dftRHN, - dftHRN, - dhtN, - dct1N, - dct2N, - dct3N, - dct4N, - dst1N, - dst2N, - dst3N, - dst4N, - -- ** Multi-dimensional transforms with possibly different transforms in each dimension - dftRRN, - -- ** General transforms - dftRRG, - - -- * Wisdom - importWisdomString, - importWisdomSystem, - exportWisdomString, -) where - -import Math.FFT.Base diff --git a/Math/FFT/Base.hsc b/Math/FFT/Base.hsc deleted file mode 100644 --- a/Math/FFT/Base.hsc +++ /dev/null @@ -1,592 +0,0 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, DeriveDataTypeable - , FlexibleContexts, NoMonomorphismRestriction #-} -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 (mallocForeignPtrArrayAligned, mapCArrayInPlace) -import Data.Ix.Shapable (shapeToStride, sBounds) -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) - -#include - --- | Our API is polymorphic over the real data type. FFTW, at least in --- principle, supports single precision 'Float', double precision 'Double' and --- long double 'CLDouble' (presumable?). -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 - --- | Using this instance requires linking with @-lfftw3@. -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 - --- | This lock must be taken during /planning/ of any transform. The FFTW --- library is not thread-safe in the planning phase. Thankfully, the lock is --- not needed during the execute phase. -lock :: MVar () -lock = unsafePerformIO $ newMVar () -{-# NOINLINE lock #-} - -withLock :: IO a -> IO a -withLock = withMVar lock . const - --- | A plan is an opaque foreign object. -type Plan = Ptr FFTWPlan - -type FFTWPlan = () - --- | The 'Flag' type is used to influence the kind of plans which are created. --- To specify multiple flags, use a bitwise '.|.'. -newtype Flag = Flag { unFlag :: FFTWFlag } - deriving (Eq, Show, Num, Bits) - -type FFTWFlag = CUInt - -#{enum FFTWFlag, - , c_measure = FFTW_MEASURE - , c_destroy_input = FFTW_DESTROY_INPUT - , c_unaligned = FFTW_UNALIGNED - , c_conserve_memory = FFTW_CONSERVE_MEMORY - , c_exhaustive = FFTW_EXHAUSTIVE - , c_preserve_input = FFTW_PRESERVE_INPUT - , c_patient = FFTW_PATIENT - , c_estimate = FFTW_ESTIMATE } - --- | Default flag. For most transforms, this is equivalent to setting 'measure' --- and 'preserveInput'. The exceptions are complex to real and half-complex to --- real transforms. -nullFlag :: Flag -nullFlag = Flag 0 - --- --- Algorithm restriction flags --- - --- | Allows FFTW to overwrite the input array with arbitrary data; this can --- sometimes allow more efficient algorithms to be employed. --- --- Setting this flag implies that two memory allocations will be done, one for --- work space, and one for the result. When 'estimate' is not set, we will be --- doing two memory allocations anyway, so we set this flag as well (since we --- don't retain the work array anyway). -destroyInput :: Flag -destroyInput = Flag c_destroy_input - --- | 'preserveInput' specifies that an out-of-place transform must not change --- its input array. This is ordinarily the default, except for complex to real --- transforms for which 'destroyInput' is the default. In the latter cases, --- passing 'preserveInput' will attempt to use algorithms that do not destroy --- the input, at the expense of worse performance; for multi-dimensional complex --- to real transforms, however, no input-preserving algorithms are implemented --- so the Haskell bindings will set 'destroyInput' and do a transform with two --- memory allocations. -preserveInput :: Flag -preserveInput = Flag c_preserve_input - --- | Instruct FFTW not to generate a plan which uses SIMD instructions, even if --- the memory you are planning with is aligned. This should only be needed if --- you are using the guru interface and want to reuse a plan with memory that --- may be unaligned (i.e. you constructed the 'CArray' with --- 'unsafeForeignPtrToCArray'). -unaligned :: Flag -unaligned = Flag c_unaligned - --- | The header claims that this flag is documented, but in reality, it is not. --- I don't know what it does and it is here only for completeness. -conserveMemory :: Flag -conserveMemory = Flag c_conserve_memory - --- --- Planning rigor flags --- - --- | 'estimate' specifies that, instead of actual measurements of different --- algorithms, a simple heuristic is used to pick a (probably sub-optimal) plan --- quickly. With this flag, the input/output arrays are not overwritten during --- planning. --- --- This is the only planner flag for which a single memory allocation is possible. -estimate :: Flag -estimate = Flag c_estimate - --- | 'measure' tells FFTW to find an optimized plan by actually computing --- several FFTs and measuring their execution time. Depending on your machine, --- this can take some time (often a few seconds). 'measure' is the default --- planning option. -measure :: Flag -measure = Flag c_measure - --- | 'patient' is like 'measure', but considers a wider range of algorithms and --- often produces a "more optimal" plan (especially for large transforms), but --- at the expense of several times longer planning time (especially for large --- transforms). -patient :: Flag -patient = Flag c_patient - --- | 'exhaustive' is like 'patient' but considers an even wider range of --- algorithms, including many that we think are unlikely to be fast, to --- produce the most optimal plan but with a substantially increased planning --- time. -exhaustive :: Flag -exhaustive = Flag c_exhaustive - --- | Determine which direction of DFT to execute. -data Sign = DFTForward | DFTBackward - deriving (Eq,Show) - -type FFTWSign = CInt - -#{enum FFTWSign, - , c_forward = FFTW_FORWARD - , c_backward = FFTW_BACKWARD } - -unSign :: Sign -> FFTWSign -unSign DFTForward = c_forward -unSign DFTBackward = c_backward - --- | Real to Real transform kinds. -data Kind = R2HC | HC2R -- half-complex transforms - | DHT -- discrete Hartley transformm - | REDFT00 | REDFT10 | REDFT01 | REDFT11 -- discrete cosine transforms - | RODFT00 | RODFT01 | RODFT10 | RODFT11 -- discrete sine transforms - 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 - -#{enum FFTWKind, - , c_r2hc = FFTW_R2HC - , c_hc2r = FFTW_HC2R - , c_dht = FFTW_DHT - , c_redft00 = FFTW_REDFT00 - , c_redft10 = FFTW_REDFT10 - , c_redft01 = FFTW_REDFT01 - , c_redft11 = FFTW_REDFT11 - , c_rodft00 = FFTW_RODFT00 - , c_rodft10 = FFTW_RODFT10 - , c_rodft01 = FFTW_RODFT01 - , c_rodft11 = FFTW_RODFT11 } - --- | Corresponds to the @fftw_iodim@ structure. It completely describes the --- layout of each dimension, before and after the transform. -data IODim = IODim { nIODim :: Int -- ^ Logical size of dimension - , isIODim :: Int -- ^ Stride along dimension in input array - , osIODim :: Int -- ^ Stride along dimension in output array - } - deriving (Eq, Show, Data, Typeable) - -instance Storable IODim where - sizeOf _ = #{size fftw_iodim} - alignment _ = alignment (undefined :: CInt) - peek p = do - n' <- #{peek fftw_iodim, n} p - is' <- #{peek fftw_iodim, is} p - os' <- #{peek fftw_iodim, os} p - return (IODim n' is' os') - poke p (IODim n' is' os') = do - #{poke fftw_iodim, n} p n' - #{poke fftw_iodim, is} p is' - #{poke fftw_iodim, os} p os' - --- | Tuple of transform dimensions and non-transform dimensions of the array. -type TSpec = ([IODim],[IODim]) - --- | Types of transforms. Used to control 'dftShape'. -data DFT = CC | RC | CR | CRO | RR - deriving (Eq, Show) - --- | Verify that a plan is valid. Thows an exception if not. -check :: Plan -> IO () -check p = when (p == nullPtr) . ioError $ userError "invalid plan" - --- | Confirm that the plan is valid, then execute the transform. -execute :: Plan -> IO () -execute p = check p >> c_execute p - --- | In-place normalization outside of IO. You must be able to prove that no --- reference to the original can be retained. -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) - --- | Normalized general complex DFT -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) - --- | Normalized general complex to real DFT where the last transformed dimension --- is logically even. -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) - --- | Normalized general complex to real DFT where the last transformed dimension --- is logicall odd. -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) - --- | Multi-dimensional forward DFT. -dftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r) -dftN = dftG DFTForward estimate --- | Multi-dimensional inverse DFT. -idftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r) -idftN = dftG DFTBackward estimate --- | Multi-dimensional forward DFT of real data. -dftRCN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i (Complex r) -dftRCN = dftRCG estimate --- | Multi-dimensional inverse DFT of Hermitian-symmetric data (where only the --- non-negative frequencies are given). -dftCRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r -dftCRN = dftCRG estimate --- | Multi-dimensional inverse DFT of Hermitian-symmetric data (where only the --- non-negative frequencies are given) and the last transformed dimension is --- logically odd. -dftCRON :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r -dftCRON = dftCROG estimate - -fzr :: b -> [a] -> [(a,b)] -fzr = flip zip . repeat -drr :: (FFTWReal r, Ix i, Shapable i) => Kind -> [Int] -> CArray i r -> CArray i r -drr = (dftRRN .) . fzr - --- | Multi-dimensional real to real transform. The result is not normalized. -dftRRN :: (FFTWReal r, Ix i, Shapable i) => [(Int,Kind)] -> CArray i r -> CArray i r -dftRRN = dftRRG estimate - --- --- The following do the same type of transform in each dimension specified. --- --- | Multi-dimensional real to half-complex transform. The result is not normalized. -dftRHN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dftRHN = drr R2HC --- | Multi-dimensional half-complex to real transform. The result is not normalized. -dftHRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dftHRN = drr HC2R --- | Multi-dimensional Discrete Hartley Transform. The result is not normalized. -dhtN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dhtN = drr DHT --- | Multi-dimensional Type 1 discrete cosine transform. -dct1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dct1N = drr REDFT00 --- | Multi-dimensional Type 2 discrete cosine transform. This is commonly known --- as /the/ DCT. -dct2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dct2N = drr REDFT10 --- | Multi-dimensional Type 3 discrete cosine transform. This is commonly known --- as /the/ inverse DCT. The result is not normalized. -dct3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dct3N = drr REDFT01 --- | Multi-dimensional Type 4 discrete cosine transform. -dct4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dct4N = drr REDFT11 --- | Multi-dimensional Type 1 discrete sine transform. -dst1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dst1N = drr RODFT00 --- | Multi-dimensional Type 2 discrete sine transform. -dst2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dst2N = drr RODFT10 --- | Multi-dimensional Type 3 discrete sine transform. -dst3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dst3N = drr RODFT01 --- | Multi-dimensional Type 4 discrete sine transform. -dst4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r -dst4N = drr RODFT11 - --- --- Transform in the first dimension only. --- - --- | 1-dimensional complex DFT. -dft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r) -dft = dftN [0] --- | 1-dimensional complex inverse DFT. Inverse of 'dft'. -idft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r) -idft = idftN [0] --- | 1-dimensional real to complex DFT. -dftRC :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i (Complex r) -dftRC = dftRCN [0] --- | 1-dimensional complex to real DFT with logically even dimension. Inverse of 'dftRC'. -dftCR :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r -dftCR = dftCRN [0] --- | 1-dimensional complex to real DFT with logically odd dimension. Inverse of 'dftRC'. -dftCRO :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r -dftCRO = dftCRON [0] --- | 1-dimensional real to half-complex DFT. -dftRH :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dftRH = dftRHN [0] --- | 1-dimensional half-complex to real DFT. Inverse of 'dftRH' after normalization. -dftHR :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dftHR = dftHRN [0] --- | 1-dimensional Discrete Hartley Transform. Self-inverse after normalization. -dht :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dht = dhtN [0] --- | 1-dimensional Type 1 discrete cosine transform. -dct1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dct1 = dct1N [0] --- | 1-dimensional Type 2 discrete cosine transform. This is commonly known as /the/ DCT. -dct2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dct2 = dct2N [0] --- | 1-dimensional Type 3 discrete cosine transform. This is commonly known as /the/ inverse DCT. -dct3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dct3 = dct3N [0] --- | 1-dimensional Type 4 discrete cosine transform. -dct4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dct4 = dct4N [0] --- | 1-dimensional Type 1 discrete sine transform. -dst1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dst1 = dst1N [0] --- | 1-dimensional Type 2 discrete sine transform. -dst2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dst2 = dst2N [0] --- | 1-dimensional Type 3 discrete sine transform. -dst3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dst3 = dst3N [0] --- | 1-dimensional Type 4 discrete sine transform. -dst4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r -dst4 = dst4N [0] - --- Check if a flag is set. -infix 7 `has` -has :: Flag -> Flag -> Bool -a `has` b = a .&. b == b - --- | Try to transform a CArray with only one memory allocation (for the result). --- If we can find a way to prove that FFTW already has a sufficiently good plan --- for this transform size and the input will not be overwritten, then we could --- call have a version of this that does not require 'estimate'. Since this is --- not currently the case, we require 'estimate' to be set. Note that we do not --- check for the 'preserveInput' flag here. This is because the default is to --- preserve input for all but the C->R and HC->R transforms. Therefore, this --- function must not be called for those transforms, unless 'preserveInput' is --- set. -{-# NOINLINE transformCArray #-} -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`) [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 - --- | Transform a CArray with two memory allocations. This is entirely safe with --- all transforms, but it must allocate a temporary array to do the planning in. -{-# NOINLINE transformCArray' #-} -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 - --- | All the logic for determining shape of resulting array, and how to do the transform. -dftShape :: (Ix i, Shapable i, Storable e) - => DFT -> [Int] -> CArray i e -> ((i,i),TSpec) -dftShape t tdims a = assert valid (oBounds,tspec) - where shp = shape a - rnk = rank a - 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 -- Physical shape of the output array - where f = case t of - RC -> (\n -> n `div` 2 + 1) - CR -> (\n -> (n - 1) * 2) - CRO -> (\n -> (n - 1) * 2 + 1) - _ -> id - lShape = adjust f ldim shp -- Logical shape of the output array - where f = case t of - CR -> (\n -> (n - 1) * 2) - CRO -> (\n -> (n - 1) * 2 + 1) - _ -> id - oBounds = sBounds oShape - oStrides = shapeToStride oShape - ldim = last tdims - --- | A simple helper. -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 - --- | A generally useful list utility -adjust :: (a -> a) -> Int -> [a] -> [a] -adjust f i = uncurry (++) . second (\(x:xs) -> f x : xs) . splitAt i - --- | Complex to Complex DFT, un-normalized. -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 - --- | Real to Complex DFT. -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 - --- | Complex to Real DFT. The first argument determines whether the last --- transformed dimension is logically odd or even. 'True' implies the dimension --- is odd. -dftCRG_ :: (FFTWReal r, Ix i, Shapable i) => Bool -> Flag -> [Int] -> CArray i (Complex r) -> CArray i r -dftCRG_ isOdd 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 isOdd then CRO else CR) tdims ain - tCArr = if length tdims == 1 && f `has` preserveInput - -- A multi-dimensional C->R transform destroys its input. - -- Also, a one-dimensional transform is faster if it can - -- destroy input. - then transformCArray - else transformCArray' - --- | Complex to Real DFT where last transformed dimension is logically even. -dftCRGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r -dftCRGU = dftCRG_ False - --- | Complex to Real DFT where last transformed dimension is logically odd. -dftCROGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r -dftCROGU = dftCRG_ True - --- | Real to Real transforms. -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 - --- | Queries the FFTW cache. The 'String' can be written to a file so the --- wisdom can be reused on a subsequent run. -exportWisdomString :: IO String -exportWisdomString = do - pc <- c_export_wisdom_string - peekCString pc `finally` c_free pc - --- | Add wisdom to the FFTW cache. Returns 'True' if it is successful. -importWisdomString :: String -> IO Bool -importWisdomString str = - (==1) <$> withCString str c_import_wisdom_string - --- | Tries to import wisdom from a global source, typically @/etc/fftw/wisdom@. --- Returns 'True' if it was successful. -importWisdomSystem :: IO Bool -importWisdomSystem = (==1) <$> c_import_wisdom_system - --- We use "safe" calls for anything which could take a while so that it won't block --- other Haskell threads. - --- | Plan a complex to complex transform using the guru interface. -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 - --- | Plan a real to complex transform using the guru interface. -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 - --- | Plan a complex to real transform using the guru interface. -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 - --- | Plan a real to real transform using the guru interface. -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 - --- | Simple plan execution -foreign import ccall safe "fftw3.h fftw_execute" c_execute - :: Plan -> IO () - --- Execute a plan on different memory than the plan was created for. --- Alignment /must/ be the same. If we parallelize a transform of --- multi-dimensional data by making separate calls within an un-transformed --- dimension, it is possible that the alignment constraint would not be --- fulfilled. However, this only poses a problem for real transforms with odd --- transform dimension. -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 - --- | Frees memory allocated by 'fftw_malloc'. Currently, we only need this to --- free the wisdom string. -foreign import ccall unsafe "fftw3.h fftw_free" c_free :: Ptr a -> IO () diff --git a/fft.cabal b/fft.cabal --- a/fft.cabal +++ b/fft.cabal @@ -1,5 +1,5 @@ name: fft -version: 0.1.7.2 +version: 0.1.7.3 synopsis: Bindings to the FFTW library. description: Bindings to the FFTW library. @@ -14,17 +14,13 @@ author: Jed Brown maintainer: Jed Brown , Henning Thielemann build-type: Simple -cabal-version: >= 1.6 +cabal-version: >= 1.14 flag splitBase flag base4 -flag buildTests - description: build test suite executable - default: False - source-repository this - tag: 0.1.7.2 + tag: 0.1.7.3 type: darcs location: http://code.haskell.org/fft/ @@ -33,33 +29,36 @@ location: http://code.haskell.org/fft/ library + build-depends: + carray >=0.1.5 && <0.2, + storable-complex >=0.2.1 && <0.3, + ix-shapable >=0.1 && <0.2 if flag(splitBase) - build-depends: base >= 3, array, carray, storable-complex, ix-shapable + build-depends: base >= 3, array else - build-depends: base < 3, carray, storable-complex, ix-shapable + build-depends: base < 3 if flag(base4) build-depends: base >= 4 && < 5, syb >= 0.1 else build-depends: base < 4 + hs-source-dirs: src exposed-modules: Math.FFT Math.FFT.Base pkgconfig-depends: fftw3 >=3.3 && <4 - extensions: ForeignFunctionInterface ghc-options: -Wall + default-language: Haskell98 -executable test-fft - if flag(buildTests) - build-depends: - QuickCheck >=2.5 && <2.6 - else - buildable: False +test-suite test-fft + type: exitcode-stdio-1.0 + build-depends: + fft, + carray, + storable-complex, + QuickCheck >=2.5 && <2.6, + base ghc-options: -Wall main-is: tests.hs - hs-source-dirs: ., tests - other-modules: Math.FFT - Math.FFT.Base - pkgconfig-depends: - fftw3 >=3.3 && <4 - extensions: ForeignFunctionInterface + hs-source-dirs: tests + default-language: Haskell98 diff --git a/src/Math/FFT.hs b/src/Math/FFT.hs new file mode 100644 --- /dev/null +++ b/src/Math/FFT.hs @@ -0,0 +1,170 @@ +-- | +-- Module : Math.FFT +-- Copyright : (c) 2008 Jed Brown +-- License : BSD-style +-- +-- Maintainer : jed@59A2.org +-- Stability : experimental +-- Portability : non-portable +-- +-- This module exposes an interface to FFTW, the Fastest Fourier Transform in +-- the West. +-- +-- These bindings present several levels of interface. All the higher level +-- functions ('dft', 'idft', 'dftN', ...) are easily derived from the general +-- functions ('dftG', 'dftRCG', ...). Only the general functions let you +-- specify planner flags. The higher levels all set 'estimate' so you should +-- not have to wait through time consuming planning (see below for more). +-- +-- The simplest interface is the one-dimensional transforms. If you supply a +-- multi-dimensional array, these will only transform the first dimension. +-- These functions only take one argument, the array to be transformed. +-- +-- At the next level, we have multi-dimensional transforms where you specify +-- which dimensions to transform in and the array to transform. For instance +-- +-- > b = dftRCN [0,2] a +-- +-- is the real to complex transform in dimensions 0 and 2 of the array @a@ which +-- must be at least rank 3. The array @b@ will be complex valued with the same +-- extent as @a@ in every dimension except @2@. If @a@ had extent @n@ in +-- dimension @2@ then the @b@ will have extent @a `div` 2 + 1@ which consists of +-- all non-negative frequency components in this dimension (the negative +-- frequencies are conjugate to the positive frequencies because of symmetry +-- since @a@ is real valued). +-- +-- The real to real transforms allow different transform kinds in each +-- transformed dimension. For example, +-- +-- > b = dftRRN [(0,DHT), (1,REDFT10), (2,RODFT11)] a +-- +-- is a Discrete Hartley Transform in dimension 0, a discrete cosine transform +-- (DCT-2) in dimension 1, and distrete sine transform (DST-4) in dimension 2 +-- where the array @a@ must have rank at least 3. +-- +-- The general interface is similar to the multi-dimensional interface, takes as +-- its first argument, a bitwise '.|.' of planning 'Flag's. (In the complex +-- version, the sign of the transform is first.) For example, +-- +-- > b = dftG DFTBackward (patient .|. destroy_input) [1,2] a +-- +-- is an inverse DFT in dimensions 1 and 2 of the complex array @a@ which has +-- rank at least 3. It will use the patient planner to generate a (near) +-- optimal transform. If you compute the same type of transform again, it +-- should be very fast since the plan is cached. +-- +-- Inverse transforms are typically normalized. The un-normalized inverse +-- transforms are 'dftGU', 'dftCRGU' and 'dftCROGU'. For example +-- +-- > b = dftCROGU measure [0,1] a +-- +-- is an un-normalized inverse DFT in dimensions 0 and 1 of the complex array +-- @a@ (representing the non-negative frequencies, where the negative +-- frequencies are conjugate) which has rank at least 2. Here, dimension 1 is +-- logically odd so if @a@ has extent @n@ in dimension 1, then @b@ will have +-- extent @(n - 1) * 2 + 1@ in dimension 1. It is more common that the logical +-- dimension is even, in which case we would use 'dftCRGU' in which case @b@ +-- would have extent @(n - 1) * 2@ in dimension @1@. +-- +-- +-- The FFTW library separates transforms into two steps. First you compute a +-- plan for a given transform, then you execute it. Often the planning stage is +-- quite time-consuming, but subsequent transforms of the same size and type +-- will be extremely fast. The planning phase actually computes transforms, so +-- it overwrites its input array. For many C codes, it is reasonable to re-use +-- the same arrays to compute a given transform on different data. This is not +-- a very useful paradigm from Haskell. Fortunately, FFTW caches its plans so +-- if try to generate a new plan for a transform size which has already been +-- planned, the planner will return immediately. Unfortunately, it is not +-- possible to consult the cache, so if a plan is cached, we may use more memory +-- than is strictly necessary since we must allocate a work array which we +-- expect to be overwritten during planning. FFTW can export its cached plans +-- to a string. This is known as wisdom. For high performance work, it is a +-- good idea to compute plans of the sizes you are interested in, using +-- aggressive options (i.e. 'patient'), use 'exportWisdomString' to get a string +-- representing these plans, and write this to a file. Then for production +-- runs, you can read this in, then add it to the cache with +-- 'importWisdomString'. Now you can use the 'estimate' planner so the Haskell +-- bindings know that FFTW will not overwrite the input array, and you will +-- still get a high quality transform (because it has wisdom). + +module Math.FFT ( + -- * Data types + Sign(..), + Kind(..), + -- * Planner flags + -- ** Algorithm restriction flags + destroyInput, + preserveInput, + -- ** Planning rigor flags + estimate, + measure, + patient, + exhaustive, + + -- * DFT of complex data + -- ** DFT in first dimension only + dft, + idft, + -- ** Multi-dimensional transforms + dftN, + idftN, + -- ** General transform + dftG, + -- ** Un-normalized general transform + dftGU, + + -- * DFT of real data + -- ** DFT in first dimension only + dftRC, + dftCR, + dftCRO, + -- ** Multi-dimensional transforms + dftRCN, + dftCRN, + dftCRON, + -- ** General transform + dftRCG, + dftCRG, + dftCROG, + -- ** Un-normalized general transform + dftCRGU, + dftCROGU, + + -- * Real to real transforms (all un-normalized) + -- ** Transforms in first dimension only + dftRH, + dftHR, + dht, + dct1, + dct2, + dct3, + dct4, + dst1, + dst2, + dst3, + dst4, + -- ** Multi-dimensional transforms with the same transform type in each dimension + dftRHN, + dftHRN, + dhtN, + dct1N, + dct2N, + dct3N, + dct4N, + dst1N, + dst2N, + dst3N, + dst4N, + -- ** Multi-dimensional transforms with possibly different transforms in each dimension + dftRRN, + -- ** General transforms + dftRRG, + + -- * Wisdom + importWisdomString, + importWisdomSystem, + exportWisdomString, +) where + +import Math.FFT.Base diff --git a/src/Math/FFT/Base.hsc b/src/Math/FFT/Base.hsc new file mode 100644 --- /dev/null +++ b/src/Math/FFT/Base.hsc @@ -0,0 +1,595 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE ForeignFunctionInterface #-} +module Math.FFT.Base where + +import Control.Applicative +import Control.Arrow +import Control.Exception +import Control.Concurrent +import Control.Monad +import Data.Array.CArray + (CArray, withCArray, unsafeForeignPtrToCArray, + Ix, Shapable, shape, rank, size, rangeSize) +import Data.Array.CArray.Base (mallocForeignPtrArrayAligned, mapCArrayInPlace) +import Data.Ix.Shapable (shapeToStride, sBounds) +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) + +#include + +-- | Our API is polymorphic over the real data type. FFTW, at least in +-- principle, supports single precision 'Float', double precision 'Double' and +-- long double 'CLDouble' (presumable?). +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 + +-- | Using this instance requires linking with @-lfftw3@. +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 + +-- | This lock must be taken during /planning/ of any transform. The FFTW +-- library is not thread-safe in the planning phase. Thankfully, the lock is +-- not needed during the execute phase. +lock :: MVar () +lock = unsafePerformIO $ newMVar () +{-# NOINLINE lock #-} + +withLock :: IO a -> IO a +withLock = withMVar lock . const + +-- | A plan is an opaque foreign object. +type Plan = Ptr FFTWPlan + +type FFTWPlan = () + +-- | The 'Flag' type is used to influence the kind of plans which are created. +-- To specify multiple flags, use a bitwise '.|.'. +newtype Flag = Flag { unFlag :: FFTWFlag } + deriving (Eq, Show, Num, Bits) + +type FFTWFlag = CUInt + +#{enum FFTWFlag, + , c_measure = FFTW_MEASURE + , c_destroy_input = FFTW_DESTROY_INPUT + , c_unaligned = FFTW_UNALIGNED + , c_conserve_memory = FFTW_CONSERVE_MEMORY + , c_exhaustive = FFTW_EXHAUSTIVE + , c_preserve_input = FFTW_PRESERVE_INPUT + , c_patient = FFTW_PATIENT + , c_estimate = FFTW_ESTIMATE } + +-- | Default flag. For most transforms, this is equivalent to setting 'measure' +-- and 'preserveInput'. The exceptions are complex to real and half-complex to +-- real transforms. +nullFlag :: Flag +nullFlag = Flag 0 + +-- +-- Algorithm restriction flags +-- + +-- | Allows FFTW to overwrite the input array with arbitrary data; this can +-- sometimes allow more efficient algorithms to be employed. +-- +-- Setting this flag implies that two memory allocations will be done, one for +-- work space, and one for the result. When 'estimate' is not set, we will be +-- doing two memory allocations anyway, so we set this flag as well (since we +-- don't retain the work array anyway). +destroyInput :: Flag +destroyInput = Flag c_destroy_input + +-- | 'preserveInput' specifies that an out-of-place transform must not change +-- its input array. This is ordinarily the default, except for complex to real +-- transforms for which 'destroyInput' is the default. In the latter cases, +-- passing 'preserveInput' will attempt to use algorithms that do not destroy +-- the input, at the expense of worse performance; for multi-dimensional complex +-- to real transforms, however, no input-preserving algorithms are implemented +-- so the Haskell bindings will set 'destroyInput' and do a transform with two +-- memory allocations. +preserveInput :: Flag +preserveInput = Flag c_preserve_input + +-- | Instruct FFTW not to generate a plan which uses SIMD instructions, even if +-- the memory you are planning with is aligned. This should only be needed if +-- you are using the guru interface and want to reuse a plan with memory that +-- may be unaligned (i.e. you constructed the 'CArray' with +-- 'unsafeForeignPtrToCArray'). +unaligned :: Flag +unaligned = Flag c_unaligned + +-- | The header claims that this flag is documented, but in reality, it is not. +-- I don't know what it does and it is here only for completeness. +conserveMemory :: Flag +conserveMemory = Flag c_conserve_memory + +-- +-- Planning rigor flags +-- + +-- | 'estimate' specifies that, instead of actual measurements of different +-- algorithms, a simple heuristic is used to pick a (probably sub-optimal) plan +-- quickly. With this flag, the input/output arrays are not overwritten during +-- planning. +-- +-- This is the only planner flag for which a single memory allocation is possible. +estimate :: Flag +estimate = Flag c_estimate + +-- | 'measure' tells FFTW to find an optimized plan by actually computing +-- several FFTs and measuring their execution time. Depending on your machine, +-- this can take some time (often a few seconds). 'measure' is the default +-- planning option. +measure :: Flag +measure = Flag c_measure + +-- | 'patient' is like 'measure', but considers a wider range of algorithms and +-- often produces a "more optimal" plan (especially for large transforms), but +-- at the expense of several times longer planning time (especially for large +-- transforms). +patient :: Flag +patient = Flag c_patient + +-- | 'exhaustive' is like 'patient' but considers an even wider range of +-- algorithms, including many that we think are unlikely to be fast, to +-- produce the most optimal plan but with a substantially increased planning +-- time. +exhaustive :: Flag +exhaustive = Flag c_exhaustive + +-- | Determine which direction of DFT to execute. +data Sign = DFTForward | DFTBackward + deriving (Eq,Show) + +type FFTWSign = CInt + +#{enum FFTWSign, + , c_forward = FFTW_FORWARD + , c_backward = FFTW_BACKWARD } + +unSign :: Sign -> FFTWSign +unSign DFTForward = c_forward +unSign DFTBackward = c_backward + +-- | Real to Real transform kinds. +data Kind = R2HC | HC2R -- half-complex transforms + | DHT -- discrete Hartley transformm + | REDFT00 | REDFT10 | REDFT01 | REDFT11 -- discrete cosine transforms + | RODFT00 | RODFT01 | RODFT10 | RODFT11 -- discrete sine transforms + 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 + +#{enum FFTWKind, + , c_r2hc = FFTW_R2HC + , c_hc2r = FFTW_HC2R + , c_dht = FFTW_DHT + , c_redft00 = FFTW_REDFT00 + , c_redft10 = FFTW_REDFT10 + , c_redft01 = FFTW_REDFT01 + , c_redft11 = FFTW_REDFT11 + , c_rodft00 = FFTW_RODFT00 + , c_rodft10 = FFTW_RODFT10 + , c_rodft01 = FFTW_RODFT01 + , c_rodft11 = FFTW_RODFT11 } + +-- | Corresponds to the @fftw_iodim@ structure. It completely describes the +-- layout of each dimension, before and after the transform. +data IODim = IODim { nIODim :: Int -- ^ Logical size of dimension + , isIODim :: Int -- ^ Stride along dimension in input array + , osIODim :: Int -- ^ Stride along dimension in output array + } + deriving (Eq, Show, Data, Typeable) + +instance Storable IODim where + sizeOf _ = #{size fftw_iodim} + alignment _ = alignment (undefined :: CInt) + peek p = do + n' <- #{peek fftw_iodim, n} p + is' <- #{peek fftw_iodim, is} p + os' <- #{peek fftw_iodim, os} p + return (IODim n' is' os') + poke p (IODim n' is' os') = do + #{poke fftw_iodim, n} p n' + #{poke fftw_iodim, is} p is' + #{poke fftw_iodim, os} p os' + +-- | Tuple of transform dimensions and non-transform dimensions of the array. +type TSpec = ([IODim],[IODim]) + +-- | Types of transforms. Used to control 'dftShape'. +data DFT = CC | RC | CR | CRO | RR + deriving (Eq, Show) + +-- | Verify that a plan is valid. Thows an exception if not. +check :: Plan -> IO () +check p = when (p == nullPtr) . ioError $ userError "invalid plan" + +-- | Confirm that the plan is valid, then execute the transform. +execute :: Plan -> IO () +execute p = check p >> c_execute p + +-- | In-place normalization outside of IO. You must be able to prove that no +-- reference to the original can be retained. +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) + +-- | Normalized general complex DFT +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) + +-- | Normalized general complex to real DFT where the last transformed dimension +-- is logically even. +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) + +-- | Normalized general complex to real DFT where the last transformed dimension +-- is logicall odd. +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) + +-- | Multi-dimensional forward DFT. +dftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r) +dftN = dftG DFTForward estimate +-- | Multi-dimensional inverse DFT. +idftN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i (Complex r) +idftN = dftG DFTBackward estimate +-- | Multi-dimensional forward DFT of real data. +dftRCN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i (Complex r) +dftRCN = dftRCG estimate +-- | Multi-dimensional inverse DFT of Hermitian-symmetric data (where only the +-- non-negative frequencies are given). +dftCRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r +dftCRN = dftCRG estimate +-- | Multi-dimensional inverse DFT of Hermitian-symmetric data (where only the +-- non-negative frequencies are given) and the last transformed dimension is +-- logically odd. +dftCRON :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i (Complex r) -> CArray i r +dftCRON = dftCROG estimate + +fzr :: b -> [a] -> [(a,b)] +fzr = flip zip . repeat +drr :: (FFTWReal r, Ix i, Shapable i) => Kind -> [Int] -> CArray i r -> CArray i r +drr = (dftRRN .) . fzr + +-- | Multi-dimensional real to real transform. The result is not normalized. +dftRRN :: (FFTWReal r, Ix i, Shapable i) => [(Int,Kind)] -> CArray i r -> CArray i r +dftRRN = dftRRG estimate + +-- +-- The following do the same type of transform in each dimension specified. +-- +-- | Multi-dimensional real to half-complex transform. The result is not normalized. +dftRHN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dftRHN = drr R2HC +-- | Multi-dimensional half-complex to real transform. The result is not normalized. +dftHRN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dftHRN = drr HC2R +-- | Multi-dimensional Discrete Hartley Transform. The result is not normalized. +dhtN :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dhtN = drr DHT +-- | Multi-dimensional Type 1 discrete cosine transform. +dct1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dct1N = drr REDFT00 +-- | Multi-dimensional Type 2 discrete cosine transform. This is commonly known +-- as /the/ DCT. +dct2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dct2N = drr REDFT10 +-- | Multi-dimensional Type 3 discrete cosine transform. This is commonly known +-- as /the/ inverse DCT. The result is not normalized. +dct3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dct3N = drr REDFT01 +-- | Multi-dimensional Type 4 discrete cosine transform. +dct4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dct4N = drr REDFT11 +-- | Multi-dimensional Type 1 discrete sine transform. +dst1N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dst1N = drr RODFT00 +-- | Multi-dimensional Type 2 discrete sine transform. +dst2N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dst2N = drr RODFT10 +-- | Multi-dimensional Type 3 discrete sine transform. +dst3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dst3N = drr RODFT01 +-- | Multi-dimensional Type 4 discrete sine transform. +dst4N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dst4N = drr RODFT11 + +-- +-- Transform in the first dimension only. +-- + +-- | 1-dimensional complex DFT. +dft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r) +dft = dftN [0] +-- | 1-dimensional complex inverse DFT. Inverse of 'dft'. +idft :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i (Complex r) +idft = idftN [0] +-- | 1-dimensional real to complex DFT. +dftRC :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i (Complex r) +dftRC = dftRCN [0] +-- | 1-dimensional complex to real DFT with logically even dimension. Inverse of 'dftRC'. +dftCR :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r +dftCR = dftCRN [0] +-- | 1-dimensional complex to real DFT with logically odd dimension. Inverse of 'dftRC'. +dftCRO :: (FFTWReal r, Ix i, Shapable i) => CArray i (Complex r) -> CArray i r +dftCRO = dftCRON [0] +-- | 1-dimensional real to half-complex DFT. +dftRH :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dftRH = dftRHN [0] +-- | 1-dimensional half-complex to real DFT. Inverse of 'dftRH' after normalization. +dftHR :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dftHR = dftHRN [0] +-- | 1-dimensional Discrete Hartley Transform. Self-inverse after normalization. +dht :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dht = dhtN [0] +-- | 1-dimensional Type 1 discrete cosine transform. +dct1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dct1 = dct1N [0] +-- | 1-dimensional Type 2 discrete cosine transform. This is commonly known as /the/ DCT. +dct2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dct2 = dct2N [0] +-- | 1-dimensional Type 3 discrete cosine transform. This is commonly known as /the/ inverse DCT. +dct3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dct3 = dct3N [0] +-- | 1-dimensional Type 4 discrete cosine transform. +dct4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dct4 = dct4N [0] +-- | 1-dimensional Type 1 discrete sine transform. +dst1 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dst1 = dst1N [0] +-- | 1-dimensional Type 2 discrete sine transform. +dst2 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dst2 = dst2N [0] +-- | 1-dimensional Type 3 discrete sine transform. +dst3 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dst3 = dst3N [0] +-- | 1-dimensional Type 4 discrete sine transform. +dst4 :: (FFTWReal r, Ix i, Shapable i) => CArray i r -> CArray i r +dst4 = dst4N [0] + +-- Check if a flag is set. +infix 7 `has` +has :: Flag -> Flag -> Bool +a `has` b = a .&. b == b + +-- | Try to transform a CArray with only one memory allocation (for the result). +-- If we can find a way to prove that FFTW already has a sufficiently good plan +-- for this transform size and the input will not be overwritten, then we could +-- call have a version of this that does not require 'estimate'. Since this is +-- not currently the case, we require 'estimate' to be set. Note that we do not +-- check for the 'preserveInput' flag here. This is because the default is to +-- preserve input for all but the C->R and HC->R transforms. Therefore, this +-- function must not be called for those transforms, unless 'preserveInput' is +-- set. +{-# NOINLINE transformCArray #-} +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`) [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 + +-- | Transform a CArray with two memory allocations. This is entirely safe with +-- all transforms, but it must allocate a temporary array to do the planning in. +{-# NOINLINE transformCArray' #-} +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 + +-- | All the logic for determining shape of resulting array, and how to do the transform. +dftShape :: (Ix i, Shapable i, Storable e) + => DFT -> [Int] -> CArray i e -> ((i,i),TSpec) +dftShape t tdims a = assert valid (oBounds,tspec) + where shp = shape a + rnk = rank a + 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 -- Physical shape of the output array + where f = case t of + RC -> (\n -> n `div` 2 + 1) + CR -> (\n -> (n - 1) * 2) + CRO -> (\n -> (n - 1) * 2 + 1) + _ -> id + lShape = adjust f ldim shp -- Logical shape of the output array + where f = case t of + CR -> (\n -> (n - 1) * 2) + CRO -> (\n -> (n - 1) * 2 + 1) + _ -> id + oBounds = sBounds oShape + oStrides = shapeToStride oShape + ldim = last tdims + +-- | A simple helper. +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 + +-- | A generally useful list utility +adjust :: (a -> a) -> Int -> [a] -> [a] +adjust f i = uncurry (++) . second (\(x:xs) -> f x : xs) . splitAt i + +-- | Complex to Complex DFT, un-normalized. +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 + +-- | Real to Complex DFT. +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 + +-- | Complex to Real DFT. The first argument determines whether the last +-- transformed dimension is logically odd or even. 'True' implies the dimension +-- is odd. +dftCRG_ :: (FFTWReal r, Ix i, Shapable i) => Bool -> Flag -> [Int] -> CArray i (Complex r) -> CArray i r +dftCRG_ isOdd 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 isOdd then CRO else CR) tdims ain + tCArr = if length tdims == 1 && f `has` preserveInput + -- A multi-dimensional C->R transform destroys its input. + -- Also, a one-dimensional transform is faster if it can + -- destroy input. + then transformCArray + else transformCArray' + +-- | Complex to Real DFT where last transformed dimension is logically even. +dftCRGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r +dftCRGU = dftCRG_ False + +-- | Complex to Real DFT where last transformed dimension is logically odd. +dftCROGU :: (FFTWReal r, Ix i, Shapable i) => Flag -> [Int] -> CArray i (Complex r) -> CArray i r +dftCROGU = dftCRG_ True + +-- | Real to Real transforms. +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 + +-- | Queries the FFTW cache. The 'String' can be written to a file so the +-- wisdom can be reused on a subsequent run. +exportWisdomString :: IO String +exportWisdomString = do + pc <- c_export_wisdom_string + peekCString pc `finally` c_free pc + +-- | Add wisdom to the FFTW cache. Returns 'True' if it is successful. +importWisdomString :: String -> IO Bool +importWisdomString str = + (==1) <$> withCString str c_import_wisdom_string + +-- | Tries to import wisdom from a global source, typically @/etc/fftw/wisdom@. +-- Returns 'True' if it was successful. +importWisdomSystem :: IO Bool +importWisdomSystem = (==1) <$> c_import_wisdom_system + +-- We use "safe" calls for anything which could take a while so that it won't block +-- other Haskell threads. + +-- | Plan a complex to complex transform using the guru interface. +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 + +-- | Plan a real to complex transform using the guru interface. +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 + +-- | Plan a complex to real transform using the guru interface. +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 + +-- | Plan a real to real transform using the guru interface. +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 + +-- | Simple plan execution +foreign import ccall safe "fftw3.h fftw_execute" c_execute + :: Plan -> IO () + +-- Execute a plan on different memory than the plan was created for. +-- Alignment /must/ be the same. If we parallelize a transform of +-- multi-dimensional data by making separate calls within an un-transformed +-- dimension, it is possible that the alignment constraint would not be +-- fulfilled. However, this only poses a problem for real transforms with odd +-- transform dimension. +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 + +-- | Frees memory allocated by 'fftw_malloc'. Currently, we only need this to +-- free the wisdom string. +foreign import ccall unsafe "fftw3.h fftw_free" c_free :: Ptr a -> IO ()