{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
module Numeric.EMD.Sift (
Sifter(..), SiftResult(..), SingleSift(..), SM
, defaultSifter
, siftStdDev
, siftTimes
, siftEnergyDiff
, siftSCond
, siftAnd
, siftOr
, envMean
, energyDiff
, normalizeProj
, siftCauchy
, siftPairs
, siftProj
, siftPairs_
, siftProj_
, sift, envelopes, rms
) where
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State
import Data.Conduino
import Data.Conduino.Internal
import Data.Default.Class
import Data.Finite
import Data.Sequence (Seq(..))
import GHC.TypeNats
import Numeric.EMD.Internal
import Numeric.EMD.Internal.Extrema
import Numeric.EMD.Internal.Spline
import qualified Data.Conduino.Combinators as C
import qualified Data.Map as M
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Sized as SVG
instance (VG.Vector v a, Fractional a, Ord a) => Default (Sifter v n a) where
def = defaultSifter
defaultSifter :: (VG.Vector v a, Fractional a, Ord a) => Sifter v n a
defaultSifter = siftStdDev 0.3 `siftOr` siftTimes 50
siftEnergyDiff
:: (VG.Vector v a, KnownNat n, Floating a, Ord a)
=> a
-> a
-> Sifter v n a
siftEnergyDiff s t = siftProj energyDiff s
`siftAnd` siftProj envMean t
data SiftResult v n a = SRResidual !(SVG.Vector v n a)
| SRIMF !(SVG.Vector v n a) !Int
siftTimes :: Int -> Sifter v n a
siftTimes n = Sifter $ C.drop (n - 1) >> void awaitSurely
siftProj
:: Ord b
=> (SingleSift v n a -> SM v n a b)
-> b
-> Sifter v n a
siftProj p t = siftProj_ $ fmap (<= t) . p
siftProj_ :: (SingleSift v n a -> SM v n a Bool) -> Sifter v n a
siftProj_ p = Sifter go
where
go = do
v <- awaitSurely
r <- lift $ p v
unless r go
siftPairs
:: Ord b
=> (SingleSift v n a -> SingleSift v n a -> SM v n a b)
-> b
-> Sifter v n a
siftPairs p t = siftPairs_ $ \x y -> (<= t) <$> p x y
siftPairs_
:: (SingleSift v n a -> SingleSift v n a -> SM v n a Bool)
-> Sifter v n a
siftPairs_ p = Sifter $ go =<< awaitSurely
where
go s = do
s' <- awaitSurely
r <- lift $ p s s'
unless r (go s')
siftStdDev
:: forall v n a. (VG.Vector v a, Fractional a, Ord a)
=> a
-> Sifter v n a
siftStdDev = siftPairs $ \(SingleSift v _ _) (SingleSift v' _ _) -> pure $
SVG.sum (SVG.zipWith (\x x' -> (x-x')^(2::Int) / (x^(2::Int) + eps)) v v')
where
eps = 0.0000001
siftCauchy
:: (Fractional b, Ord b)
=> (SingleSift v n a -> b)
-> b
-> Sifter v n a
siftCauchy p = siftPairs $ \s s' ->
let ps = p s
ps' = p s'
δ = ps' - ps
in pure $ (δ * δ) / (ps * ps)
siftSCond
:: (VG.Vector v a, KnownNat n, Fractional a, Ord a)
=> Int
-> Sifter v (n + 1) a
siftSCond n = Sifter $ C.map (crossCount . ssResult)
.| C.consecutive n
.| C.concatMap pick
.| C.dropWhile notGood
where
pick Empty = Nothing
pick (xs :|> x) = (xs, x) <$ guard (length xs == (n - 1))
notGood (xs, x) = all ((<= 1) . abs . subtract x) xs
crossCount xs = M.size mins + M.size maxs + crosses
where
(mins, maxs) = extrema xs
crosses = fst . flip execState (0, Nothing) . flip SVG.mapM_ xs $ \x -> modify $ \(!i, !y) ->
let xPos = x > 0
i' = case y of
Nothing -> i
Just y'
| xPos == y' -> i
| otherwise -> i + 1
in (i', Just xPos)
siftOr :: Sifter v n a -> Sifter v n a -> Sifter v n a
siftOr (Sifter p) (Sifter q) = Sifter $ altSink p q
infixr 2 `siftOr`
siftAnd :: Sifter v n a -> Sifter v n a -> Sifter v n a
siftAnd (Sifter p) (Sifter q) = Sifter $ zipSink (id <$ p) q
infixr 3 `siftAnd`
envMean
:: (VG.Vector v a, KnownNat n, Floating a)
=> SingleSift v n a
-> SM v n a a
envMean SingleSift{..} = pure $
rms $ SVG.zipWith (\x y -> (x + y) / 2) ssMinEnv ssMaxEnv
energyDiff
:: (VG.Vector v a, Floating a)
=> SingleSift v n a
-> SM v n a a
energyDiff SingleSift{..} = do
v0 <- ask
pure . sqrt . abs . SVG.sum
$ SVG.zipWith (\x c -> c * (x - c)) v0 ssResult
normalizeProj
:: (VG.Vector v a, KnownNat n, Floating a)
=> (SingleSift v n a -> SM v n a a)
-> (SingleSift v n a -> SM v n a a)
normalizeProj f ss = do
v0 <- asks rms
r <- f ss
pure $ r / v0
rms :: (VG.Vector v a, KnownNat n, Floating a) => SVG.Vector v n a -> a
rms xs = sqrt $ SVG.foldl' (\s x -> s + x*x) 0 xs / fromIntegral (SVG.length xs)
sift
:: forall v n a. (VG.Vector v a, KnownNat n, Floating a, Ord a)
=> EMDOpts v (n + 1) a
-> SVG.Vector v (n + 1) a
-> SiftResult v (n + 1) a
sift EO{..} v0 = case execStateT (runPipe sifterPipe) (0, v0) of
Left v -> SRResidual v
Right (!i, !v) -> SRIMF v i
where
sifterPipe = C.repeatM go
.| hoistPipe
(pure . (`runReader` v0))
(sPipe eoSifter)
go = StateT $ \(!i, !v) ->
case sift' eoSplineEnd eoBoundaryHandler v of
Nothing -> Left v
Just ss@SingleSift{..} -> Right (ss, (i + 1, ssResult))
sift'
:: (VG.Vector v a, KnownNat n, Fractional a, Ord a)
=> SplineEnd a
-> Maybe BoundaryHandler
-> SVG.Vector v (n + 1) a
-> Maybe (SingleSift v (n + 1) a)
sift' se bh v = do
(mins, maxs) <- envelopes se bh v
pure SingleSift
{ ssResult = SVG.zipWith3 (\x mi ma -> x - (mi + ma)/2) v mins maxs
, ssMinEnv = mins
, ssMaxEnv = maxs
}
envelopes
:: (VG.Vector v a, KnownNat n, Fractional a, Ord a)
=> SplineEnd a
-> Maybe BoundaryHandler
-> SVG.Vector v (n + 1) a
-> Maybe (SVG.Vector v (n + 1) a, SVG.Vector v (n + 1) a)
envelopes se bh xs = do
when (bh == Just BHClamp) $ do
guard (M.size mins > 1)
guard (M.size maxs > 1)
(,) <$> splineAgainst se emin mins
<*> splineAgainst se emax maxs
where
(mins,maxs) = extrema xs
(emin,emax) = case bh of
Nothing -> mempty
Just bh' -> extendExtrema xs bh' (mins,maxs)
extendExtrema
:: forall v n a. (VG.Vector v a, KnownNat n)
=> SVG.Vector v (n + 1) a
-> BoundaryHandler
-> (M.Map (Finite (n + 1)) a, M.Map (Finite (n + 1)) a)
-> (M.Map Int a, M.Map Int a)
extendExtrema xs = \case
BHClamp -> const (firstLast, firstLast)
BHSymmetric -> \(mins, maxs) ->
let addFirst = case (flippedMin, flippedMax) of
(Nothing , Nothing ) -> mempty
(Just (_,mn) , Nothing ) -> (mn , firstPoint)
(Nothing , Just (_,mx) ) -> (firstPoint, mx )
(Just (mni,mn), Just (mxi,mx))
| mni < mxi -> (mn , firstPoint)
| otherwise -> (firstPoint, mx )
where
flippedMin = flip fmap (M.lookupMin mins) $ \(minIx, minVal) ->
(minIx, M.singleton (negate (fromIntegral minIx)) minVal)
flippedMax = flip fmap (M.lookupMin maxs) $ \(maxIx, maxVal) ->
(maxIx, M.singleton (negate (fromIntegral maxIx)) maxVal)
addLast = case (flippedMin, flippedMax) of
(Nothing , Nothing ) -> mempty
(Just (_,mn) , Nothing ) -> (mn , lastPoint )
(Nothing , Just (_,mx) ) -> (lastPoint , mx )
(Just (mni,mn), Just (mxi,mx))
| mni > mxi -> (mn , lastPoint )
| otherwise -> (lastPoint , mx )
where
flippedMin = flip fmap (M.lookupMax mins) $ \(minIx, minVal) ->
(minIx, M.singleton (extendSym (fromIntegral minIx)) minVal)
flippedMax = flip fmap (M.lookupMax maxs) $ \(maxIx, maxVal) ->
(maxIx, M.singleton (extendSym (fromIntegral maxIx)) maxVal)
in addFirst `mappend` addLast
where
lastIx = fromIntegral $ maxBound @(Finite n)
firstPoint = M.singleton 0 (SVG.head xs)
lastPoint = M.singleton lastIx (SVG.last xs)
firstLast = firstPoint `mappend` lastPoint
extendSym i = 2 * lastIx - i
splineAgainst
:: (VG.Vector v a, KnownNat n, Fractional a, Ord a)
=> SplineEnd a
-> M.Map Int a
-> M.Map (Finite n) a
-> Maybe (SVG.Vector v n a)
splineAgainst se ext = fmap go
. makeSpline se
. mappend (M.mapKeysMonotonic fromIntegral ext)
. M.mapKeysMonotonic fromIntegral
where
go spline = SVG.generate (sampleSpline spline . fromIntegral)