{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
module Numeric.HHT (
hhtEmd
, hht
, hhtSpectrum
, marginal, instantaneousEnergy, degreeOfStationarity
, HHT(..), HHTLine(..)
, EMDOpts(..), defaultEO, BoundaryHandler(..), SiftCondition(..), defaultSC, SplineEnd(..)
, hilbert
, hilbertIm
, hilbertMagFreq
) where
import Data.Complex
import Data.Finite
import Data.Fixed
import Data.List
import Data.Proxy
import Data.Semigroup
import GHC.Generics (Generic)
import GHC.TypeNats
import Numeric.EMD
import qualified Data.Binary as Bi
import qualified Data.Map as M
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Sized as SVG
import qualified Data.Vector.Sized as SV
data HHTLine v n a = HHTLine
{
hlMags :: !(SVG.Vector v n a)
, hlFreqs :: !(SVG.Vector v n a)
}
deriving (Show, Eq, Ord, Generic)
instance (VG.Vector v a, KnownNat n, Bi.Binary (v a)) => Bi.Binary (HHTLine v n a) where
put HHTLine{..} = Bi.put (SVG.fromSized hlMags )
*> Bi.put (SVG.fromSized hlFreqs)
get = do
Just hlMags <- SVG.toSized <$> Bi.get
Just hlFreqs <- SVG.toSized <$> Bi.get
pure HHTLine{..}
newtype HHT v n a = HHT { hhtLines :: [HHTLine v n a] }
deriving (Show, Eq, Ord, Generic)
instance (VG.Vector v a, KnownNat n, Bi.Binary (v a)) => Bi.Binary (HHT v n a)
hht :: forall v n a. (VG.Vector v a, KnownNat n, RealFloat a)
=> EMDOpts a
-> SVG.Vector v (n + 1) a
-> HHT v n a
hht eo = hhtEmd . emd eo
hhtEmd
:: forall v n a. (VG.Vector v a, KnownNat n, RealFloat a)
=> EMD v (n + 1) a
-> HHT v n a
hhtEmd EMD{..} = HHT $ map go emdIMFs
where
go i = HHTLine (SVG.init m) f
where
(m, f) = hilbertMagFreq i
hhtSpectrum
:: forall n a k. (KnownNat n, Ord k, Num a)
=> (a -> k)
-> HHT V.Vector n a
-> SV.Vector n (M.Map k a)
hhtSpectrum f = foldl' ((SV.zipWith . M.unionWith) (+)) (pure mempty) . map go . hhtLines
where
go :: HHTLine V.Vector n a -> SV.Vector n (M.Map k a)
go HHTLine{..} = SV.generate $ \i ->
M.singleton (f $ hlFreqs `SVG.index` i) (hlMags `SVG.index` i)
marginal
:: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
=> (a -> k)
-> HHT v n a
-> M.Map k a
marginal f = M.unionsWith (+) . concatMap go . hhtLines
where
go :: HHTLine v n a -> [M.Map k a]
go HHTLine{..} = flip fmap (finites @n) $ \i ->
M.singleton (f $ hlFreqs `SVG.index` i) (hlMags `SVG.index` i)
instantaneousEnergy
:: forall v n a. (VG.Vector v a, KnownNat n, Num a)
=> HHT v n a
-> SVG.Vector v n a
instantaneousEnergy = sum . map (SVG.map (^ (2 :: Int)) . hlMags) . hhtLines
degreeOfStationarity
:: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Fractional a)
=> (a -> k)
-> HHT v n a
-> M.Map k a
degreeOfStationarity f h = M.unionsWith (+)
. concatMap go
. hhtLines
$ h
where
meanMarg = (/ fromIntegral (natVal (Proxy @n))) <$> marginal f h
go :: HHTLine v n a -> [M.Map k a]
go HHTLine{..} = flip fmap (finites @n) $ \i ->
let fr = f $ hlFreqs `SVG.index` i
in M.singleton fr $
(1 - (hlMags `SVG.index` i / meanMarg M.! fr)) ^ (2 :: Int)
hilbertMagFreq
:: forall v n a. (VG.Vector v a, KnownNat n, RealFloat a)
=> SVG.Vector v (n + 1) a
-> (SVG.Vector v (n + 1) a, SVG.Vector v n a)
hilbertMagFreq v = (hilbertMag, hilbertFreq)
where
v' = hilbertIm v
hilbertMag = SVG.zipWith (\x x' -> magnitude (x :+ x')) v v'
hilbertPhase = SVG.zipWith (\x x' -> phase (x :+ x')) v v'
hilbertFreq = SVG.map ((`mod'` 1) . (/ (2 * pi))) $ SVG.tail hilbertPhase - SVG.init hilbertPhase
hilbert
:: forall v n a. (VG.Vector v a, VG.Vector v (Complex a), KnownNat n, Floating a)
=> SVG.Vector v n a
-> SVG.Vector v n (Complex a)
hilbert v = SVG.zipWith (:+) v (hilbertIm v)
hilbertIm
:: forall v n a. (VG.Vector v a, KnownNat n, Floating a)
=> SVG.Vector v n a
-> SVG.Vector v n a
hilbertIm v = SVG.generate $ \i -> getSum . foldMap (Sum . go i) $ finites @n
where
go :: Finite n -> Finite n -> a
go i j
| even k = 0
| otherwise = 2 * (v `SVG.index` j) / pi / fromIntegral k
where
k :: Int
k = fromIntegral i - fromIntegral j