diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,31 @@ +Copyright Jed Brown 2008 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Jed Brown nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + diff --git a/Math/FFT.hs b/Math/FFT.hs new file mode 100644 --- /dev/null +++ b/Math/FFT.hs @@ -0,0 +1,166 @@ +-- | +-- 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 +) where + +import Math.FFT.Base +import Data.Array.CArray diff --git a/Math/FFT/Base.hsc b/Math/FFT/Base.hsc new file mode 100644 --- /dev/null +++ b/Math/FFT/Base.hsc @@ -0,0 +1,583 @@ +{-# 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 ( 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 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 = 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 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 { n :: Int, is :: Int, os :: Int } + deriving (Eq, Show) + +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 = flip zip . repeat +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 REDFT01 +-- | 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 REDFT10 +-- | 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 RODFT01 +-- | Multi-dimensional Type 3 discrete sine transform. +dst3N :: (FFTWReal r, Ix i, Shapable i) => [Int] -> CArray i r -> CArray i r +dst3N = drr RODFT10 +-- | 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` +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`) [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 + +-- | 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, 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 -- Physical shape of the output array + 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 -- Logical shape of the output array + 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 + +-- | 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_ 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 + -- 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/README b/README new file mode 100644 --- /dev/null +++ b/README @@ -0,0 +1,25 @@ +This package provides bindings to the FFTW library. + +You will need to install FFTW version 3, including development files before this +package. Consult your package manager or visit http://fftw.org to install FFTW. +In addition, the Haskell package carray is required. + + runhaskell Setup.lhs configure + runhaskell Setup.lhs build + runhaskell Setup.lhs haddock (optional) + runhaskell Setup.lhs install + +Then run the tests: + + runhaskell tests/tests.hs + +We use the CArray package for multi-dimensional arrays. It allocates pinned +memory on the GC'd heap, which is 16-byte aligned by default, allowing SIMD +instructions. If you get a CArray from a foreign source using +unsafeForeignPtrToCArray (an O(1) operation) then you must be sure that the +memory is aligned if you expect SIMD code to be used by FFTW. + +A note regarding licensing: FFTW is generally distributed under the GPL, +although a different license can be purchased. Therefore, the fact that these +bindings are BSD licensed does not mean you can link against a GPL'd copy of +FFTW without complying with the GPL. diff --git a/Setup.lhs b/Setup.lhs new file mode 100644 --- /dev/null +++ b/Setup.lhs @@ -0,0 +1,3 @@ +#!/usr/bin/env runhaskell +> import Distribution.Simple +> main = defaultMain diff --git a/fft.cabal b/fft.cabal new file mode 100644 --- /dev/null +++ b/fft.cabal @@ -0,0 +1,22 @@ +name: fft +version: 0.1.0.0 +synopsis: Bindings to the FFTW library. +description: + Bindings to the FFTW library. + . + Provides high performance discrete fourier transforms in + arbitrary dimensions. Include transforms of complex data, + real data, and real to real transforms. + . +category: Math +license: BSD3 +license-file: LICENSE +author: Jed Brown +maintainer: +build-Depends: base, array, carray +extra-libraries: fftw3 +extensions: ForeignFunctionInterface +exposed-modules: Math.FFT + Math.FFT.Base +ghc-options: +build-type: Simple diff --git a/tests/tests.hs b/tests/tests.hs new file mode 100644 --- /dev/null +++ b/tests/tests.hs @@ -0,0 +1,119 @@ +{-# LANGUAGE FlexibleInstances, FlexibleContexts #-} +import Test.QuickCheck +import Data.Array.CArray +import Data.Complex +import Math.FFT +import Foreign.Storable +import Text.Printf +import System.Environment (getArgs) +import System.IO +import System.Random + +instance Arbitrary (Complex Double) where + arbitrary = do + r <- arbitrary + i <- arbitrary + return $ r :+ i + coarbitrary = error "no coarbitrary for Complex" + +instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray Int e) where + arbitrary = do + u <- choose (1,100) + es <- vector (u+1) + return $ listArray (0,u) es + coarbitrary = error "no coarbitrary for CArray" + +instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray (Int,Int) e) where + arbitrary = do + u0 <- choose (1,30) + u1 <- choose (1,30) + es <- vector ((u0 + 1) * (u1 + 1)) + return $ listArray ((0,0),(u0,u1)) es + coarbitrary = error "no coarbitrary for CArray" + +instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray (Int,Int,Int) e) where + arbitrary = do + u0 <- choose (1,20) + u1 <- choose (1,20) + u2 <- choose (1,20) + es <- vector ((u0 + 1) * (u1 + 1) * (u2 + 1)) + return $ listArray ((0,0,0),(u0,u1,u2)) es + coarbitrary = error "no coarbitrary for CArray" + + +-- about :: (Ix i, FFTWFloat e) => CArray i e -> CArray i e -> Bool +about x y = small $ normSup (liftArray2 (-) x y) / (1 + normSup (liftArray2 (+) x y)) + where small a = a < 1e-15 + +partAbout a b = about a (slice ba ba b) + where ba = bounds a + +aboutIdem f x = f x `about` x + +prop_dft = aboutIdem $ idft . dft +prop_dftRC a = aboutIdem ((if odd (shape a !! 0) then dftCRO else dftCR) . dftRC) a +prop_dftRC_dft a = partAbout (dftRC a) (dft . amap (:+0) $ a) +prop_dht_idem a = aboutIdem (amap (/ fromIntegral (shape a !! 0)) . dht . dht) a + + +prop_dft2 = aboutIdem $ idft . dft +prop_dft22 = aboutIdem $ idftN [0,1] . dftN [0,1] +prop_dft22' = aboutIdem $ idftN [1,0] . dftN [1,0] + +prop_dftRC2 a = aboutIdem ((if odd (shape a !! 0) then dftCRO else dftCR) . dftRC) a +prop_dftRC_dft2 a = partAbout (dftRC a) (dft . amap (:+0) $ a) +prop_dftRC_dft22 a = partAbout (dftRCN [0,1] a) (dftN [0,1] . amap (:+0) $ a) +prop_dht_idem2 a = aboutIdem (amap (/ fromIntegral (shape a !! 0)) . dht . dht) a + +prop_dft3 = aboutIdem $ idft . dft +prop_dft32 = aboutIdem $ idftN [0,1] . dftN [0,1] +prop_dft32' = aboutIdem $ idftN [1,0] . dftN [1,0] +prop_dft33 = aboutIdem $ idftN [0,1,2] . dftN [0,1,2] +prop_dft33' = aboutIdem $ idftN [0,2,1] . dftN [0,2,1] +prop_dft33'' = aboutIdem $ idftN [2,0,1] . dftN [2,0,1] + +c_tests :: [(String, CArray Int (Complex Double) -> Bool)] +c_tests = [ ("dft idem 1D" , prop_dft) + ] + +c_tests2 :: [(String, CArray (Int,Int) (Complex Double) -> Bool)] +c_tests2 = [ ("dft idem 2D" , prop_dft2) + , ("dft idem 2D/2" , prop_dft22) + , ("dft idem 2D/2'" , prop_dft22') + ] + +c_tests3 :: [(String, CArray (Int,Int,Int) (Complex Double) -> Bool)] +c_tests3 = [ ("dft idem 3D" , prop_dft3) + , ("dft idem 3D/2" , prop_dft32) + , ("dft idem 3D/2'" , prop_dft32') + , ("dft idem 3D/3" , prop_dft33) + , ("dft idem 3D/3'" , prop_dft33') + , ("dft idem 3D/3''" , prop_dft33'') + ] + +r_tests :: [(String, CArray Int Double -> Bool)] +r_tests = [ ("dftRC/CR idem 1D" , prop_dftRC) + , ("dftRC dft 1D" , prop_dftRC_dft) + , ("dht idem 1D" , prop_dht_idem) + ] + +r_tests2 :: [(String, CArray (Int,Int) Double -> Bool)] +r_tests2 = [ ("dftRC/CR idem 2D" , prop_dftRC2) + , ("dftRC dft 2D" , prop_dftRC_dft2) + , ("dftRC dft 2D/2" , prop_dftRC_dft22) + , ("dht idem 2D" , prop_dht_idem2) + ] + +main = do + x <- getArgs + let n = if null x then 20 else read . head $ x + conf = Config { configMaxTest = n + , configMaxFail = 1000 + , configSize = (+ 3) . (`div` 2) + , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s] + } + mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests + mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) r_tests + mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests2 + mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) r_tests2 + mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests3