{-# LANGUAGE BangPatterns #-}
{-# 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 (
HHT(..), HHTLine(..)
, hhtEmd
, hht
, hhtSpectrum, hhtSparseSpectrum, hhtDenseSpectrum
, marginal, instantaneousEnergy, degreeOfStationarity
, expectedFreq, dominantFreq
, EMDOpts(..), defaultEO, BoundaryHandler(..), SiftCondition(..), defaultSC, SplineEnd(..)
, hilbert
, hilbertIm
, hilbertPolar
, hilbertMagFreq
) where
import Control.DeepSeq
import Data.Complex
import Data.Finite
import Data.Fixed
import Data.Foldable
import Data.Maybe
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.List.NonEmpty as NE
import qualified Data.Map as M
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{..}
instance NFData (v a) => NFData (HHTLine v n a)
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)
instance NFData (v a) => NFData (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 v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
=> (a -> k)
-> HHT v 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 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)
hhtSparseSpectrum
:: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
=> (a -> k)
-> HHT v n a
-> M.Map (Finite n, k) a
hhtSparseSpectrum f = M.unionsWith (+) . concatMap go . hhtLines
where
go :: HHTLine v n a -> [M.Map (Finite n, k) a]
go HHTLine{..} = flip fmap (finites @n) $ \i ->
M.singleton (i, f $ hlFreqs `SVG.index` i) $
hlMags `SVG.index` i
hhtDenseSpectrum
:: forall v n m a. (VG.Vector v a, KnownNat n, KnownNat m, Num a)
=> (a -> Finite m)
-> HHT v n a
-> SV.Vector n (SV.Vector m a)
hhtDenseSpectrum f h = SV.generate $ \i -> SV.generate $ \j ->
M.findWithDefault 0 (i, j) ss
where
ss = hhtSparseSpectrum f h
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)
expectedFreq
:: forall v n a. (VG.Vector v a, KnownNat n, Fractional a)
=> HHT v n a
-> SVG.Vector v n a
expectedFreq HHT{..} = SVG.generate $ \i -> weightedAverage . map (go i) $ hhtLines
where
go :: Finite n -> HHTLine v n a -> (a, a)
go i HHTLine{..} = (hlFreqs `SVG.index` i, hlMags `SVG.index` i)
weightedAverage
:: (Foldable t, Fractional a)
=> t (a, a)
-> a
weightedAverage = uncurry (/) . foldl' go (0, 0)
where
go (!sx, !sw) (!x, !w) = (sx + x, sw + w)
dominantFreq
:: forall v n a. (VG.Vector v a, KnownNat n, Ord a)
=> HHT v n a
-> SVG.Vector v n a
dominantFreq HHT{..} = SVG.generate $ \i -> (\(Max (Arg _ x)) -> x)
. sconcat
. fromMaybe err
. NE.nonEmpty
. map (go i)
$ hhtLines
where
go :: Finite n -> HHTLine v n a -> ArgMax a a
go i HHTLine{..} = Max $ Arg (hlMags `SVG.index` i)
(hlFreqs `SVG.index` i)
err = errorWithoutStackTrace "Numeric.HHT.dominantFreq: HHT was formed with no Intrinsic Mode Functions"
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
hilbertPolar
:: 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 + 1) a)
hilbertPolar v = (hilbertMag, hilbertPhase)
where
hilbertMag :: SVG.Vector v (n + 1) a
hilbertFreq :: SVG.Vector v n a
(hilbertMag, hilbertFreq) = hilbertMagFreq v
hilbertPhase :: SVG.Vector v (n + 1) a
hilbertPhase = SVG.scanl' (+) 0 hilbertFreq
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