{-# LINE 1 "Numeric/FFT/Vector/Base.hsc" #-}
-- | A basic interface between Vectors and the fftw library.
module Numeric.FFT.Vector.Base(
            -- * Transforms
            Transform(..),
            planOfType,
            PlanType(..),
            plan,
            run,
            -- * multi-demensional Transforms
            TransformND(..),
            planOfTypeND,
            planND,
            runND,
            -- * Plans
            Plan(),
            planInputSize,
            planOutputSize,
            execute,
            executeM,
            withPlanner,
            -- * Unsafe C stuff
            CFlags,
            CPlan,
            -- * Normalization helpers
            Scalable(..),
            modifyInput,
            modifyOutput,
            constMultOutput,
            multC,
            unsafeModify,
            ) where

import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MS
import Data.Vector.Generic as V hiding (forM_)
import Data.Vector.Generic.Mutable as M hiding (forM_, unsafeModify)
import Data.List as L
import Control.Concurrent.MVar
import Control.Monad.Primitive (RealWorld,PrimMonad(..), PrimBase,
            unsafePrimToPrim, unsafePrimToIO)
import Control.Monad(forM_)
import Foreign (Storable(..), Ptr, FunPtr,
                ForeignPtr, withForeignPtr, newForeignPtr)
import Foreign.C (CInt(..), CUInt, CSize(..))
import Data.Bits ( (.|.) )
import Data.Complex(Complex(..))
import Foreign.Storable.Complex()
import System.IO.Unsafe (unsafePerformIO)





---------------------
-- Creating FFTW plans

-- First, the Transform flags:
data PlanType = Estimate | Measure | Patient | Exhaustive
data Preservation = PreserveInput | DestroyInput

type CFlags = CUInt

-- | Marshal the Transform flags for use by fftw.
planInitFlags :: PlanType -> Preservation -> CFlags
planInitFlags :: PlanType -> Preservation -> CFlags
planInitFlags PlanType
pt Preservation
pr = CFlags
planTypeInt CFlags -> CFlags -> CFlags
forall a. Bits a => a -> a -> a
.|. CFlags
preservationInt
  where
    planTypeInt :: CFlags
planTypeInt = case PlanType
pt of
                    PlanType
Estimate -> CFlags
64
{-# LINE 69 "Numeric/FFT/Vector/Base.hsc" #-}
                    PlanType
Measure -> CFlags
0
{-# LINE 70 "Numeric/FFT/Vector/Base.hsc" #-}
                    PlanType
Patient -> CFlags
32
{-# LINE 71 "Numeric/FFT/Vector/Base.hsc" #-}
                    PlanType
Exhaustive -> CFlags
8
{-# LINE 72 "Numeric/FFT/Vector/Base.hsc" #-}
    preservationInt = case pr of
                    PreserveInput -> 16
{-# LINE 74 "Numeric/FFT/Vector/Base.hsc" #-}
                    DestroyInput -> 1
{-# LINE 75 "Numeric/FFT/Vector/Base.hsc" #-}

newtype CPlan = CPlan {CPlan -> ForeignPtr CPlan
unCPlan :: ForeignPtr CPlan}

withPlan :: CPlan -> (Ptr CPlan -> IO a) -> IO a
withPlan :: forall a. CPlan -> (Ptr CPlan -> IO a) -> IO a
withPlan = ForeignPtr CPlan -> (Ptr CPlan -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (ForeignPtr CPlan -> (Ptr CPlan -> IO a) -> IO a)
-> (CPlan -> ForeignPtr CPlan)
-> CPlan
-> (Ptr CPlan -> IO a)
-> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CPlan -> ForeignPtr CPlan
unCPlan

foreign import ccall unsafe fftw_execute :: Ptr CPlan -> IO ()
foreign import ccall "&" fftw_destroy_plan :: FunPtr (Ptr CPlan -> IO ())

newPlan :: Ptr CPlan -> IO CPlan
newPlan :: Ptr CPlan -> IO CPlan
newPlan = (ForeignPtr CPlan -> CPlan) -> IO (ForeignPtr CPlan) -> IO CPlan
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForeignPtr CPlan -> CPlan
CPlan (IO (ForeignPtr CPlan) -> IO CPlan)
-> (Ptr CPlan -> IO (ForeignPtr CPlan)) -> Ptr CPlan -> IO CPlan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FinalizerPtr CPlan -> Ptr CPlan -> IO (ForeignPtr CPlan)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr CPlan
fftw_destroy_plan

----------------------------------------
-- vector-fftw plans

-- | A 'Plan' can be used to run an @fftw@ algorithm for a specific input/output size.
data Plan a b = Plan {
                    forall a b. Plan a b -> MVector RealWorld a
planInput :: {-# UNPACK #-} !(VS.MVector RealWorld a),
                    forall a b. Plan a b -> MVector RealWorld b
planOutput :: {-# UNPACK #-} !(VS.MVector RealWorld b),
                    forall a b. Plan a b -> IO ()
planExecute :: IO ()
                }

-- | The (only) valid input size for this plan.
planInputSize :: Storable a => Plan a b -> Int
planInputSize :: forall a b. Storable a => Plan a b -> Int
planInputSize = MVector RealWorld a -> Int
forall a s. Storable a => MVector s a -> Int
MS.length (MVector RealWorld a -> Int)
-> (Plan a b -> MVector RealWorld a) -> Plan a b -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Plan a b -> MVector RealWorld a
forall a b. Plan a b -> MVector RealWorld a
planInput

-- | The (only) valid output size for this plan.
planOutputSize :: Storable b => Plan a b -> Int
planOutputSize :: forall b a. Storable b => Plan a b -> Int
planOutputSize = MVector RealWorld b -> Int
forall a s. Storable a => MVector s a -> Int
MS.length (MVector RealWorld b -> Int)
-> (Plan a b -> MVector RealWorld b) -> Plan a b -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Plan a b -> MVector RealWorld b
forall a b. Plan a b -> MVector RealWorld b
planOutput

-- | Run a plan on the given 'Vector'.
--
-- If @'planInputSize' p /= length v@, then calling
-- @execute p v@ will throw an exception.
execute :: (Vector v a, Vector v b, Storable a, Storable b)
            => Plan a b -> v a -> v b
execute :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Storable a, Storable b) =>
Plan a b -> v a -> v b
execute Plan{IO ()
MVector RealWorld a
MVector RealWorld b
planInput :: forall a b. Plan a b -> MVector RealWorld a
planOutput :: forall a b. Plan a b -> MVector RealWorld b
planExecute :: forall a b. Plan a b -> IO ()
planInput :: MVector RealWorld a
planOutput :: MVector RealWorld b
planExecute :: IO ()
..} = \v a
v -> -- fudge the arity to make sure it's always inlined
    if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v
        then [Char] -> v b
forall a. HasCallStack => [Char] -> a
error ([Char] -> v b) -> [Char] -> v b
forall a b. (a -> b) -> a -> b
$ [Char]
"execute: size mismatch; expected " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ [Char]
", got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ Int -> [Char]
forall a. Show a => a -> [Char]
show (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v)
        else IO (v b) -> v b
forall a. IO a -> a
unsafePerformIO (IO (v b) -> v b) -> IO (v b) -> v b
forall a b. (a -> b) -> a -> b
$ do
                        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
planInput Int
k
                                                (a -> IO ()) -> a -> IO ()
forall a b. (a -> b) -> a -> b
$ v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.unsafeIndex v a
v Int
k
                        IO ()
planExecute
                        v' <- Int -> IO (Mutable v (PrimState IO) b)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
unsafeNew Int
m
                        forM_ [0..m-1] $ \Int
k -> MVector (PrimState IO) b -> Int -> IO b
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
M.unsafeRead MVector RealWorld b
MVector (PrimState IO) b
planOutput Int
k
                                                IO b -> (b -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v (PrimState IO) b -> Int -> b -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.unsafeWrite Mutable v RealWorld b
Mutable v (PrimState IO) b
v' Int
k
                        V.unsafeFreeze v'
  where
    n :: Int
n = MVector RealWorld a -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld a
planInput
    m :: Int
m = MVector RealWorld b -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld b
planOutput
{-# INLINE execute #-}

-- TODO: decide whether this is actually unsafe.
-- | Run a plan on the given mutable vectors.  The same vector may be used for both
-- input and output.
--
-- If @'planInputSize' p \/= length vIn@ or @'planOutputSize' p \/= length vOut@,
-- then calling @unsafeExecuteM p vIn vOut@ will throw an exception.
executeM :: forall m v a b .
        (PrimBase m, MVector v a, MVector v b, Storable a, Storable b)
            => Plan a b -- ^ The plan to run.
            -> v (PrimState m) a  -- ^ The input vector.
                    -> v (PrimState m) b -- ^ The output vector.
                    -> m ()
executeM :: forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimBase m, MVector v a, MVector v b, Storable a, Storable b) =>
Plan a b -> v (PrimState m) a -> v (PrimState m) b -> m ()
executeM Plan{IO ()
MVector RealWorld a
MVector RealWorld b
planInput :: forall a b. Plan a b -> MVector RealWorld a
planOutput :: forall a b. Plan a b -> MVector RealWorld b
planExecute :: forall a b. Plan a b -> IO ()
planInput :: MVector RealWorld a
planOutput :: MVector RealWorld b
planExecute :: IO ()
..} = \v (PrimState m) a
vIn v (PrimState m) b
vOut ->
    if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vIn Bool -> Bool -> Bool
|| Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= v (PrimState m) b -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) b
vOut
        then [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"executeM: size mismatch; expected " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Int
n,Int
m)
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ [Char]
", got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
L.++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vIn, v (PrimState m) b -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) b
vOut)
        else IO () -> m ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> v (PrimState m) b -> IO ()
act v (PrimState m) a
vIn v (PrimState m) b
vOut
  where
    n :: Int
n = MVector RealWorld a -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld a
planInput
    m :: Int
m = MVector RealWorld b -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld b
planOutput

    act :: v (PrimState m) a -> v (PrimState m) b -> IO ()
    act :: v (PrimState m) a -> v (PrimState m) b -> IO ()
act v (PrimState m) a
vIn v (PrimState m) b
vOut = do
            [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> m a -> IO a
forall (m :: * -> *) a. PrimBase m => m a -> IO a
unsafePrimToIO (v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
M.unsafeRead v (PrimState m) a
vIn Int
k :: m a)
                                    IO a -> (a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
planInput Int
k
            IO () -> IO ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim IO ()
planExecute
            [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> MVector (PrimState IO) b -> Int -> IO b
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
M.unsafeRead MVector RealWorld b
MVector (PrimState IO) b
planOutput Int
k
                                    IO b -> (b -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m () -> IO ()
forall (m :: * -> *) a. PrimBase m => m a -> IO a
unsafePrimToIO (m () -> IO ()) -> (b -> m ()) -> b -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v (PrimState m) b -> Int -> b -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.unsafeWrite v (PrimState m) b
vOut Int
k
                                                            :: b -> m ())
{-# INLINE executeM #-}

------------------
-- Malloc/free of fftw array

foreign import ccall unsafe fftw_malloc :: CSize -> IO (Ptr a)
foreign import ccall "&" fftw_free :: FunPtr (Ptr a -> IO ())

newFFTVector :: forall a . Storable a => Int -> IO (MS.MVector RealWorld a)
newFFTVector :: forall a. Storable a => Int -> IO (MVector RealWorld a)
newFFTVector Int
n = do
    p <- CSize -> IO (Ptr a)
forall a. CSize -> IO (Ptr a)
fftw_malloc (CSize -> IO (Ptr a)) -> CSize -> IO (Ptr a)
forall a b. (a -> b) -> a -> b
$ Int -> CSize
forall a. Enum a => Int -> a
toEnum (Int -> CSize) -> Int -> CSize
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    fp <- newForeignPtr fftw_free p
    return $ MS.MVector n fp
{-# INLINE newFFTVector #-}


-----------------------
-- Transforms: methods of plan creation.

-- | A transform which may be applied to vectors of different sizes.
data Transform a b = Transform {
                        forall a b. Transform a b -> Int -> Int
inputSize :: Int -> Int,
                        forall a b. Transform a b -> Int -> Int
outputSize :: Int -> Int,
                        forall a b. Transform a b -> Int -> Int
creationSizeFromInput :: Int -> Int,
                        forall a b.
Transform a b -> CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
makePlan :: CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan),
                        forall a b. Transform a b -> Int -> Plan a b -> Plan a b
normalization :: Int -> Plan a b -> Plan a b
                    }

-- | Create a 'Plan' of a specific size for this transform.
planOfType :: (Storable a, Storable b) => PlanType
                                -> Transform a b -> Int -> Plan a b
planOfType :: forall a b.
(Storable a, Storable b) =>
PlanType -> Transform a b -> Int -> Plan a b
planOfType PlanType
ptype Transform{Int -> Int
Int -> Plan a b -> Plan a b
CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
inputSize :: forall a b. Transform a b -> Int -> Int
outputSize :: forall a b. Transform a b -> Int -> Int
creationSizeFromInput :: forall a b. Transform a b -> Int -> Int
makePlan :: forall a b.
Transform a b -> CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
normalization :: forall a b. Transform a b -> Int -> Plan a b -> Plan a b
inputSize :: Int -> Int
outputSize :: Int -> Int
creationSizeFromInput :: Int -> Int
makePlan :: CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
normalization :: Int -> Plan a b -> Plan a b
..} Int
n
  | Int
m_in Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
m_out Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Char] -> Plan a b
forall a. HasCallStack => [Char] -> a
error [Char]
"Can't (yet) plan for empty arrays!"
  | Bool
otherwise  = IO (Plan a b) -> Plan a b
forall a. IO a -> a
unsafePerformIO (IO (Plan a b) -> Plan a b) -> IO (Plan a b) -> Plan a b
forall a b. (a -> b) -> a -> b
$ do
    planInput <- Int -> IO (MVector RealWorld a)
forall a. Storable a => Int -> IO (MVector RealWorld a)
newFFTVector Int
m_in
    planOutput <- newFFTVector m_out
    MS.unsafeWith planInput $ \Ptr a
inP -> MVector RealWorld b -> (Ptr b -> IO (Plan a b)) -> IO (Plan a b)
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld b
planOutput ((Ptr b -> IO (Plan a b)) -> IO (Plan a b))
-> (Ptr b -> IO (Plan a b)) -> IO (Plan a b)
forall a b. (a -> b) -> a -> b
$ \Ptr b
outP -> do
        pPlan <- CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
makePlan (Int -> CInt
forall a. Enum a => Int -> a
toEnum Int
n) Ptr a
inP Ptr b
outP (CFlags -> IO (Ptr CPlan)) -> CFlags -> IO (Ptr CPlan)
forall a b. (a -> b) -> a -> b
$ PlanType -> Preservation -> CFlags
planInitFlags PlanType
ptype Preservation
DestroyInput
        cPlan <- newPlan pPlan
        -- Use unsafeWith here to ensure that the Storable MVectors' ForeignPtrs
        -- aren't released too soon:
        let planExecute = MVector RealWorld a -> (Ptr a -> IO ()) -> IO ()
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld a
planInput ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
_ ->
                            MVector RealWorld b -> (Ptr b -> IO ()) -> IO ()
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld b
planOutput ((Ptr b -> IO ()) -> IO ()) -> (Ptr b -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr b
_ ->
                                CPlan -> (Ptr CPlan -> IO ()) -> IO ()
forall a. CPlan -> (Ptr CPlan -> IO a) -> IO a
withPlan CPlan
cPlan Ptr CPlan -> IO ()
fftw_execute
        return $ normalization n $ Plan {..}
  where
    m_in :: Int
m_in = Int -> Int
inputSize Int
n
    m_out :: Int
m_out = Int -> Int
outputSize Int
n
{-# INLINE planOfType #-}

-- | Create a 'Plan' of a specific size.  This function is equivalent to
-- @'planOfType' 'Estimate'@.
plan :: (Storable a, Storable b) => Transform a b -> Int -> Plan a b
plan :: forall a b.
(Storable a, Storable b) =>
Transform a b -> Int -> Plan a b
plan = PlanType -> Transform a b -> Int -> Plan a b
forall a b.
(Storable a, Storable b) =>
PlanType -> Transform a b -> Int -> Plan a b
planOfType PlanType
Estimate
{-# INLINE plan #-}

-- | Create and run a 'Plan' for the given transform.
run :: (Vector v a, Vector v b, Storable a, Storable b)
            => Transform a b -> v a -> v b
run :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Storable a, Storable b) =>
Transform a b -> v a -> v b
run Transform a b
p = \v a
v -> Plan a b -> v a -> v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Storable a, Storable b) =>
Plan a b -> v a -> v b
execute
            (PlanType -> Transform a b -> Int -> Plan a b
forall a b.
(Storable a, Storable b) =>
PlanType -> Transform a b -> Int -> Plan a b
planOfType PlanType
Estimate Transform a b
p (Int -> Plan a b) -> Int -> Plan a b
forall a b. (a -> b) -> a -> b
$ Transform a b -> Int -> Int
forall a b. Transform a b -> Int -> Int
creationSizeFromInput Transform a b
p (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v)
            v a
v
{-# INLINE run #-}

-----------------------
-- TransformND: methods of plan creation for multi-dimensional plans.

-- | A transform which may be applied to vectors of different sizes.
--
-- @since 0.2
data TransformND a b = TransformND {
                        forall a b. TransformND a b -> Int -> Int
inputSizeND :: Int -> Int,
                        forall a b. TransformND a b -> Int -> Int
outputSizeND :: Int -> Int,
                        forall a b. TransformND a b -> Int -> Int
creationSizeFromInputND :: Int -> Int,
                        forall a b.
TransformND a b
-> CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
makePlanND :: CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan),
                        forall a b. TransformND a b -> Vector Int -> Plan a b -> Plan a b
normalizationND :: VS.Vector Int -> Plan a b -> Plan a b
                    }

-- | Create a 'Plan' of a specific size for this transform.
-- 'dims' must have rank greater or equal to 1
--
-- @since 0.2
planOfTypeND :: (Storable a, Storable b) => PlanType
                                -> TransformND a b -> VS.Vector Int -> Plan a b
planOfTypeND :: forall a b.
(Storable a, Storable b) =>
PlanType -> TransformND a b -> Vector Int -> Plan a b
planOfTypeND PlanType
ptype TransformND{Int -> Int
CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
Vector Int -> Plan a b -> Plan a b
inputSizeND :: forall a b. TransformND a b -> Int -> Int
outputSizeND :: forall a b. TransformND a b -> Int -> Int
creationSizeFromInputND :: forall a b. TransformND a b -> Int -> Int
makePlanND :: forall a b.
TransformND a b
-> CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
normalizationND :: forall a b. TransformND a b -> Vector Int -> Plan a b -> Plan a b
inputSizeND :: Int -> Int
outputSizeND :: Int -> Int
creationSizeFromInputND :: Int -> Int
makePlanND :: CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
normalizationND :: Vector Int -> Plan a b -> Plan a b
..} Vector Int
dims
  | Int
m_in Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
m_out Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Char] -> Plan a b
forall a. HasCallStack => [Char] -> a
error [Char]
"Can't (yet) plan for empty arrays!"
  | Bool
otherwise  = IO (Plan a b) -> Plan a b
forall a. IO a -> a
unsafePerformIO (IO (Plan a b) -> Plan a b) -> IO (Plan a b) -> Plan a b
forall a b. (a -> b) -> a -> b
$ do
    mdims <- Vector CInt -> IO (Mutable Vector (PrimState IO) CInt)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
unsafeThaw (Vector CInt -> IO (Mutable Vector (PrimState IO) CInt))
-> Vector CInt -> IO (Mutable Vector (PrimState IO) CInt)
forall a b. (a -> b) -> a -> b
$ (Int -> CInt) -> Vector Int -> Vector CInt
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map Int -> CInt
forall a. Enum a => Int -> a
toEnum Vector Int
dims
    planInput <- newFFTVector m_in
    planOutput <- newFFTVector m_out
    MS.unsafeWith mdims $ \Ptr CInt
dimsP -> MVector RealWorld a -> (Ptr a -> IO (Plan a b)) -> IO (Plan a b)
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld a
planInput ((Ptr a -> IO (Plan a b)) -> IO (Plan a b))
-> (Ptr a -> IO (Plan a b)) -> IO (Plan a b)
forall a b. (a -> b) -> a -> b
$ \Ptr a
inP -> MVector RealWorld b -> (Ptr b -> IO (Plan a b)) -> IO (Plan a b)
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld b
planOutput ((Ptr b -> IO (Plan a b)) -> IO (Plan a b))
-> (Ptr b -> IO (Plan a b)) -> IO (Plan a b)
forall a b. (a -> b) -> a -> b
$ \Ptr b
outP -> do
        pPlan <- CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan)
makePlanND (Int -> CInt
forall a. Enum a => Int -> a
toEnum (Int -> CInt) -> Int -> CInt
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length Vector Int
dims) Ptr CInt
dimsP Ptr a
inP Ptr b
outP (CFlags -> IO (Ptr CPlan)) -> CFlags -> IO (Ptr CPlan)
forall a b. (a -> b) -> a -> b
$ PlanType -> Preservation -> CFlags
planInitFlags PlanType
ptype Preservation
DestroyInput
        cPlan <- newPlan pPlan
        -- Use unsafeWith here to ensure that the Storable MVectors' ForeignPtrs
        -- aren't released too soon:
        let planExecute = MVector RealWorld a -> (Ptr a -> IO ()) -> IO ()
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld a
planInput ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
_ ->
                                MVector RealWorld b -> (Ptr b -> IO ()) -> IO ()
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MS.unsafeWith MVector RealWorld b
planOutput ((Ptr b -> IO ()) -> IO ()) -> (Ptr b -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr b
_ ->
                                CPlan -> (Ptr CPlan -> IO ()) -> IO ()
forall a. CPlan -> (Ptr CPlan -> IO a) -> IO a
withPlan CPlan
cPlan Ptr CPlan -> IO ()
fftw_execute
        return $ normalizationND dims $ Plan {..}
  where
    m :: Int
m = Vector Int -> Int
forall (v :: * -> *) a. (Vector v a, Num a) => v a -> a
V.product (Vector Int -> Int) -> Vector Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int
forall (v :: * -> *) a. Vector v a => v a -> v a
V.init Vector Int
dims
    m_in :: Int
m_in = Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
inputSizeND (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> a
V.last Vector Int
dims)
    m_out :: Int
m_out = Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
outputSizeND (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> a
V.last Vector Int
dims)
{-# INLINE planOfTypeND #-}

-- | Create a 'Plan' of a specific size.  This function is equivalent to
-- @'planOfType' 'Estimate'@.
--
-- @since 0.2
planND :: (Storable a, Storable b) => TransformND a b -> VS.Vector Int -> Plan a b
planND :: forall a b.
(Storable a, Storable b) =>
TransformND a b -> Vector Int -> Plan a b
planND = PlanType -> TransformND a b -> Vector Int -> Plan a b
forall a b.
(Storable a, Storable b) =>
PlanType -> TransformND a b -> Vector Int -> Plan a b
planOfTypeND PlanType
Estimate
{-# INLINE planND #-}

-- | Create and run a 'Plan' for the given transform.
--
-- @since 0.2
runND :: (Vector v a, Vector v b, Storable a, Storable b)
            => TransformND a b -> VS.Vector Int ->  v a -> v b
runND :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Storable a, Storable b) =>
TransformND a b -> Vector Int -> v a -> v b
runND TransformND a b
p = \Vector Int
dims v a
v ->
  let creationSize :: Vector Int
creationSize = Vector Int -> Vector Int
forall (v :: * -> *) a. Vector v a => v a -> v a
V.init Vector Int
dims Vector Int -> Int -> Vector Int
forall (v :: * -> *) a. Vector v a => v a -> a -> v a
`V.snoc` TransformND a b -> Int -> Int
forall a b. TransformND a b -> Int -> Int
creationSizeFromInputND TransformND a b
p (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> a
V.last Vector Int
dims) in
    Plan a b -> v a -> v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Storable a, Storable b) =>
Plan a b -> v a -> v b
execute
    (PlanType -> TransformND a b -> Vector Int -> Plan a b
forall a b.
(Storable a, Storable b) =>
PlanType -> TransformND a b -> Vector Int -> Plan a b
planOfTypeND PlanType
Estimate TransformND a b
p Vector Int
creationSize)
    v a
v
{-# INLINE runND #-}

---------------------------
-- For scaling input/output:

class Scalable a where
    scaleByD :: Double -> a -> a

instance Scalable Double where
    scaleByD :: Double -> Double -> Double
scaleByD = Double -> Double -> Double
forall a. Num a => a -> a -> a
(*)
    {-# INLINE scaleByD #-}

instance Scalable (Complex Double) where
    scaleByD :: Double -> Complex Double -> Complex Double
scaleByD Double
s (Double
x:+Double
y) = Double
sDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
x Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
sDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
y
    {-# INLINE scaleByD #-}


{-# INLINE modifyInput #-}
modifyInput :: (MS.MVector RealWorld a -> IO ()) -> Plan a b -> Plan a b
modifyInput :: forall a b. (MVector RealWorld a -> IO ()) -> Plan a b -> Plan a b
modifyInput MVector RealWorld a -> IO ()
f p :: Plan a b
p@Plan{IO ()
MVector RealWorld a
MVector RealWorld b
planInput :: forall a b. Plan a b -> MVector RealWorld a
planOutput :: forall a b. Plan a b -> MVector RealWorld b
planExecute :: forall a b. Plan a b -> IO ()
planInput :: MVector RealWorld a
planOutput :: MVector RealWorld b
planExecute :: IO ()
..} = Plan a b
p {planExecute = f planInput >> planExecute}

{-# INLINE modifyOutput #-}
modifyOutput :: (MS.MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput :: forall b a. (MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput MVector RealWorld b -> IO ()
f p :: Plan a b
p@Plan{IO ()
MVector RealWorld b
MVector RealWorld a
planInput :: forall a b. Plan a b -> MVector RealWorld a
planOutput :: forall a b. Plan a b -> MVector RealWorld b
planExecute :: forall a b. Plan a b -> IO ()
planInput :: MVector RealWorld a
planOutput :: MVector RealWorld b
planExecute :: IO ()
..} = Plan a b
p {planExecute = planExecute >> f planOutput}

{-# INLINE constMultOutput #-}
constMultOutput :: (Storable b, Scalable b) => Double -> Plan a b -> Plan a b
constMultOutput :: forall b a.
(Storable b, Scalable b) =>
Double -> Plan a b -> Plan a b
constMultOutput !Double
s = (MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
forall b a. (MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput (Double -> MVector RealWorld b -> IO ()
forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC Double
s)

{-# INLINE multC #-}
multC :: (Storable a, Scalable a) => Double -> MS.MVector RealWorld a -> IO ()
multC :: forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC !Double
s MVector RealWorld a
v = [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> MVector RealWorld a -> Int -> (a -> a) -> IO ()
forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld a
v Int
k (Double -> a -> a
forall a. Scalable a => Double -> a -> a
scaleByD Double
s)
  where !n :: Int
n = MVector RealWorld a -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld a
v

-- | Helper function; seems like it should be in the vector package...
{-# INLINE unsafeModify #-}
unsafeModify :: (Storable a)
                => MS.MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify :: forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld a
v Int
k a -> a
f = MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MS.unsafeRead MVector RealWorld a
MVector (PrimState IO) a
v Int
k IO a -> (a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MS.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
v Int
k (a -> IO ()) -> (a -> a) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f

plannerLock :: MVar ()
plannerLock :: MVar ()
plannerLock = IO (MVar ()) -> MVar ()
forall a. IO a -> a
unsafePerformIO (IO (MVar ()) -> MVar ()) -> IO (MVar ()) -> MVar ()
forall a b. (a -> b) -> a -> b
$ () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
{-# NOINLINE plannerLock #-}

-- | Calls to the FFTW planner are non-reentrant. Here we take a mutex to
-- ensure thread safety.
withPlanner :: IO a -> IO a
withPlanner :: forall a. IO a -> IO a
withPlanner IO a
action = do
    MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
plannerLock
    res <- IO a
action
    putMVar plannerLock ()
    return res