{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wall #-}
#ifdef TESTING
{-# OPTIONS_GHC -fno-warn-unused-binds #-}
#endif
module ConCat.FFT
( dft, FFT(..), DFTTy, genericFft, GFFT
, twiddle, twiddles, omega, cis
, o8sq
) where
import Prelude hiding (zipWith)
import Data.Complex
import GHC.Generics hiding (C,S)
import Data.Key (Zip(..))
import Data.Pointed
#ifdef TESTING
import Data.Foldable (toList)
import Test.QuickCheck.All (quickCheckAll)
import ShapedTypes.ApproxEq
#endif
import ConCat.Misc (Unop,transpose,C,inGeneric1,inComp)
import ConCat.Sized
import ConCat.Scan (LScan,lproducts)
import ConCat.Pair
import ConCat.Free.LinearRow (($*))
type AS h = (Pointed h, Zip h, LScan h)
type ASZ h = (AS h, Sized h)
dft :: forall f a. (ASZ f, Foldable f, RealFloat a) => Unop (f (Complex a))
dft :: forall (f :: * -> *) a.
(ASZ f, Foldable f, RealFloat a) =>
Unop (f (Complex a))
dft f (Complex a)
xs = Int -> f (f (Complex a))
forall (f :: * -> *) (g :: * -> *) a.
(AS f, AS g, RealFloat a) =>
Int -> g (f (Complex a))
omegas (forall (f :: * -> *). Sized f => Int
size @f) f (f (Complex a)) -> f (Complex a) -> f (Complex a)
forall s (a :: * -> *) (b :: * -> *).
(Zip a, Foldable a, Functor b, Num s) =>
(:-*) a b s -> a s -> b s
$* f (Complex a)
xs
{-# INLINE dft #-}
omegas :: (AS f, AS g, RealFloat a) => Int -> g (f (Complex a))
omegas :: forall (f :: * -> *) (g :: * -> *) a.
(AS f, AS g, RealFloat a) =>
Int -> g (f (Complex a))
omegas = (Complex a -> f (Complex a)) -> g (Complex a) -> g (f (Complex a))
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Complex a -> f (Complex a)
forall (f :: * -> *) a. (LScan f, Pointed f, Num a) => a -> f a
powers (g (Complex a) -> g (f (Complex a)))
-> (Int -> g (Complex a)) -> Int -> g (f (Complex a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Complex a -> g (Complex a)
forall (f :: * -> *) a. (LScan f, Pointed f, Num a) => a -> f a
powers (Complex a -> g (Complex a))
-> (Int -> Complex a) -> Int -> g (Complex a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Complex a
forall a. RealFloat a => Int -> Complex a
omega
{-# INLINE omegas #-}
omegas' :: forall f g a. (ASZ f, ASZ g, RealFloat a) => g (f (Complex a))
omegas' :: forall (f :: * -> *) (g :: * -> *) a.
(ASZ f, ASZ g, RealFloat a) =>
g (f (Complex a))
omegas' = (Complex a -> f (Complex a)) -> g (Complex a) -> g (f (Complex a))
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Complex a -> f (Complex a)
forall (f :: * -> *) a. (LScan f, Pointed f, Num a) => a -> f a
powers (Complex a -> g (Complex a)
forall (f :: * -> *) a. (LScan f, Pointed f, Num a) => a -> f a
powers (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. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (f :: * -> *). Sized f => Int
size @(g :.: f)))))
omega :: RealFloat a => Int -> Complex a
omega :: forall a. RealFloat a => Int -> Complex a
omega Int
n = 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. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
{-# INLINE omega #-}
type DFTTy f = forall a. RealFloat a => f (Complex a) -> FFO f (Complex a)
class FFT f where
type FFO f :: * -> *
fft :: DFTTy f
fftDummy :: f a
fftDummy = f a
forall a. HasCallStack => a
undefined
twiddle :: forall g f a. (ASZ g, ASZ f, RealFloat a) => Unop (g (f (Complex a)))
twiddle :: forall (g :: * -> *) (f :: * -> *) a.
(ASZ g, ASZ f, RealFloat a) =>
Unop (g (f (Complex a)))
twiddle = ((f (Complex a) -> f (Complex a) -> f (Complex a))
-> g (f (Complex a)) -> g (f (Complex a)) -> g (f (Complex a))
forall a b c. (a -> b -> c) -> g a -> g b -> g c
forall (f :: * -> *) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith((f (Complex a) -> f (Complex a) -> f (Complex a))
-> g (f (Complex a)) -> g (f (Complex a)) -> g (f (Complex a)))
-> ((Complex a -> Complex a -> Complex a)
-> f (Complex a) -> f (Complex a) -> f (Complex a))
-> (Complex a -> Complex a -> Complex a)
-> g (f (Complex a))
-> g (f (Complex a))
-> g (f (Complex a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Complex a -> Complex a -> Complex a)
-> f (Complex a) -> f (Complex a) -> f (Complex a)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith) Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
(*) g (f (Complex a))
forall (f :: * -> *) (g :: * -> *) a.
(ASZ f, ASZ g, RealFloat a) =>
g (f (Complex a))
omegas'
{-# INLINE twiddle #-}
twiddles :: forall g f a. (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))
twiddles :: forall (g :: * -> *) (f :: * -> *) a.
(ASZ g, ASZ f, RealFloat a) =>
g (f (Complex a))
twiddles = Int -> g (f (Complex a))
forall (f :: * -> *) (g :: * -> *) a.
(AS f, AS g, RealFloat a) =>
Int -> g (f (Complex a))
omegas (forall (f :: * -> *). Sized f => Int
size @(g :.: f))
{-# INLINE twiddles #-}
o8sq :: C
o8sq :: C
o8sq = Int -> C
forall a. RealFloat a => Int -> Complex a
omega (Int
8 :: Int) C -> Int -> C
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)
powers :: (LScan f, Pointed f, Num a) => a -> f a
powers :: forall (f :: * -> *) a. (LScan f, Pointed f, Num a) => a -> f a
powers = (f a, a) -> f a
forall a b. (a, b) -> a
fst ((f a, a) -> f a) -> (a -> (f a, a)) -> a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> (f a, a)
forall (f :: * -> *) a. (LScan f, Num a) => f a -> f a :* a
lproducts (f a -> (f a, a)) -> (a -> f a) -> a -> (f a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> f a
forall a. a -> f a
forall (p :: * -> *) a. Pointed p => a -> p a
point
{-# INLINE powers #-}
instance FFT Par1 where
type FFO Par1 = Par1
fft :: DFTTy Par1
fft = Par1 (Complex a) -> Par1 (Complex a)
Par1 (Complex a) -> FFO Par1 (Complex a)
forall a. a -> a
id
#if 0
inTranspose :: (Traversable f', Traversable g, Applicative g', Applicative f)
=> (f (g a) -> f' (g' a)) -> g (f a) -> g' (f' a)
inTranspose = transpose <-- transpose
ffts' :: ( FFT g, Traversable f, Traversable g
, Applicative (FFO g), Applicative f, RealFloat a) =>
g (f (Complex a)) -> FFO g (f (Complex a))
ffts' = transpose . fmap fft . transpose
#endif
#if 0
transpose :: g (f C) -> f (g C)
fmap fft :: f (g C) -> f (FFO g C)
transpose :: f (FFO g C) -> FFO g (f C)
#endif
instance ( Zip f, Traversable f , Traversable g
, Applicative f, Applicative (FFO f), Applicative (FFO g), Zip (FFO g)
, Pointed f, Traversable (FFO g), Pointed (FFO g)
, FFT f, FFT g, LScan f, LScan (FFO g), Sized f, Sized (FFO g) )
=> FFT (g :.: f) where
type FFO (g :.: f) = FFO f :.: FFO g
fft :: DFTTy (g :.: f)
fft = (g (f (Complex a)) -> FFO f (FFO g (Complex a)))
-> (:.:) g f (Complex a) -> (:.:) (FFO f) (FFO g) (Complex a)
forall {k4} {k5} {k6} {k7} (g :: k4 -> *) (f :: k5 -> k4) (a :: k5)
(g' :: k6 -> *) (f' :: k7 -> k6) (a' :: k7).
(g (f a) -> g' (f' a')) -> (:.:) g f a -> (:.:) g' f' a'
inComp ((f (Complex a) -> FFO f (Complex a))
-> FFO g (f (Complex a)) -> FFO f (FFO g (Complex a))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> FFO g a -> f (FFO g b)
traverse f (Complex a) -> FFO f (Complex a)
DFTTy f
forall (f :: * -> *). FFT f => DFTTy f
fft (FFO g (f (Complex a)) -> FFO f (FFO g (Complex a)))
-> (g (f (Complex a)) -> FFO g (f (Complex a)))
-> g (f (Complex a))
-> FFO f (FFO g (Complex a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Unop (FFO g (f (Complex a)))
forall (g :: * -> *) (f :: * -> *) a.
(ASZ g, ASZ f, RealFloat a) =>
Unop (g (f (Complex a)))
twiddle Unop (FFO g (f (Complex a)))
-> (g (f (Complex a)) -> FFO g (f (Complex a)))
-> g (f (Complex a))
-> FFO g (f (Complex a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (g (Complex a) -> FFO g (Complex a))
-> f (g (Complex a)) -> FFO g (f (Complex a))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f b)
traverse g (Complex a) -> FFO g (Complex a)
DFTTy g
forall (f :: * -> *). FFT f => DFTTy f
fft (f (g (Complex a)) -> FFO g (f (Complex a)))
-> (g (f (Complex a)) -> f (g (Complex a)))
-> g (f (Complex a))
-> FFO g (f (Complex a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g (f (Complex a)) -> f (g (Complex a))
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
transpose)
{-# INLINE fft #-}
#if 0
fft = Comp1 . transpose . fmap fft . twiddle . transpose . fmap fft . transpose . unComp1
fft = Comp1 . traverse fft . twiddle . traverse fft . transpose . unComp1
unComp1 :: (g :. f) a -> g (f a)
transpose :: g (f a) -> f (g a)
fmap fft :: f (g a) -> f (g' a)
transpose :: f (g' a) -> g' (f a)
twiddle :: g' (f a) -> g' (f a)
fmap fft :: g' (f a) -> g' (f' a)
transpose :: g' (f' a) -> f' (g' a)
Comp1 :: f' (g' a) -> (f' :. g') a
#endif
#if 0
ffts' :: g (f C) -> FFO g (f C)
twiddle :: FFO g (f C) -> FFO g (f C)
transpose :: FFO g (f C) -> f (FFO g C)
ffts' :: f (FFO g C) -> FFO f (FFO g C)
#endif
genericFft :: ( Generic1 f, Generic1 (FFO f)
, FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f) ) => DFTTy f
genericFft :: forall (f :: * -> *).
(Generic1 f, Generic1 (FFO f), FFT (Rep1 f),
FFO (Rep1 f) ~ Rep1 (FFO f)) =>
DFTTy f
genericFft = (Rep1 f (Complex a) -> Rep1 (FFO f) (Complex a))
-> f (Complex a) -> FFO f (Complex a)
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> *) (a :: k1) (b :: k2).
(Generic1 f, Generic1 g) =>
(Rep1 f a -> Rep1 g b) -> f a -> g b
inGeneric1 Rep1 f (Complex a) -> Rep1 (FFO f) (Complex a)
Rep1 f (Complex a) -> FFO (Rep1 f) (Complex a)
DFTTy (Rep1 f)
forall (f :: * -> *). FFT f => DFTTy f
fft
type GFFT f = (Generic1 f, Generic1 (FFO f), FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f))
#define GenericFFT(f,g) instance GFFT (f) => FFT (f) where { type FFO (f) = (g); fft = genericFft }
instance FFT Pair where
type FFO Pair = Pair
fft :: DFTTy Pair
fft = \ (Complex a
a :# Complex a
b) -> (Complex a
a Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
+ Complex a
b) Complex a -> Complex a -> Pair (Complex a)
forall a. a -> a -> Pair a
:# (Complex a
a Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
- Complex a
b)
{-# INLINE fft #-}
#ifdef TESTING
#if 0
twiddles :: (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))
as :: f C
(<.> as) :: f C -> C
twiddles :: f (f C)
(<.> as) <$> twiddles :: f C
#endif
dftL :: RealFloat a => Unop [Complex a]
dftL xs = [ sum [ x * ok^n | x <- xs | n <- [0 :: Int ..] ]
| k <- [0 .. length xs - 1], let ok = om ^ k ]
where
om = omega (length xs)
fftl :: (FFT f, Foldable (FFO f), RealFloat a) => f (Complex a) -> [Complex a]
fftl = toList . fft
type LTree = L.Pow Pair
type RTree = R.Pow Pair
type LC n = LTree n C
type RC n = RTree n C
p1 :: Pair C
p1 = 1 :# 0
tw1 :: LTree N1 (Pair C)
tw1 = twiddles
tw2 :: RTree N2 (Pair C)
tw2 = twiddles
tw3 :: RTree N2 (RTree N2 C)
tw3 = twiddles
tw3' :: [[C]]
tw3' = toList (toList <$> tw3)
#if 0
t0 :: LC N0
t0 = L.fromList [1]
t1 :: LC N1
t1 = L.fromList [1, 0]
t2s :: [LC N2]
t2s = L.fromList <$>
[ [1, 0, 0, 0]
, [1, 1, 1, 1]
, [1, -1, 1, -1]
, [1, 0, -1, 0]
, [0, 1, 0, -1]
]
tests :: IO ()
tests = do test p1
test t0
test t1
mapM_ test t2s
#endif
infix 4 ===
(===) :: Eq b => (a -> b) -> (a -> b) -> a -> Bool
(f === g) x = f x == g x
infix 4 =~=
(=~=) :: ApproxEq b => (a -> b) -> (a -> b) -> a -> Bool
(f =~= g) x = f x =~ g x
fftIsDftL :: (FFT f, Foldable f, Foldable (FFO f), RealFloat a, ApproxEq a) =>
f (Complex a) -> Bool
fftIsDftL = toList . fft =~= dftL . toList
dftTIsDftL :: (ASZ f, Traversable f, RealFloat a, ApproxEq a) =>
f (Complex a) -> Bool
dftTIsDftL = toList . dftT =~= dftL . toList
dftIsDftL :: (ASZ f, Foldable f, RealFloat a, ApproxEq a) =>
f (Complex a) -> Bool
dftIsDftL = toList . dft =~= dftL . toList
transposeTwiddleCommutes :: (ASZ g, Traversable g, ASZ f, (ApproxEq (f (g C))))
=> g (f C) -> Bool
transposeTwiddleCommutes =
twiddle . transpose =~= transpose . twiddle
prop_transposeTwiddle_L3P :: LTree N3 (Pair C) -> Bool
prop_transposeTwiddle_L3P = transposeTwiddleCommutes
prop_transposeTwiddle_R3P :: RTree N3 (Pair C) -> Bool
prop_transposeTwiddle_R3P = transposeTwiddleCommutes
prop_dftT_p :: Pair C -> Bool
prop_dftT_p = dftTIsDftL
prop_dftT_L3 :: LTree N3 C -> Bool
prop_dftT_L3 = dftTIsDftL
prop_dftT_R3 :: RTree N3 C -> Bool
prop_dftT_R3 = dftTIsDftL
prop_fft_p :: Pair C -> Bool
prop_fft_p = fftIsDftL
prop_fft_L1 :: LTree N1 C -> Bool
prop_fft_L1 = fftIsDftL
prop_fft_L2 :: LTree N2 C -> Bool
prop_fft_L2 = fftIsDftL
prop_fft_L3 :: LTree N3 C -> Bool
prop_fft_L3 = fftIsDftL
prop_fft_L4 :: LTree N4 C -> Bool
prop_fft_L4 = fftIsDftL
prop_fft_R1 :: RTree N1 C -> Bool
prop_fft_R1 = fftIsDftL
prop_fft_R2 :: RTree N2 C -> Bool
prop_fft_R2 = fftIsDftL
prop_fft_R3 :: RTree N3 C -> Bool
prop_fft_R3 = fftIsDftL
prop_fft_R4 :: RTree N4 C -> Bool
prop_fft_R4 = fftIsDftL
return []
runTests :: IO Bool
runTests = $quickCheckAll
#endif