{-# LINE 1 "Numeric/FFT/Vector/Base.hsc" #-}
module Numeric.FFT.Vector.Base(
Transform(..),
planOfType,
PlanType(..),
plan,
run,
TransformND(..),
planOfTypeND,
planND,
runND,
Plan(),
planInputSize,
planOutputSize,
execute,
executeM,
withPlanner,
CFlags,
CPlan,
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)
data PlanType = Estimate | Measure | Patient | Exhaustive
data Preservation = PreserveInput | DestroyInput
type CFlags = CUInt
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
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 ()
}
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
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
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 ->
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 #-}
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 :: 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 #-}
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 #-}
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
}
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
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 #-}
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 #-}
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 #-}
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
}
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
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 #-}
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 #-}
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 #-}
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
{-# 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 #-}
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