-- SPDX-License-Identifier: MPL-2.0

{- | CPU reference Fourier transforms. Forward exponent is negative; inverse
exponent is positive and divided by N. ND data are contiguous row-major with
the last axis fastest. Plans are immutable, reusable and cache N roots; no
global cache or thread-local state is retained. Scalar transforms accept
1..65536 samples (direct quadratic transforms at most 4096).

DCT/DST I--IV use orthonormal matrices, so inverse means transpose. Definitions
follow <https://fftw.org/fftw3_doc/What-FFTW-Really-Computes.html> with explicit
endpoint weights added for orthogonality. Algorithms are reference CPU paths,
not hardware FFTW bindings.
-}
module HMath.Fourier (
    Direction (..),
    FFTAlgorithm (..),
    FFTPlan,
    planFFT,
    executeFFT,
    fft,
    ifft,
    dft,
    fftND,
    fftBatch,
    realFFT,
    inverseRealFFT,
    fftInPlace,
    frequencies,
    fftShift,
    ifftShift,
    TransformType (..),
    dct,
    idct,
    dst,
    idst,
    hartley,
    STFT,
    stft,
    stftFrames,
    istft,
    magnitudeSpectrum,
    phaseSpectrum,
    powerSpectrum,
    periodogram,
    welch,
    crossSpectrum,
    spectrogram,
    cepstrum,
    analyticSignal,
    hilbert,
    chirpZ,
    goertzel,
    fftConvolve,
    fftCorrelate,
) where

import Data.Complex (Complex ((:+)), cis, conjugate, imagPart, magnitude, phase, realPart)
import qualified Data.Foldable as Foldable (foldl')
import Foreign.Marshal.Array (peekArray, pokeArray)
import Foreign.Ptr (Ptr, nullPtr)
import HMath.Error (MathError (..))
import HMath.Signal.Internal (finite, require)

-- | Transform exponent convention; Inverse additionally divides by sample count.
data Direction = Forward | Inverse deriving (Direction -> Direction -> Bool
(Direction -> Direction -> Bool)
-> (Direction -> Direction -> Bool) -> Eq Direction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Direction -> Direction -> Bool
== :: Direction -> Direction -> Bool
$c/= :: Direction -> Direction -> Bool
/= :: Direction -> Direction -> Bool
Eq, Int -> Direction -> ShowS
[Direction] -> ShowS
Direction -> String
(Int -> Direction -> ShowS)
-> (Direction -> String)
-> ([Direction] -> ShowS)
-> Show Direction
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Direction -> ShowS
showsPrec :: Int -> Direction -> ShowS
$cshow :: Direction -> String
show :: Direction -> String
$cshowList :: [Direction] -> ShowS
showList :: [Direction] -> ShowS
Show)

-- | Explicit reference factorization, or Automatic mixed-radix with Bluestein for larger prime factors.
data FFTAlgorithm = Automatic | Direct | Radix2 | Radix4 | MixedRadix | SplitRadix | Bluestein | Rader
    deriving (FFTAlgorithm -> FFTAlgorithm -> Bool
(FFTAlgorithm -> FFTAlgorithm -> Bool)
-> (FFTAlgorithm -> FFTAlgorithm -> Bool) -> Eq FFTAlgorithm
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FFTAlgorithm -> FFTAlgorithm -> Bool
== :: FFTAlgorithm -> FFTAlgorithm -> Bool
$c/= :: FFTAlgorithm -> FFTAlgorithm -> Bool
/= :: FFTAlgorithm -> FFTAlgorithm -> Bool
Eq, Int -> FFTAlgorithm -> ShowS
[FFTAlgorithm] -> ShowS
FFTAlgorithm -> String
(Int -> FFTAlgorithm -> ShowS)
-> (FFTAlgorithm -> String)
-> ([FFTAlgorithm] -> ShowS)
-> Show FFTAlgorithm
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FFTAlgorithm -> ShowS
showsPrec :: Int -> FFTAlgorithm -> ShowS
$cshow :: FFTAlgorithm -> String
show :: FFTAlgorithm -> String
$cshowList :: [FFTAlgorithm] -> ShowS
showList :: [FFTAlgorithm] -> ShowS
Show)

-- | Opaque immutable plan with cached unit-circle roots and a fixed input length.
data FFTPlan = FFTPlan Direction FFTAlgorithm Int [Complex Double] deriving (FFTPlan -> FFTPlan -> Bool
(FFTPlan -> FFTPlan -> Bool)
-> (FFTPlan -> FFTPlan -> Bool) -> Eq FFTPlan
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FFTPlan -> FFTPlan -> Bool
== :: FFTPlan -> FFTPlan -> Bool
$c/= :: FFTPlan -> FFTPlan -> Bool
/= :: FFTPlan -> FFTPlan -> Bool
Eq, Int -> FFTPlan -> ShowS
[FFTPlan] -> ShowS
FFTPlan -> String
(Int -> FFTPlan -> ShowS)
-> (FFTPlan -> String) -> ([FFTPlan] -> ShowS) -> Show FFTPlan
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FFTPlan -> ShowS
showsPrec :: Int -> FFTPlan -> ShowS
$cshow :: FFTPlan -> String
show :: FFTPlan -> String
$cshowList :: [FFTPlan] -> ShowS
showList :: [FFTPlan] -> ShowS
Show)

-- | Orthonormal cosine/sine basis type I, II, III or IV.
data TransformType = TypeI | TypeII | TypeIII | TypeIV deriving (TransformType -> TransformType -> Bool
(TransformType -> TransformType -> Bool)
-> (TransformType -> TransformType -> Bool) -> Eq TransformType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TransformType -> TransformType -> Bool
== :: TransformType -> TransformType -> Bool
$c/= :: TransformType -> TransformType -> Bool
/= :: TransformType -> TransformType -> Bool
Eq, Int -> TransformType -> ShowS
[TransformType] -> ShowS
TransformType -> String
(Int -> TransformType -> ShowS)
-> (TransformType -> String)
-> ([TransformType] -> ShowS)
-> Show TransformType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TransformType -> ShowS
showsPrec :: Int -> TransformType -> ShowS
$cshow :: TransformType -> String
show :: TransformType -> String
$cshowList :: [TransformType] -> ShowS
showList :: [TransformType] -> ShowS
Show)

finiteC :: Complex Double -> Bool
finiteC :: Complex Double -> Bool
finiteC Complex Double
x = Double -> Bool
finite (Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
x) Bool -> Bool -> Bool
&& Double -> Bool
finite (Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
x)
checkC :: [Complex Double] -> Either MathError [Complex Double]
checkC :: [Complex Double] -> Either MathError [Complex Double]
checkC [Complex Double]
xs = if (Complex Double -> Bool) -> [Complex Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Complex Double -> Bool
finiteC [Complex Double]
xs then [Complex Double] -> Either MathError [Complex Double]
forall a b. b -> Either a b
Right [Complex Double]
xs else MathError -> Either MathError [Complex Double]
forall a b. a -> Either a b
Left MathError
NonFiniteResult
powerOf :: Int -> Int -> Bool
powerOf :: Int -> Int -> Bool
powerOf Int
r Int
n | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Bool
True | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Bool
False | Bool
otherwise = Int -> Int -> Bool
powerOf Int
r (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
r)
prime :: Int -> Bool
prime :: Int -> Bool
prime Int
n = Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Int -> Bool) -> (Int -> Int) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Integral a => a -> a -> a
mod Int
n) [Int
2 .. Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Double))]

-- | Create a reusable plan after checking length, prime or radix requirements. No input data or external buffers are retained.
planFFT :: Direction -> FFTAlgorithm -> Int -> Either MathError FFTPlan
planFFT :: Direction -> FFTAlgorithm -> Int -> Either MathError FFTPlan
planFFT Direction
direction FFTAlgorithm
method Int
n = do
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536 Bool -> Bool -> Bool
&& (FFTAlgorithm
method FFTAlgorithm -> FFTAlgorithm -> Bool
forall a. Eq a => a -> a -> Bool
/= FFTAlgorithm
Direct Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096))
        String
"FFT length outside reference budget"
    Bool -> String -> Either MathError ()
require
        (FFTAlgorithm
method FFTAlgorithm -> FFTAlgorithm -> Bool
forall a. Eq a => a -> a -> Bool
/= FFTAlgorithm
Radix2 Bool -> Bool -> Bool
&& FFTAlgorithm
method FFTAlgorithm -> FFTAlgorithm -> Bool
forall a. Eq a => a -> a -> Bool
/= FFTAlgorithm
SplitRadix Bool -> Bool -> Bool
|| Int -> Int -> Bool
powerOf Int
2 Int
n)
        String
"radix-2/split-radix length must be a power of two"
    Bool -> String -> Either MathError ()
require (FFTAlgorithm
method FFTAlgorithm -> FFTAlgorithm -> Bool
forall a. Eq a => a -> a -> Bool
/= FFTAlgorithm
Radix4 Bool -> Bool -> Bool
|| Int -> Int -> Bool
powerOf Int
4 Int
n) String
"radix-4 length must be a power of four"
    Bool -> String -> Either MathError ()
require (FFTAlgorithm
method FFTAlgorithm -> FFTAlgorithm -> Bool
forall a. Eq a => a -> a -> Bool
/= FFTAlgorithm
Rader Bool -> Bool -> Bool
|| Int -> Bool
prime Int
n) String
"Rader length must be prime"
    let sign :: Double
sign = if Direction
direction Direction -> Direction -> Bool
forall a. Eq a => a -> a -> Bool
== Direction
Forward then -Double
1 else Double
1
    FFTPlan -> Either MathError FFTPlan
forall a b. b -> Either a b
Right
        ( Direction -> FFTAlgorithm -> Int -> [Complex Double] -> FFTPlan
FFTPlan
            Direction
direction
            FFTAlgorithm
method
            Int
n
            [Double -> Complex Double
forall a. Floating a => a -> Complex a
cis (Double
sign Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        )

-- | Apply a plan to exactly its declared number of finite complex samples; reject nonfinite intermediate results.
executeFFT :: FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
executeFFT :: FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
executeFFT (FFTPlan Direction
direction FFTAlgorithm
method Int
n [Complex Double]
roots) [Complex Double]
xs = do
    Bool -> String -> Either MathError ()
require ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n) String
"FFT plan/input length mismatch"
    _ <- [Complex Double] -> Either MathError [Complex Double]
checkC [Complex Double]
xs
    let raw = case FFTAlgorithm
method of
            FFTAlgorithm
Direct -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}. Num a => [a] -> [a] -> [a]
direct [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
Radix2 -> Int -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}. Num a => Int -> [a] -> [a] -> [a]
radix Int
2 [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
Radix4 -> Int -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}. Num a => Int -> [a] -> [a] -> [a]
radix Int
4 [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
SplitRadix -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}. Num a => [a] -> [a] -> [a]
split [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
Bluestein -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a} {a}.
(RealFloat a, Ord a, Num a) =>
[Complex a] -> [Complex a] -> [Complex a]
blue [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
Rader -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
rader [Complex Double]
roots [Complex Double]
xs
            FFTAlgorithm
_ -> [Complex Double] -> [Complex Double] -> [Complex Double]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
mixed [Complex Double]
roots [Complex Double]
xs
        scaled = if Direction
direction Direction -> Direction -> Bool
forall a. Eq a => a -> a -> Bool
== Direction
Inverse then (Complex Double -> Complex Double)
-> [Complex Double] -> [Complex Double]
forall a b. (a -> b) -> [a] -> [b]
map (Complex Double -> Complex Double -> Complex Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0)) [Complex Double]
raw else [Complex Double]
raw
    checkC scaled
  where
    direct :: [a] -> [a] -> [a]
direct [a]
rs [a]
ys =
        [[a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a
y a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys)) | (Int
j, a
y) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [a]
ys] | Int
k <- [Int
0 .. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    radix :: Int -> [a] -> [a] -> [a]
radix Int
r [a]
rs [a]
ys
        | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = [a]
ys
        | Bool
otherwise =
            let m :: Int
m = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
r
                sub :: [[a]]
sub =
                    [ Int -> [a] -> [a] -> [a]
radix Int
r [[a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [[a]
ys [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                    | Int
q <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
             in [ [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[a]]
sub [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
q [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m) a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys)) | Int
q <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                | Int
k <- [Int
0 .. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
    mixed :: [Complex a] -> [Complex a] -> [Complex a]
mixed [Complex a]
rs [Complex a]
ys
        | [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = [Complex a]
ys
        | Bool
otherwise = case [Int
p | Int
p <- [Int
2 .. Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys) :: Double))], [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0] of
            [] -> if [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
13 then [Complex a] -> [Complex a] -> [Complex a]
forall {a}. Num a => [a] -> [a] -> [a]
direct [Complex a]
rs [Complex a]
ys else [Complex a] -> [Complex a] -> [Complex a]
forall {a} {a}.
(RealFloat a, Ord a, Num a) =>
[Complex a] -> [Complex a] -> [Complex a]
blue [Complex a]
rs [Complex a]
ys
            Int
p : [Int]
_ ->
                let m :: Int
m = [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
p
                    sub :: [[Complex a]]
sub =
                        [ [Complex a] -> [Complex a] -> [Complex a]
mixed [[Complex a]
rs [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [[Complex a]
ys [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                        | Int
q <- [Int
0 .. Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        ]
                 in [ [Complex a] -> Complex a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[Complex a]]
sub [[Complex a]] -> Int -> [Complex a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
q [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m) Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
* ([Complex a]
rs [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys)) | Int
q <- [Int
0 .. Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                    | Int
k <- [Int
0 .. [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
    split :: [a] -> [a] -> [a]
split [a]
rs [a]
ys
        | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2 = [a] -> [a] -> [a]
forall {a}. Num a => [a] -> [a] -> [a]
direct [a]
rs [a]
ys
        | Bool
otherwise =
            let m :: Int
m = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4
                evenPart :: [a]
evenPart = [a] -> [a] -> [a]
split [[a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [[a]
ys [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                odd1 :: [a]
odd1 = [a] -> [a] -> [a]
split [[a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [[a]
ys [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                odd3 :: [a]
odd3 = [a] -> [a] -> [a]
split [[a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [[a]
ys [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
             in [ [a]
evenPart [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m))
                    a -> a -> a
forall a. Num a => a -> a -> a
+ ([a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k) a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
odd1 [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m))
                    a -> a -> a
forall a. Num a => a -> a -> a
+ ([a]
rs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m))) a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
odd3 [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m))
                | Int
k <- [Int
0 .. Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
    -- Bluestein writes jk=(j²+k²-(j-k)²)/2 and uses padded convolution.
    blue :: [Complex a] -> [Complex a] -> [Complex a]
blue [Complex a]
_ [Complex a
y] = [Complex a
y]
    blue [Complex a]
rs [Complex a]
ys =
        let len :: Int
len = [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys
            sign :: a
sign = if Complex a -> a
forall a. Complex a -> a
imagPart ([Complex a]
rs [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 then -a
1 else a
1
            chirp :: a -> Complex a
chirp a
j = a -> Complex a
forall a. Floating a => a -> Complex a
cis (a
sign a -> a -> a
forall a. Num a => a -> a -> a
* a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
j a -> a -> a
forall a. Num a => a -> a -> a
* a
j) a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
            m :: Int
m = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) Int
1
            a :: [Complex a]
a = (Complex a -> Complex a -> Complex a)
-> [Complex a] -> [Complex a] -> [Complex a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
(*) [Complex a]
ys [Int -> Complex a
forall {a}. Integral a => a -> Complex a
chirp Int
j | Int
j <- [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [Complex a] -> [Complex a] -> [Complex a]
forall a. [a] -> [a] -> [a]
++ Int -> Complex a -> [Complex a]
forall a. Int -> a -> [a]
replicate (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len) Complex a
0
            b :: [Complex a]
b =
                [ if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len then Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate (Int -> Complex a
forall {a}. Integral a => a -> Complex a
chirp Int
j) else if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len then Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate (Int -> Complex a
forall {a}. Integral a => a -> Complex a
chirp (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j)) else Complex a
0
                | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
            conv :: [Complex a]
conv = [Complex a] -> [Complex a] -> [Complex a]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
convolutionRaw [Complex a]
a [Complex a]
b
         in [[Complex a]
conv [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
* Int -> Complex a
forall {a}. Integral a => a -> Complex a
chirp Int
k | Int
k <- [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    convolutionRaw :: [Complex a] -> [Complex a] -> [Complex a]
convolutionRaw [Complex a]
a [Complex a]
b =
        let m :: Int
m = [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
a
            rootsF :: [Complex a]
rootsF = [a -> Complex a
forall a. Floating a => a -> Complex a
cis (-a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            fa :: [Complex a]
fa = Int -> [Complex a] -> [Complex a] -> [Complex a]
forall {a}. Num a => Int -> [a] -> [a] -> [a]
radix Int
2 [Complex a]
rootsF [Complex a]
a
            fb :: [Complex a]
fb = Int -> [Complex a] -> [Complex a] -> [Complex a]
forall {a}. Num a => Int -> [a] -> [a] -> [a]
radix Int
2 [Complex a]
rootsF [Complex a]
b
         in (Complex a -> Complex a) -> [Complex a] -> [Complex a]
forall a b. (a -> b) -> [a] -> [b]
map (Complex a -> Complex a -> Complex a
forall a. Fractional a => a -> a -> a
/ (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m a -> a -> Complex a
forall a. a -> a -> Complex a
:+ a
0)) (Int -> [Complex a] -> [Complex a] -> [Complex a]
forall {a}. Num a => Int -> [a] -> [a] -> [a]
radix Int
2 ((Complex a -> Complex a) -> [Complex a] -> [Complex a]
forall a b. (a -> b) -> [a] -> [b]
map Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate [Complex a]
rootsF) ((Complex a -> Complex a -> Complex a)
-> [Complex a] -> [Complex a] -> [Complex a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
(*) [Complex a]
fa [Complex a]
fb))
    -- Primitive-root permutation reduces a prime DFT to cyclic correlation.
    rader :: [Complex a] -> [Complex a] -> [Complex a]
rader [Complex a]
rs [Complex a]
ys
        | [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 = [Complex a] -> [Complex a] -> [Complex a]
forall {a}. Num a => [a] -> [a] -> [a]
direct [Complex a]
rs [Complex a]
ys
        | Bool
otherwise =
            let p :: Int
p = [Complex a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex a]
ys
                m :: Int
m = Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                order :: Int -> [Int]
order Int
generator = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
m ((Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (\Int
v -> Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
generator Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p) Int
1)
                g :: Int
g = case [Int
v | Int
v <- [Int
2 .. Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
q -> Int -> Int -> Int -> Int
forall {a} {p}. (Integral a, Integral p) => p -> a -> p -> p
powMod Int
v (Int
m Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
q) Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1) (Int -> [Int]
forall {t}. Integral t => t -> [t]
primeDivisors Int
m)] of
                    Int
v : [Int]
_ -> Int
v
                    [] -> Int
1
                perm :: [Int]
perm = Int -> [Int]
order Int
g
                a :: [Complex a]
a = [[Complex a]
ys [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
perm [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ((-Int
j) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m)) | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                b :: [Complex a]
b = [[Complex a]
rs [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | Int
k <- [Int]
perm]
                rm :: [Complex a]
rm =
                    [ a -> Complex a
forall a. Floating a => a -> Complex a
cis ((if Direction
direction Direction -> Direction -> Bool
forall a. Eq a => a -> a -> Bool
== Direction
Forward then -a
1 else a
1) a -> a -> a
forall a. Num a => a -> a -> a
* a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m)
                    | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
                fa :: [Complex a]
fa = [Complex a] -> [Complex a] -> [Complex a]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
mixed [Complex a]
rm [Complex a]
a
                fb :: [Complex a]
fb = [Complex a] -> [Complex a] -> [Complex a]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
mixed [Complex a]
rm [Complex a]
b
                c :: [Complex a]
c = (Complex a -> Complex a) -> [Complex a] -> [Complex a]
forall a b. (a -> b) -> [a] -> [b]
map (Complex a -> Complex a -> Complex a
forall a. Fractional a => a -> a -> a
/ (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m a -> a -> Complex a
forall a. a -> a -> Complex a
:+ a
0)) ([Complex a] -> [Complex a] -> [Complex a]
forall {a}.
RealFloat a =>
[Complex a] -> [Complex a] -> [Complex a]
mixed ((Complex a -> Complex a) -> [Complex a] -> [Complex a]
forall a b. (a -> b) -> [a] -> [b]
map Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate [Complex a]
rm) ((Complex a -> Complex a -> Complex a)
-> [Complex a] -> [Complex a] -> [Complex a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
(*) [Complex a]
fa [Complex a]
fb))
             in [Complex a] -> Complex a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Complex a]
ys Complex a -> [Complex a] -> [Complex a]
forall a. a -> [a] -> [a]
: [[Complex a]
ys [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
+ [Complex a]
c [Complex a] -> Int -> Complex a
forall a. HasCallStack => [a] -> Int -> a
!! Int -> [Int] -> Int
forall {a}. Eq a => a -> [a] -> Int
index Int
k [Int]
perm | Int
k <- [Int
1 .. Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    primeDivisors :: t -> [t]
primeDivisors t
value = t -> t -> [t]
forall {t}. Integral t => t -> t -> [t]
go t
value t
2
      where
        go :: t -> t -> [t]
go t
1 t
_ = []
        go t
z t
q | t
q t -> t -> t
forall a. Num a => a -> a -> a
* t
q t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
z = [t
z] | t
z t -> t -> t
forall a. Integral a => a -> a -> a
`mod` t
q t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = t
q t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> t -> [t]
go (t -> t -> t
forall a. Integral a => a -> a -> a
strip t
z t
q) (t
q t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) | Bool
otherwise = t -> t -> [t]
go t
z (t
q t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
        strip :: t -> t -> t
strip t
z t
q | t
z t -> t -> t
forall a. Integral a => a -> a -> a
`mod` t
q t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = t -> t -> t
strip (t
z t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
q) t
q | Bool
otherwise = t
z
    index :: a -> [a] -> Int
index a
k [a]
ps = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
k) [a]
ps)
    powMod :: p -> a -> p -> p
powMod p
a a
e p
m = p -> a -> p -> p
forall {a}. Integral a => p -> a -> p -> p
go p
a a
e p
1
      where
        go :: p -> a -> p -> p
go p
_ a
0 p
z = p
z; go p
b a
k p
z = p -> a -> p -> p
go (p
b p -> p -> p
forall a. Num a => a -> a -> a
* p
b p -> p -> p
forall a. Integral a => a -> a -> a
`mod` p
m) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (if a -> Bool
forall a. Integral a => a -> Bool
odd a
k then p
z p -> p -> p
forall a. Num a => a -> a -> a
* p
b p -> p -> p
forall a. Integral a => a -> a -> a
`mod` p
m else p
z)

-- | Automatic forward/inverse FFT and bounded direct forward DFT, respectively.
fft, ifft, dft :: [Complex Double] -> Either MathError [Complex Double]
fft :: [Complex Double] -> Either MathError [Complex Double]
fft [Complex Double]
xs = Direction -> FFTAlgorithm -> Int -> Either MathError FFTPlan
planFFT Direction
Forward FFTAlgorithm
Automatic ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs) Either MathError FFTPlan
-> (FFTPlan -> Either MathError [Complex Double])
-> Either MathError [Complex Double]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
`executeFFT` [Complex Double]
xs)
ifft :: [Complex Double] -> Either MathError [Complex Double]
ifft [Complex Double]
xs = Direction -> FFTAlgorithm -> Int -> Either MathError FFTPlan
planFFT Direction
Inverse FFTAlgorithm
Automatic ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs) Either MathError FFTPlan
-> (FFTPlan -> Either MathError [Complex Double])
-> Either MathError [Complex Double]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
`executeFFT` [Complex Double]
xs)
dft :: [Complex Double] -> Either MathError [Complex Double]
dft [Complex Double]
xs = Direction -> FFTAlgorithm -> Int -> Either MathError FFTPlan
planFFT Direction
Forward FFTAlgorithm
Direct ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs) Either MathError FFTPlan
-> (FFTPlan -> Either MathError [Complex Double])
-> Either MathError [Complex Double]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
`executeFFT` [Complex Double]
xs)

-- | Full complex spectrum; the real-spectrum inverse checks Hermitian symmetry.
realFFT :: [Double] -> Either MathError [Complex Double]
realFFT :: [Double] -> Either MathError [Complex Double]
realFFT = [Complex Double] -> Either MathError [Complex Double]
fft ([Complex Double] -> Either MathError [Complex Double])
-> ([Double] -> [Complex Double])
-> [Double]
-> Either MathError [Complex Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Complex Double) -> [Double] -> [Complex Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0)

-- | Inverse of a full real spectrum; require conjugate symmetry within the supplied absolute tolerance.
inverseRealFFT :: Double -> [Complex Double] -> Either MathError [Double]
inverseRealFFT :: Double -> [Complex Double] -> Either MathError [Double]
inverseRealFFT Double
tolerance [Complex Double]
xs = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
tolerance Bool -> Bool -> Bool
&& Double
tolerance Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Bool -> Bool
not ([Complex Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Complex Double]
xs)) String
"invalid real FFT tolerance/length"
    let n :: Int
n = [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs
    Bool -> String -> Either MathError ()
require
        ([Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude ([Complex Double]
xs [Complex Double] -> Int -> Complex Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
- Complex Double -> Complex Double
forall a. Num a => Complex a -> Complex a
conjugate ([Complex Double]
xs [Complex Double] -> Int -> Complex Double
forall a. HasCallStack => [a] -> Int -> a
!! ((-Int
k) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n))) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
tolerance | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
        String
"spectrum is not Hermitian"
    (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Double
forall a. Complex a -> a
realPart ([Complex Double] -> [Double])
-> Either MathError [Complex Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Complex Double] -> Either MathError [Complex Double]
ifft [Complex Double]
xs

-- | Apply one immutable plan independently to each batch; any inconsistent row length fails.
fftBatch :: FFTPlan -> [[Complex Double]] -> Either MathError [[Complex Double]]
fftBatch :: FFTPlan
-> [[Complex Double]] -> Either MathError [[Complex Double]]
fftBatch FFTPlan
plan = ([Complex Double] -> Either MathError [Complex Double])
-> [[Complex Double]] -> Either MathError [[Complex Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
executeFFT FFTPlan
plan)

-- | Transform all axes. Product budget is 65536; last axis varies fastest.
fftND :: Direction -> [Int] -> [Complex Double] -> Either MathError [Complex Double]
fftND :: Direction
-> [Int] -> [Complex Double] -> Either MathError [Complex Double]
fftND Direction
direction [Int]
dims [Complex Double]
xs = do
    Bool -> String -> Either MathError ()
require
        ( Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
dims)
            Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) [Int]
dims
            Bool -> Bool -> Bool
&& [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
dims) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
65536
            Bool -> Bool -> Bool
&& [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs
        )
        String
"invalid ND FFT shape"
    (Either MathError [Complex Double]
 -> Int -> Either MathError [Complex Double])
-> Either MathError [Complex Double]
-> [Int]
-> Either MathError [Complex Double]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Either MathError [Complex Double]
-> Int -> Either MathError [Complex Double]
axis ([Complex Double] -> Either MathError [Complex Double]
forall a b. b -> Either a b
Right [Complex Double]
xs) [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  where
    axis :: Either MathError [Complex Double]
-> Int -> Either MathError [Complex Double]
axis Either MathError [Complex Double]
result Int
a = do
        values <- Either MathError [Complex Double]
result
        let n = [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
a; stride = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
dims); outer = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
a [Int]
dims)
        p <- planFFT direction Automatic n
        blocks <-
            mapM
                ( \Int
o ->
                    (Int -> Either MathError [Complex Double])
-> [Int] -> Either MathError [[Complex Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                        (\Int
s -> FFTPlan -> [Complex Double] -> Either MathError [Complex Double]
executeFFT FFTPlan
p [[Complex Double]
values [Complex Double] -> Int -> Complex Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s) | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
                        [Int
0 .. Int
stride Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                )
                [0 .. outer - 1]
        Right [blocks !! o !! s !! j | o <- [0 .. outer - 1], j <- [0 .. n - 1], s <- [0 .. stride - 1]]

{- | The caller exclusively owns a valid, aligned, writable buffer of exactly
N Complex Double elements for the duration of this call. No pointers escape;
the entire input is copied and checked before a successful write-back.
-}
fftInPlace :: FFTPlan -> Ptr (Complex Double) -> Int -> IO (Either MathError ())
fftInPlace :: FFTPlan -> Ptr (Complex Double) -> Int -> IO (Either MathError ())
fftInPlace plan :: FFTPlan
plan@(FFTPlan Direction
_ FFTAlgorithm
_ Int
n [Complex Double]
_) Ptr (Complex Double)
pointer Int
count
    | Ptr (Complex Double)
pointer Ptr (Complex Double) -> Ptr (Complex Double) -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr (Complex Double)
forall a. Ptr a
nullPtr Bool -> Bool -> Bool
|| Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n =
        Either MathError () -> IO (Either MathError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid in-place FFT buffer/length"))
    | Bool
otherwise = do
        xs <- Int -> Ptr (Complex Double) -> IO [Complex Double]
forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray Int
count Ptr (Complex Double)
pointer
        case executeFFT plan xs of
            Left MathError
e -> Either MathError () -> IO (Either MathError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
e)
            Right [Complex Double]
ys -> Ptr (Complex Double) -> [Complex Double] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr (Complex Double)
pointer [Complex Double]
ys IO () -> IO (Either MathError ()) -> IO (Either MathError ())
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Either MathError () -> IO (Either MathError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either MathError ()
forall a b. b -> Either a b
Right ())

-- | Signed DFT frequency bins in cycles per coordinate unit, given sample count and positive sample spacing.
frequencies :: Int -> Double -> Either MathError [Double]
frequencies :: Int -> Double -> Either MathError [Double]
frequencies Int
n Double
spacing = do
    Bool -> String -> Either MathError ()
require (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
spacing Bool -> Bool -> Bool
&& Double
spacing Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) String
"invalid sampling spacing"
    let period :: Double
period = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
spacing
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
period) String
"sampling period overflow"
    let result :: [Double]
result = [Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 then Int
k else Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
period | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    Bool -> String -> Either MathError ()
require ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
result) String
"frequency bins overflow"
    [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Double]
result

-- | Rotate zero frequency to/from the center; these operations are exact inverses for odd and even lengths.
fftShift, ifftShift :: [a] -> [a]
fftShift :: forall a. [a] -> [a]
fftShift [a]
xs = let k :: Int
k = ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 in Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [a]
xs
ifftShift :: forall a. [a] -> [a]
ifftShift [a]
xs = let k :: Int
k = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 in Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [a]
xs

realTransform :: Bool -> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform :: Bool
-> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform Bool
sine Bool
inverse TransformType
kind [Double]
xs = do
    let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
        nd :: Double
nd = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
        endpoint :: Int -> a
endpoint Int
j = if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 then a
1 a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
sqrt a
2 else a
1
        base :: Int -> Int -> Double
base Int
j Int
k =
            if Bool
sine
                then case TransformType
kind of
                    TransformType
TypeI -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nd Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nd Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1))
                    TransformType
TypeII ->
                        Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 else Double
1)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                    TransformType
TypeIII ->
                        Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 else Double
1)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                    TransformType
TypeIV -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                else case TransformType
kind of
                    TransformType
TypeI -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1)) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall {a}. Floating a => Int -> a
endpoint Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall {a}. Floating a => Int -> a
endpoint Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1))
                    TransformType
TypeII ->
                        Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 else Double
1)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                    TransformType
TypeIII ->
                        Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 else Double
1)
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
                    TransformType
TypeIV -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd)
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096 Bool -> Bool -> Bool
&& (Bool
sine Bool -> Bool -> Bool
|| TransformType
kind TransformType -> TransformType -> Bool
forall a. Eq a => a -> a -> Bool
/= TransformType
TypeI Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs)
        String
"invalid cosine/sine transform input"
    let ys :: [Double]
ys =
            [[Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Bool
inverse then Int -> Int -> Double
base Int
k Int
j else Int -> Int -> Double
base Int
j Int
k) | (Int
j, Double
x) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
xs] | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    Bool -> String -> Either MathError ()
require ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
ys) String
"cosine/sine transform overflow"
    [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Double]
ys

-- | Apply an orthonormal cosine/sine basis or its transpose; type-I cosine needs at least two samples.
dct, idct, dst, idst :: TransformType -> [Double] -> Either MathError [Double]
dct :: TransformType -> [Double] -> Either MathError [Double]
dct = Bool
-> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform Bool
False Bool
False
idct :: TransformType -> [Double] -> Either MathError [Double]
idct = Bool
-> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform Bool
False Bool
True
dst :: TransformType -> [Double] -> Either MathError [Double]
dst = Bool
-> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform Bool
True Bool
False
idst :: TransformType -> [Double] -> Either MathError [Double]
idst = Bool
-> Bool -> TransformType -> [Double] -> Either MathError [Double]
realTransform Bool
True Bool
True

-- | Unnormalized real Hartley transform, Re(DFT)-Im(DFT); applying twice multiplies by N.
hartley :: [Double] -> Either MathError [Double]
hartley :: [Double] -> Either MathError [Double]
hartley [Double]
xs = (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Complex Double
z -> Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
z) ([Complex Double] -> [Double])
-> Either MathError [Complex Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double] -> Either MathError [Complex Double]
realFFT [Double]
xs

{- | Frames include zero padding on both sides by one full window, so windows
with zero endpoints reconstruct boundary samples. Inverse divides by summed
window squares and rejects uncovered samples. No implicit amplitude scaling.
-}
data STFT = STFT Int Int [Double] [[Complex Double]] deriving (STFT -> STFT -> Bool
(STFT -> STFT -> Bool) -> (STFT -> STFT -> Bool) -> Eq STFT
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: STFT -> STFT -> Bool
== :: STFT -> STFT -> Bool
$c/= :: STFT -> STFT -> Bool
/= :: STFT -> STFT -> Bool
Eq, Int -> STFT -> ShowS
[STFT] -> ShowS
STFT -> String
(Int -> STFT -> ShowS)
-> (STFT -> String) -> ([STFT] -> ShowS) -> Show STFT
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> STFT -> ShowS
showsPrec :: Int -> STFT -> ShowS
$cshow :: STFT -> String
show :: STFT -> String
$cshowList :: [STFT] -> ShowS
showList :: [STFT] -> ShowS
Show)

-- | Read the full complex spectra in increasing frame-start order.
stftFrames :: STFT -> [[Complex Double]]
stftFrames :: STFT -> [[Complex Double]]
stftFrames (STFT Int
_ Int
_ [Double]
_ [[Complex Double]]
frames) = [[Complex Double]]
frames

-- | Windowed analysis with an explicit hop in samples, padded endpoint frames, and a 1048576-coefficient storage budget.
stft :: [Double] -> Int -> [Double] -> Either MathError STFT
stft :: [Double] -> Int -> [Double] -> Either MathError STFT
stft [Double]
window Int
hop [Double]
xs = do
    let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window
        padded :: [Double]
padded = Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0 [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
xs [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0
        starts :: [Int]
starts = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\Int
k -> Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
padded) [Int
0, Int
hop ..]
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536 Bool -> Bool -> Bool
&& Int
hop Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
hop Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite ([Double]
window [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
xs) Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536)
        String
"invalid STFT/window/hop"
    Bool -> String -> Either MathError ()
require
        (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* ((Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hop Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1048576)
        String
"STFT frame storage budget exceeded"
    frames <- (Int -> Either MathError [Complex Double])
-> [Int] -> Either MathError [[Complex Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Double] -> Either MathError [Complex Double]
realFFT ([Double] -> Either MathError [Complex Double])
-> (Int -> [Double]) -> Int -> Either MathError [Complex Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
window ([Double] -> [Double]) -> (Int -> [Double]) -> Int -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\Int
k -> Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
k [Double]
padded))) [Int]
starts
    Right (STFT (length xs) hop window frames)

-- | Overlap-add inverse with squared-window normalization; reject every uncovered original sample.
istft :: STFT -> Either MathError [Double]
istft :: STFT -> Either MathError [Double]
istft (STFT Int
count Int
hop [Double]
window [[Complex Double]]
frames) = do
    time <- ([Complex Double] -> Either MathError [Complex Double])
-> [[Complex Double]] -> Either MathError [[Complex Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM [Complex Double] -> Either MathError [Complex Double]
ifft [[Complex Double]]
frames
    let n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window
        at Int
k =
            [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                [ Complex Double -> Double
forall a. Complex a -> a
realPart ([Complex Double]
row [Complex Double] -> Int -> Complex Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
window [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) | (Int
s, [Complex Double]
row) <- [Int] -> [[Complex Double]] -> [(Int, [Complex Double])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0, Int
hop ..] [[Complex Double]]
time, let j :: Int
j = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s, Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0, Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
                ]
        weight Int
k =
            [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                [([Double]
window [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Int
s <- Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([[Complex Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Complex Double]]
frames) [Int
0, Int
hop ..], let j :: Int
j = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s, Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0, Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n]
        ws = (Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
weight [Int
n .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    require (all (> 1e-300) ws) "STFT overlap leaves uncovered samples"
    let ys = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Fractional a => a -> a -> a
(/) ((Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
at [Int
n .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) [Double]
ws
    require (all finite ys) "STFT reconstruction overflow"
    Right ys

-- | Elementwise complex magnitude, principal phase in radians, and squared magnitude; no PSD normalization.
magnitudeSpectrum, phaseSpectrum, powerSpectrum :: [Complex Double] -> [Double]
magnitudeSpectrum :: [Complex Double] -> [Double]
magnitudeSpectrum = (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude
phaseSpectrum :: [Complex Double] -> [Double]
phaseSpectrum = (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Double
forall a. RealFloat a => Complex a -> a
phase
powerSpectrum :: [Complex Double] -> [Double]
powerSpectrum = (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Complex Double
z -> Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude Complex Double
z Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int))

-- | Two-sided PSD in amplitude-squared/Hz: squared FFT magnitude divided by sample rate and window squared energy.
periodogram :: Double -> [Double] -> [Double] -> Either MathError [Double]
periodogram :: Double -> [Double] -> [Double] -> Either MathError [Double]
periodogram Double
rate [Double]
window [Double]
xs = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
finite Double
rate Bool -> Bool -> Bool
&& Double
rate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
window)
        String
"invalid periodogram input"
    let energy :: Double
energy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
window)
    Bool -> String -> Either MathError ()
require
        (Double
energy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
energy Bool -> Bool -> Bool
&& Double -> Bool
finite (Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy) Bool -> Bool -> Bool
&& Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0)
        String
"zero/overflow window energy or PSD scale"
    spectrum <- [Double] -> Either MathError [Complex Double]
realFFT ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
window [Double]
xs)
    let ys = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy)) ([Complex Double] -> [Double]
powerSpectrum [Complex Double]
spectrum)
    require (all finite ys) "PSD overflow"
    Right ys

-- | Average full-window two-sided periodograms at the supplied hop; discard the incomplete final frame.
welch :: Double -> [Double] -> Int -> [Double] -> Either MathError [Double]
welch :: Double -> [Double] -> Int -> [Double] -> Either MathError [Double]
welch Double
rate [Double]
window Int
hop [Double]
xs = do
    Bool -> String -> Either MathError ()
require
        ( Int
hop Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
window)
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1048576
            Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hop Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1048576
        )
        String
"invalid Welch frames"
    let frames :: [[Double]]
frames =
            ([Double] -> Bool) -> [[Double]] -> [[Double]]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile
                ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length)
                [Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
k [Double]
xs) | Int
k <- [Int
0, Int
hop .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    spectra <- ([Double] -> Either MathError [Double])
-> [[Double]] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Double -> [Double] -> [Double] -> Either MathError [Double]
periodogram Double
rate [Double]
window) [[Double]]
frames
    Right (map (/ fromIntegral (length spectra)) (foldl1 (zipWith (+)) spectra))

-- | Two-sided conjugate(X)*Y cross PSD using the same sample-rate/window-energy normalization as periodogram.
crossSpectrum :: Double -> [Double] -> [Double] -> [Double] -> Either MathError [Complex Double]
crossSpectrum :: Double
-> [Double]
-> [Double]
-> [Double]
-> Either MathError [Complex Double]
crossSpectrum Double
rate [Double]
window [Double]
x [Double]
y = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
finite Double
rate Bool -> Bool -> Bool
&& Double
rate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
y Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
window)
        String
"invalid cross spectrum"
    let energy :: Double
energy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
window)
    Bool -> String -> Either MathError ()
require
        (Double
energy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
energy Bool -> Bool -> Bool
&& Double -> Bool
finite (Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy) Bool -> Bool -> Bool
&& Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0)
        String
"invalid cross-spectrum scale"
    a <- [Double] -> Either MathError [Complex Double]
realFFT ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
window [Double]
x)
    b <- realFFT (zipWith (*) window y)
    checkC (zipWith (\Complex Double
u Complex Double
v -> Complex Double -> Complex Double
forall a. Num a => Complex a -> Complex a
conjugate Complex Double
u Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Complex Double
v Complex Double -> Complex Double -> Complex Double
forall a. Fractional a => a -> a -> a
/ (Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
energy Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0)) a b)

-- | Squared-magnitude STFT frames, retaining the explicit padding/frame order of stft.
spectrogram :: [Double] -> Int -> [Double] -> Either MathError [[Double]]
spectrogram :: [Double] -> Int -> [Double] -> Either MathError [[Double]]
spectrogram [Double]
w Int
h [Double]
xs = ([Complex Double] -> [Double]) -> [[Complex Double]] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map [Complex Double] -> [Double]
powerSpectrum ([[Complex Double]] -> [[Double]])
-> (STFT -> [[Complex Double]]) -> STFT -> [[Double]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STFT -> [[Complex Double]]
stftFrames (STFT -> [[Double]])
-> Either MathError STFT -> Either MathError [[Double]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double] -> Int -> [Double] -> Either MathError STFT
stft [Double]
w Int
h [Double]
xs

-- | Real cepstrum from inverse FFT of log(max(floor,abs(FFT(x)))); floor must be positive.
cepstrum :: Double -> [Double] -> Either MathError [Double]
cepstrum :: Double -> [Double] -> Either MathError [Double]
cepstrum Double
floorValue [Double]
xs = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
floorValue Bool -> Bool -> Bool
&& Double
floorValue Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) String
"invalid cepstral log floor"
    a <- [Double] -> Either MathError [Complex Double]
realFFT [Double]
xs
    map realPart <$> ifft [log (max floorValue (magnitude z)) :+ 0 | z <- a]

-- | Real signal plus imaginary Hilbert transform, formed by doubling positive-frequency bins and retaining DC/Nyquist.
analyticSignal :: [Double] -> Either MathError [Complex Double]
analyticSignal :: [Double] -> Either MathError [Complex Double]
analyticSignal [Double]
xs = do
    spectrum <- [Double] -> Either MathError [Complex Double]
realFFT [Double]
xs
    let n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
        factor Int
k | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int -> Bool
forall a. Integral a => a -> Bool
even Int
n Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 = a
1 | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 = a
2 | Bool
otherwise = a
0
    ifft [z * (factor k :+ 0) | (k, z) <- zip [0 ..] spectrum]

-- | Imaginary part of the analytic signal; periodic finite-record boundary convention.
hilbert :: [Double] -> Either MathError [Double]
hilbert :: [Double] -> Either MathError [Double]
hilbert [Double]
xs = (Complex Double -> Double) -> [Complex Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Double
forall a. Complex a -> a
imagPart ([Complex Double] -> [Double])
-> Either MathError [Complex Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double] -> Either MathError [Complex Double]
analyticSignal [Double]
xs

-- | Direct bounded Z-transform along z_k=a*w^(-k); nonzero a,w required.
chirpZ ::
    Int -> Complex Double -> Complex Double -> [Complex Double] -> Either MathError [Complex Double]
chirpZ :: Int
-> Complex Double
-> Complex Double
-> [Complex Double]
-> Either MathError [Complex Double]
chirpZ Int
count Complex Double
a Complex Double
w [Complex Double]
xs = do
    Bool -> String -> Either MathError ()
require
        (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096 Bool -> Bool -> Bool
&& [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096 Bool -> Bool -> Bool
&& Complex Double
a Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Complex Double
0 Bool -> Bool -> Bool
&& Complex Double
w Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Complex Double
0 Bool -> Bool -> Bool
&& Complex Double -> Bool
finiteC Complex Double
a Bool -> Bool -> Bool
&& Complex Double -> Bool
finiteC Complex Double
w)
        String
"invalid chirp-Z arguments"
    _ <- [Complex Double] -> Either MathError [Complex Double]
checkC [Complex Double]
xs
    checkC [sum [x * a ^^ (-j) * w ^^ (j * k) | (j, x) <- zip [0 ..] xs] | k <- [0 .. count - 1]]

-- | One arbitrary-frequency DFT value from a second-order real resonator; omega is radians/sample.
goertzel :: Double -> [Double] -> Either MathError (Complex Double)
goertzel :: Double -> [Double] -> Either MathError (Complex Double)
goertzel Double
omega [Double]
xs = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
omega Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs)) String
"invalid Goertzel input"
    let c :: Double
c = Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos Double
omega
        (Double
s1, Double
s2) = ((Double, Double) -> Double -> (Double, Double))
-> (Double, Double) -> [Double] -> (Double, Double)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\(Double
a, Double
b) Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b, Double
a)) (Double
0, Double
0) [Double]
xs
        -- Final phase rotates the usual resonator to arbitrary-frequency DFT.
        value :: Complex Double
value = ((Double
s1 Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0) Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
- Double -> Complex Double
forall a. Floating a => a -> Complex a
cis (-Double
omega) Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* (Double
s2 Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0)) Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Double -> Complex Double
forall a. Floating a => a -> Complex a
cis (-Double
omega Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
    _ <- [Complex Double] -> Either MathError [Complex Double]
checkC [Complex Double
value]
    Right value

-- | Full zero-padded complex linear convolution; padded transform length is limited to 65536.
fftConvolve :: [Complex Double] -> [Complex Double] -> Either MathError [Complex Double]
fftConvolve :: [Complex Double]
-> [Complex Double] -> Either MathError [Complex Double]
fftConvolve [Complex Double]
a [Complex Double]
b = do
    Bool -> String -> Either MathError ()
require
        (Bool -> Bool
not ([Complex Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Complex Double]
a) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Complex Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Complex Double]
b) Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
a) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
b) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
65537)
        String
"invalid FFT convolution lengths"
    let count :: Int
count = [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1; n :: Int
n = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
count) (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) Int
1
    Bool -> String -> Either MathError ()
require (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536) String
"FFT convolution exceeds padded budget"
    fa <- [Complex Double] -> Either MathError [Complex Double]
fft ([Complex Double]
a [Complex Double] -> [Complex Double] -> [Complex Double]
forall a. [a] -> [a] -> [a]
++ Int -> Complex Double -> [Complex Double]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Complex Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Complex Double]
a) Complex Double
0)
    fb <- fft (b ++ replicate (n - length b) 0)
    take count <$> ifft (zipWith (*) fa fb)

-- | Correlation lags -(length b-1)..length a-1, conjugating b.
fftCorrelate :: [Complex Double] -> [Complex Double] -> Either MathError [Complex Double]
fftCorrelate :: [Complex Double]
-> [Complex Double] -> Either MathError [Complex Double]
fftCorrelate [Complex Double]
a [Complex Double]
b = [Complex Double]
-> [Complex Double] -> Either MathError [Complex Double]
fftConvolve [Complex Double]
a ([Complex Double] -> [Complex Double]
forall a. [a] -> [a]
reverse ((Complex Double -> Complex Double)
-> [Complex Double] -> [Complex Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Complex Double
forall a. Num a => Complex a -> Complex a
conjugate [Complex Double]
b))