{- |
This module provides normalized versions of the transforms in @fftw@.

All of the transforms are normalized so that

 - Each transform is unitary, i.e., preserves the inner product and the sum-of-squares norm of its input.

 - Each backwards transform is the inverse of the corresponding forwards transform.

(Both conditions only hold approximately, due to floating point precision.)

For more information on the underlying transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
-}
module Numeric.FFT.Vector.Unitary(
                -- * Creating and executing 'Plan's
                run,
                plan,
                execute,
                -- * Complex-to-complex transforms
                dft,
                idft,
                -- * Real-to-complex transforms
                dftR2C,
                dftC2R,
                -- * Discrete cosine transforms
                -- $dct_size
                dct2,
                idct2,
                dct4,
                ) where

import Numeric.FFT.Vector.Base
import qualified Numeric.FFT.Vector.Unnormalized as U
import Data.Complex
import qualified Data.Vector.Storable.Mutable as MS
import Control.Monad.Primitive(RealWorld)

-- | A discrete Fourier transform. The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(-2pi i j k\/n)@
dft :: Transform (Complex Double) (Complex Double)
dft :: Transform (Complex Double) (Complex Double)
dft = Transform (Complex Double) (Complex Double)
U.dft {normalization = \Int
n -> Double
-> Plan (Complex Double) (Complex Double)
-> Plan (Complex Double) (Complex Double)
forall b a.
(Storable b, Scalable b) =>
Double -> Plan a b -> Plan a b
constMultOutput (Double
 -> Plan (Complex Double) (Complex Double)
 -> Plan (Complex Double) (Complex Double))
-> Double
-> Plan (Complex Double) (Complex Double)
-> Plan (Complex Double) (Complex Double)
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n)}

-- | An inverse discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(2pi i j k\/n)@
idft :: Transform (Complex Double) (Complex Double)
idft :: Transform (Complex Double) (Complex Double)
idft = Transform (Complex Double) (Complex Double)
U.idft {normalization = \Int
n -> Double
-> Plan (Complex Double) (Complex Double)
-> Plan (Complex Double) (Complex Double)
forall b a.
(Storable b, Scalable b) =>
Double -> Plan a b -> Plan a b
constMultOutput (Double
 -> Plan (Complex Double) (Complex Double)
 -> Plan (Complex Double) (Complex Double))
-> Double
-> Plan (Complex Double) (Complex Double)
-> Plan (Complex Double) (Complex Double)
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n)}

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: Transform Double (Complex Double)
dftR2C :: Transform Double (Complex Double)
dftR2C = Transform Double (Complex Double)
U.dftR2C {normalization = \Int
n -> (MVector RealWorld (Complex Double) -> IO ())
-> Plan Double (Complex Double) -> Plan Double (Complex Double)
forall b a. (MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput ((MVector RealWorld (Complex Double) -> IO ())
 -> Plan Double (Complex Double) -> Plan Double (Complex Double))
-> (MVector RealWorld (Complex Double) -> IO ())
-> Plan Double (Complex Double)
-> Plan Double (Complex Double)
forall a b. (a -> b) -> a -> b
$
                    Double -> Int -> MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling (Double -> Double
forall a. Floating a => a -> a
sqrt Double
2) Int
n
        }

-- | A normalized backward discrete Fourier transform which is the left inverse of
-- 'U.dftR2C'.  (Specifically, @run dftC2R . run dftR2C == id@.)
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
--
dftC2R :: Transform (Complex Double) Double
dftC2R :: Transform (Complex Double) Double
dftC2R = Transform (Complex Double) Double
U.dftC2R {normalization = \Int
n -> (MVector RealWorld (Complex Double) -> IO ())
-> Plan (Complex Double) Double -> Plan (Complex Double) Double
forall a b. (MVector RealWorld a -> IO ()) -> Plan a b -> Plan a b
modifyInput ((MVector RealWorld (Complex Double) -> IO ())
 -> Plan (Complex Double) Double -> Plan (Complex Double) Double)
-> (MVector RealWorld (Complex Double) -> IO ())
-> Plan (Complex Double) Double
-> Plan (Complex Double) Double
forall a b. (a -> b) -> a -> b
$
                    Double -> Int -> MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling (Double -> Double
forall a. Floating a => a -> a
sqrt Double
0.5) Int
n
        }

complexR2CScaling :: Double -> Int -> MS.MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling :: Double -> Int -> MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling !Double
t !Int
n !MVector RealWorld (Complex Double)
a = do
    let !s1 :: Double
s1 = Double -> Double
forall a. Floating a => a -> a
sqrt (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n)
    let !s2 :: Double
s2 = Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s1
    let len :: Int
len = MVector RealWorld (Complex Double) -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld (Complex Double)
a
    -- Justification for the use of unsafeModify:
    -- The output size is 2n+1; so if n>0 then the output size is >=1;
    -- and if n even then the output size is >=3.
    MVector RealWorld (Complex Double)
-> Int -> (Complex Double -> Complex Double) -> IO ()
forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld (Complex Double)
a Int
0 ((Complex Double -> Complex Double) -> IO ())
-> (Complex Double -> Complex Double) -> IO ()
forall a b. (a -> b) -> a -> b
$ Double -> Complex Double -> Complex Double
forall a. Scalable a => Double -> a -> a
scaleByD Double
s1
    if Int -> Bool
forall a. Integral a => a -> Bool
odd Int
n
        then Double -> MVector RealWorld (Complex Double) -> IO ()
forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC Double
s2 (Int
-> Int
-> MVector RealWorld (Complex Double)
-> MVector RealWorld (Complex Double)
forall a s. Storable a => Int -> Int -> MVector s a -> MVector s a
MS.unsafeSlice Int
1 (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) MVector RealWorld (Complex Double)
a)
        else do
            MVector RealWorld (Complex Double)
-> Int -> (Complex Double -> Complex Double) -> IO ()
forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld (Complex Double)
a (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((Complex Double -> Complex Double) -> IO ())
-> (Complex Double -> Complex Double) -> IO ()
forall a b. (a -> b) -> a -> b
$ Double -> Complex Double -> Complex Double
forall a. Scalable a => Double -> a -> a
scaleByD Double
s1
            Double -> MVector RealWorld (Complex Double) -> IO ()
forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC Double
s2 (Int
-> Int
-> MVector RealWorld (Complex Double)
-> MVector RealWorld (Complex Double)
forall a s. Storable a => Int -> Int -> MVector s a -> MVector s a
MS.unsafeSlice Int
1 (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2) MVector RealWorld (Complex Double)
a)


-- $dct_size
-- Some normalized real-even (DCT).  The input and output sizes
-- are the same (@n@).


-- | A type-4 discrete cosine transform.  It is its own inverse.
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)(k+1\/2)\/n)@
dct4 :: Transform Double Double
dct4 :: Transform Double Double
dct4 = Transform Double Double
U.dct4 {normalization = \Int
n -> Double -> Plan Double Double -> Plan Double Double
forall b a.
(Storable b, Scalable b) =>
Double -> Plan a b -> Plan a b
constMultOutput (Double -> Plan Double Double -> Plan Double Double)
-> Double -> Plan Double Double -> Plan Double Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n)}

-- | A type-2 discrete cosine transform.  Its inverse is 'dct3'.
--
-- @y_k = w(k) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)k\/n);@
-- where
-- @w(0)=1\/sqrt n@, and @w(k)=sqrt(2\/n)@ for @k>0@.
dct2 :: Transform Double Double
dct2 :: Transform Double Double
dct2 = Transform Double Double
U.dct2 {normalization = \Int
n -> (MVector RealWorld Double -> IO ())
-> Plan Double Double -> Plan Double Double
forall b a. (MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput ((MVector RealWorld Double -> IO ())
 -> Plan Double Double -> Plan Double Double)
-> (MVector RealWorld Double -> IO ())
-> Plan Double Double
-> Plan Double Double
forall a b. (a -> b) -> a -> b
$ \MVector RealWorld Double
a -> do
    let n' :: Double
n' = Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n
    let !s1 :: Double
s1 = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
4Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
n')
    let !s2 :: Double
s2 = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
n')
    MVector RealWorld Double -> Int -> (Double -> Double) -> IO ()
forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld Double
a Int
0 (Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
s1)
    Double -> MVector RealWorld Double -> IO ()
forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC Double
s2 (Int -> Int -> MVector RealWorld Double -> MVector RealWorld Double
forall a s. Storable a => Int -> Int -> MVector s a -> MVector s a
MS.unsafeSlice Int
1 (MVector RealWorld Double -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld Double
aInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) MVector RealWorld Double
a)
    }

-- | A type-3 discrete cosine transform which is the inverse of 'dct2'.
--
-- @y_k = (-1)^k w(n-1) x_(n-1) + 2 sum_(j=0)^(n-2) w(j) x_j sin(pi(j+1)(k+1\/2)/n);@
-- where
-- @w(0)=1\/sqrt(n)@, and @w(k)=1/sqrt(2n)@ for @k>0@.
idct2 :: Transform Double Double
idct2 :: Transform Double Double
idct2 = Transform Double Double
U.dct3 {normalization = \Int
n -> (MVector RealWorld Double -> IO ())
-> Plan Double Double -> Plan Double Double
forall a b. (MVector RealWorld a -> IO ()) -> Plan a b -> Plan a b
modifyInput ((MVector RealWorld Double -> IO ())
 -> Plan Double Double -> Plan Double Double)
-> (MVector RealWorld Double -> IO ())
-> Plan Double Double
-> Plan Double Double
forall a b. (a -> b) -> a -> b
$ \MVector RealWorld Double
a -> do
    let n' :: Double
n' = Int -> Double
forall a. Enum a => Int -> a
toEnum Int
n
    let !s1 :: Double
s1 = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n'
    let !s2 :: Double
s2 = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
n')
    MVector RealWorld Double -> Int -> (Double -> Double) -> IO ()
forall a.
Storable a =>
MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify MVector RealWorld Double
a Int
0 (Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
s1)
    Double -> MVector RealWorld Double -> IO ()
forall a.
(Storable a, Scalable a) =>
Double -> MVector RealWorld a -> IO ()
multC Double
s2 (Int -> Int -> MVector RealWorld Double -> MVector RealWorld Double
forall a s. Storable a => Int -> Int -> MVector s a -> MVector s a
MS.unsafeSlice Int
1 (MVector RealWorld Double -> Int
forall a s. Storable a => MVector s a -> Int
MS.length MVector RealWorld Double
aInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) MVector RealWorld Double
a)
    }