-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Transform.Fourier.FFTUtils
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- Utility functions based on the FFT
--
-----------------------------------------------------------------------------

module Numeric.Transform.Fourier.FFTUtils (
   fft_mag, fft_db, fft_phase, fft_grd, fft_info,
   rfft_mag, rfft_db, rfft_phase, rfft_grd, rfft_info,
   write_fft_info, write_rfft_info,
   ) where

import System.IO
import Data.Array
import Data.Complex

import Numeric.Transform.Fourier.FFT
import DSP.Unwrap

magsq :: RealFloat a => Complex a -> a
magsq :: forall a. RealFloat a => Complex a -> a
magsq (a
x:+a
y) = a
xforall a. Num a => a -> a -> a
*a
x forall a. Num a => a -> a -> a
+ a
yforall a. Num a => a -> a -> a
*a
y

log10 :: (Floating a, Eq a) => a -> a
log10 :: forall a. (Floating a, Eq a) => a -> a
log10 a
0 = -a
1.0e9
log10 a
x = forall a. Floating a => a -> a -> a
logBase a
10 a
x

dot :: RealFloat a => Complex a -> Complex a -> a
dot :: forall a. RealFloat a => Complex a -> Complex a -> a
dot Complex a
a Complex a
b = forall a. Complex a -> a
realPart Complex a
a forall a. Num a => a -> a -> a
* forall a. Complex a -> a
realPart Complex a
b forall a. Num a => a -> a -> a
+ forall a. Complex a -> a
imagPart Complex a
a forall a. Num a => a -> a -> a
* forall a. Complex a -> a
imagPart Complex a
b

eps :: Double
eps :: Double
eps = Double
1.0e-1 :: Double

-- General functions

fft_mag :: (RealFloat b, Integral a, Ix a) =>
           Array a (Complex b) -> Array a b
fft_mag :: forall b a.
(RealFloat b, Integral a, Ix a) =>
Array a (Complex b) -> Array a b
fft_mag Array a (Complex b)
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magnitude forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ Array a (Complex b)
x

fft_db :: (RealFloat b, Integral a, Ix a) =>
          Array a (Complex b) -> Array a b
fft_db :: forall b a.
(RealFloat b, Integral a, Ix a) =>
Array a (Complex b) -> Array a b
fft_db Array a (Complex b)
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b
10 forall a. Num a => a -> a -> a
*) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. (Floating a, Eq a) => a -> a
log10 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magsq forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ Array a (Complex b)
x

fft_phase :: (Integral a, Ix a) =>
             Array a (Complex Double) -> Array a Double
fft_phase :: forall a.
(Integral a, Ix a) =>
Array a (Complex Double) -> Array a Double
fft_phase Array a (Complex Double)
x = forall a b.
(Ix a, Integral a, Ord b, Floating b) =>
b -> Array a b -> Array a b
unwrap Double
eps forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
phase forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ Array a (Complex Double)
x

fft_grd :: (Integral i, RealFloat a, Ix i) =>
           Array i (Complex a) -> Array i a
fft_grd :: forall i a.
(Integral i, RealFloat a, Ix i) =>
Array i (Complex a) -> Array i a
fft_grd Array i (Complex a)
x = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex a)
x') [ forall a. RealFloat a => Complex a -> Complex a -> a
dot (Array i (Complex a)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) (Array i (Complex a)
dx'forall i e. Ix i => Array i e -> i -> e
!i
i) forall a. Fractional a => a -> a -> a
/ forall a. RealFloat a => Complex a -> a
magsq (Array i (Complex a)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex a)
x' ]
    where x' :: Array i (Complex a)
x'  = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft Array i (Complex a)
x
          dx' :: Array i (Complex a)
dx' = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex a)
x) [ forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i forall a. Num a => a -> a -> a
* Array i (Complex a)
xforall i e. Ix i => Array i e -> i -> e
!i
i | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex a)
x ]

fft_info :: (Integral i, Ix i) =>
            Array i (Complex Double)
            -> (Array i Double, Array i Double, Array i Double, Array i Double)
fft_info :: forall i.
(Integral i, Ix i) =>
Array i (Complex Double)
-> (Array i Double, Array i Double, Array i Double, Array i Double)
fft_info Array i (Complex Double)
x = (Array i Double
mag,Array i Double
db,Array i Double
arg,Array i Double
grd)
    where x' :: Array i (Complex Double)
x'  = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft Array i (Complex Double)
x
          dx' :: Array i (Complex Double)
dx' = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a (Complex b) -> Array a (Complex b)
fft forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex Double)
x) [ forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i forall a. Num a => a -> a -> a
* Array i (Complex Double)
xforall i e. Ix i => Array i e -> i -> e
!i
i | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex Double)
x ]
          mag :: Array i Double
mag = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magnitude forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          db :: Array i Double
db  = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Double
10 forall a. Num a => a -> a -> a
*) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. (Floating a, Eq a) => a -> a
log10 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magsq forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          arg :: Array i Double
arg = forall a b.
(Ix a, Integral a, Ord b, Floating b) =>
b -> Array a b -> Array a b
unwrap Double
eps forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
phase forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          grd :: Array i Double
grd = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex Double)
x') [ forall a. RealFloat a => Complex a -> Complex a -> a
dot (Array i (Complex Double)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) (Array i (Complex Double)
dx'forall i e. Ix i => Array i e -> i -> e
!i
i) forall a. Fractional a => a -> a -> a
/ forall a. RealFloat a => Complex a -> a
magsq (Array i (Complex Double)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex Double)
x' ]

rfft_mag :: (RealFloat b, Integral a, Ix a) =>
            Array a b -> Array a b
rfft_mag :: forall b a.
(RealFloat b, Integral a, Ix a) =>
Array a b -> Array a b
rfft_mag Array a b
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magnitude forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft forall a b. (a -> b) -> a -> b
$ Array a b
x

rfft_db :: (RealFloat b, Integral a, Ix a) =>
           Array a b -> Array a b
rfft_db :: forall b a.
(RealFloat b, Integral a, Ix a) =>
Array a b -> Array a b
rfft_db Array a b
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b
10 forall a. Num a => a -> a -> a
*) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. (Floating a, Eq a) => a -> a
log10 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magsq forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft forall a b. (a -> b) -> a -> b
$ Array a b
x

rfft_phase :: (Integral a, Ix a) =>
              Array a Double -> Array a Double
rfft_phase :: forall a. (Integral a, Ix a) => Array a Double -> Array a Double
rfft_phase Array a Double
x = forall a b.
(Ix a, Integral a, Ord b, Floating b) =>
b -> Array a b -> Array a b
unwrap Double
eps forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
phase forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft forall a b. (a -> b) -> a -> b
$ Array a Double
x

rfft_grd :: (Integral i, Ix i, RealFloat a) =>
            Array i a -> Array i a
rfft_grd :: forall i a.
(Integral i, Ix i, RealFloat a) =>
Array i a -> Array i a
rfft_grd Array i a
x = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex a)
x') [ forall a. RealFloat a => Complex a -> Complex a -> a
dot (Array i (Complex a)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) (Array i (Complex a)
dx'forall i e. Ix i => Array i e -> i -> e
!i
i) forall a. Fractional a => a -> a -> a
/ forall a. RealFloat a => Complex a -> a
magsq (Array i (Complex a)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex a)
x' ]
    where x' :: Array i (Complex a)
x'  = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft Array i a
x
          dx' :: Array i (Complex a)
dx' = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i a
x) [ forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i forall a. Num a => a -> a -> a
* Array i a
xforall i e. Ix i => Array i e -> i -> e
!i
i | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i a
x ]

-- I/O

rfft_info :: (Integral i, Ix i) =>
             Array i Double
             -> (Array i Double, Array i Double, Array i Double, Array i Double)
rfft_info :: forall i.
(Integral i, Ix i) =>
Array i Double
-> (Array i Double, Array i Double, Array i Double, Array i Double)
rfft_info Array i Double
x = (Array i Double
mag,Array i Double
db,Array i Double
arg,Array i Double
grd)
    where x' :: Array i (Complex Double)
x'  = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft Array i Double
x
          dx' :: Array i (Complex Double)
dx' = forall a b.
(Ix a, Integral a, RealFloat b) =>
Array a b -> Array a (Complex b)
rfft forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i Double
x) [ forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i forall a. Num a => a -> a -> a
* Array i Double
xforall i e. Ix i => Array i e -> i -> e
!i
i | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i Double
x ]
          mag :: Array i Double
mag = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magnitude forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          db :: Array i Double
db  = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Double
10 forall a. Num a => a -> a -> a
*) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. (Floating a, Eq a) => a -> a
log10 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
magsq forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          arg :: Array i Double
arg = forall a b.
(Ix a, Integral a, Ord b, Floating b) =>
b -> Array a b -> Array a b
unwrap Double
eps forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. RealFloat a => Complex a -> a
phase forall a b. (a -> b) -> a -> b
$ Array i (Complex Double)
x'
          grd :: Array i Double
grd = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (forall i e. Array i e -> (i, i)
bounds Array i (Complex Double)
x') [ forall a. RealFloat a => Complex a -> Complex a -> a
dot (Array i (Complex Double)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) (Array i (Complex Double)
dx'forall i e. Ix i => Array i e -> i -> e
!i
i) forall a. Fractional a => a -> a -> a
/ forall a. RealFloat a => Complex a -> a
magsq (Array i (Complex Double)
x'forall i e. Ix i => Array i e -> i -> e
!i
i) | i
i <- forall i e. Ix i => Array i e -> [i]
indices Array i (Complex Double)
x' ]

hPrintIndex :: (Integral a, Integral i, Show b) =>
               Handle -> i -> (a, b) -> IO ()
hPrintIndex :: forall a i b.
(Integral a, Integral i, Show b) =>
Handle -> i -> (a, b) -> IO ()
hPrintIndex Handle
h i
n (a
i,b
x) = do
   Handle -> String -> IO ()
hPutStr   Handle
h forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral i
n :: Double)
   Handle -> String -> IO ()
hPutStr   Handle
h forall a b. (a -> b) -> a -> b
$ String
" "
   Handle -> String -> IO ()
hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show b
x

write_cvector :: (Show e, Integral i, Ix i) =>
                 FilePath -> Array i e -> IO ()
write_cvector :: forall e i.
(Show e, Integral i, Ix i) =>
String -> Array i e -> IO ()
write_cvector String
f Array i e
x = do
   let n :: i
n = (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall i e. Array i e -> (i, i)
bounds Array i e
x) forall a. Num a => a -> a -> a
+ i
1
   forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
f IOMode
WriteMode forall a b. (a -> b) -> a -> b
$ \Handle
h ->
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a i b.
(Integral a, Integral i, Show b) =>
Handle -> i -> (a, b) -> IO ()
hPrintIndex Handle
h i
n) forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => Array i e -> [(i, e)]
assocs Array i e
x

write_fft_info :: (Ix i, Integral i) =>
                  String -> Array i (Complex Double) -> IO ()
write_fft_info :: forall i.
(Ix i, Integral i) =>
String -> Array i (Complex Double) -> IO ()
write_fft_info String
b Array i (Complex Double)
x = do
   let (Array i Double
mag,Array i Double
db,Array i Double
arg,Array i Double
grd) = forall i.
(Integral i, Ix i) =>
Array i (Complex Double)
-> (Array i Double, Array i Double, Array i Double, Array i Double)
fft_info Array i (Complex Double)
x
   forall e i.
(Show e, Integral i, Ix i) =>
String -> Array i e -> IO ()
write_cvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_mag.out") Array i Double
mag
   forall e i.
(Show e, Integral i, Ix i) =>
String -> Array i e -> IO ()
write_cvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_db.out")  Array i Double
db
   forall e i.
(Show e, Integral i, Ix i) =>
String -> Array i e -> IO ()
write_cvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_arg.out") Array i Double
arg
   forall e i.
(Show e, Integral i, Ix i) =>
String -> Array i e -> IO ()
write_cvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_grd.out") Array i Double
grd

write_rvector :: Show e => FilePath -> Array Int e -> IO ()
write_rvector :: forall e. Show e => String -> Array Int e -> IO ()
write_rvector String
f Array Int e
x = do
   let n :: Int
n = (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall i e. Array i e -> (i, i)
bounds Array Int e
x) forall a. Num a => a -> a -> a
+ Int
1
   forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
f IOMode
WriteMode forall a b. (a -> b) -> a -> b
$ \Handle
h ->
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a i b.
(Integral a, Integral i, Show b) =>
Handle -> i -> (a, b) -> IO ()
hPrintIndex Handle
h Int
n) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (Int
n forall a. Integral a => a -> a -> a
`div` Int
2) forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => Array i e -> [(i, e)]
assocs Array Int e
x

write_rfft_info :: String -> Array Int Double -> IO ()
write_rfft_info :: String -> Array Int Double -> IO ()
write_rfft_info String
b Array Int Double
x = do
   let (Array Int Double
mag,Array Int Double
db,Array Int Double
arg,Array Int Double
grd) = forall i.
(Integral i, Ix i) =>
Array i Double
-> (Array i Double, Array i Double, Array i Double, Array i Double)
rfft_info Array Int Double
x
   forall e. Show e => String -> Array Int e -> IO ()
write_rvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_mag.out") Array Int Double
mag
   forall e. Show e => String -> Array Int e -> IO ()
write_rvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_db.out")  Array Int Double
db
   forall e. Show e => String -> Array Int e -> IO ()
write_rvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_arg.out") Array Int Double
arg
   forall e. Show e => String -> Array Int e -> IO ()
write_rvector (String
b forall a. [a] -> [a] -> [a]
++ String
"_grd.out") Array Int Double
grd