diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,11 @@ +Changelog +========= + +Version 0.1.0.0 +--------------- + +*July 25, 2018* + + + +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Justin Le (c) 2018 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Justin Le nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,11 @@ +# [emd][] + +[![emd on Hackage](https://img.shields.io/hackage/v/emd.svg?maxAge=86400)](https://hackage.haskell.org/package/emd) +[![Build Status](https://travis-ci.org/mstksg/emd.svg?branch=master)](https://travis-ci.org/mstksg/emd) + +[Empirical Mode decomposition][wiki] (Hilbert-Huang transform) in pure Haskell. + +[emd]: http://hackage.haskell.org/package/emd +[wiki]: https://en.wikipedia.org/wiki/Hilbert%E2%80%93Huang_transform + + diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/emd.cabal b/emd.cabal new file mode 100644 --- /dev/null +++ b/emd.cabal @@ -0,0 +1,66 @@ +-- This file has been generated from package.yaml by hpack version 0.28.2. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: 4cb66d27d7d3fcb9f625c87b73448433636262d9476571553b667d1a0f998496 + +name: emd +version: 0.1.0.0 +synopsis: Empirical Mode Decomposition (Hilbert-Huang Transform) +description: Please see the README on GitHub at +category: Math +homepage: https://github.com/mstksg/emd#readme +bug-reports: https://github.com/mstksg/emd/issues +author: Justin Le +maintainer: justin@jle.im +copyright: (c) Justin Le 2018 +license: BSD3 +license-file: LICENSE +tested-with: GHC >= 8.2 +build-type: Simple +cabal-version: >= 1.10 +extra-source-files: + CHANGELOG.md + README.md + test-data/sintest.csv + +source-repository head + type: git + location: https://github.com/mstksg/emd + +library + exposed-modules: + Numeric.EMD + Numeric.EMD.Internal.Spline + other-modules: + Numeric.EMD.Internal.Tridiagonal + Numeric.EMD.Internal.Extrema + hs-source-dirs: + src + ghc-options: -Wall -fwarn-redundant-constraints + build-depends: + base >=4.7 && <5 + , containers + , finite-typelits + , ghc-typelits-knownnat + , ghc-typelits-natnormalise + , transformers + , typelits-witnesses + , vector + , vector-sized + default-language: Haskell2010 + +test-suite emd-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Paths_emd + hs-source-dirs: + test + ghc-options: -Wall -fwarn-redundant-constraints -threaded -rtsopts -with-rtsopts=-N + build-depends: + HUnit + , base >=4.7 && <5 + , containers + , emd + default-language: Haskell2010 diff --git a/src/Numeric/EMD.hs b/src/Numeric/EMD.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/EMD.hs @@ -0,0 +1,192 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} + +-- | +-- Module : Numeric.EMD +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Empirical Mode Decomposition (Hilbert-Huang Transform) in pure Haskell. +-- +-- Main interface is 'emd', with 'defaultEO'. A tracing version that +-- outputs a log to stdout is also available, as 'emdTrace'. This can be +-- used to help track down a specific IMF that might be taking more time +-- than desired. +-- + +module Numeric.EMD ( + -- * EMD (Hilbert-Huang Transform) + emd + , emdTrace + , emd' + , EMD(..) + , EMDOpts(..), defaultEO, SiftCondition(..), defaultSC, SplineEnd(..) + -- * Internal + , sift, SiftResult(..) + , envelopes + ) where + +import Control.Monad.IO.Class +import Data.Finite +import Data.Functor.Identity +import GHC.TypeNats +import Numeric.EMD.Internal.Extrema +import Numeric.EMD.Internal.Spline +import Text.Printf +import qualified Data.Map as M +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Sized as SVG + +-- | Options for EMD composition. +data EMDOpts a = EO { eoSiftCondition :: SiftCondition a -- ^ stop condition for sifting + , eoSplineEnd :: SplineEnd -- ^ end conditions for envelope splines + , eoClampEnvelope :: Bool -- ^ if 'True', use time series endpoints as part of min and max envelopes + } + deriving (Show, Eq, Ord) + +-- | Default 'EMDOpts' +defaultEO :: Fractional a => EMDOpts a +defaultEO = EO { eoSiftCondition = defaultSC + , eoSplineEnd = SENotAKnot + , eoClampEnvelope = True + } + + +-- | Stop conditions for sifting process +data SiftCondition a = SCStdDev a -- ^ Stop using standard "SD" method + | SCTimes Int -- ^ Stop after a fixed number of iterations + | SCOr (SiftCondition a) (SiftCondition a) -- ^ one or the other + | SCAnd (SiftCondition a) (SiftCondition a) -- ^ both conditions met + deriving (Show, Eq, Ord) + +-- | Default 'SiftCondition' +defaultSC :: Fractional a => SiftCondition a +defaultSC = SCStdDev 0.3 + +-- | 'True' if stop +testCondition + :: (VG.Vector v a, Fractional a, Ord a) + => SiftCondition a + -> Int + -> SVG.Vector v n a + -> SVG.Vector v n a + -> Bool +testCondition = \case + SCStdDev t -> \_ v v' -> + let sd = SVG.sum $ SVG.zipWith (\x x' -> (x-x')^(2::Int) / (x^(2::Int) + eps)) v v' + in sd <= t + SCTimes l -> \i _ _ -> i >= l + SCOr f g -> \i v v' -> testCondition f i v v' || testCondition g i v v' + SCAnd f g -> \i v v' -> testCondition f i v v' && testCondition g i v v' + where + eps = 0.0000001 + +-- | An @'EMD' v n a@ is a Hilbert-Huang transform of a time series with +-- @n@ items of type @a@ stored in a vector @v@. +data EMD v n a = EMD { emdIMFs :: ![SVG.Vector v n a] + , emdResidual :: !(SVG.Vector v n a) + } + deriving Show + +-- | EMD decomposition (Hilbert-Huang Transform) of a given time series +-- with a given sifting stop condition. +emd :: (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => EMDOpts a + -> SVG.Vector v (n + 2) a + -> EMD v (n + 2) a +emd eo = runIdentity . emd' (const (pure ())) eo + +-- | 'emd', but tracing results to stdout as IMFs are found. Useful for +-- debugging to see how long each step is taking. +emdTrace + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a, MonadIO m) + => EMDOpts a + -> SVG.Vector v (n + 2) a + -> m (EMD v (n + 2) a) +emdTrace = emd' $ \case + SRResidual _ -> liftIO $ putStrLn "Residual found." + SRIMF _ i -> liftIO $ printf "IMF found (%d iterations)\n" i + +-- | 'emd' with a callback for each found IMF. +emd' + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a, Applicative m) + => (SiftResult v (n + 2) a -> m r) + -> EMDOpts a + -> SVG.Vector v (n + 2) a + -> m (EMD v (n + 2) a) +emd' cb eo = go id + where + go !imfs !v = cb res *> case res of + SRResidual r -> pure $ EMD (imfs []) r + SRIMF v' _ -> go (imfs . (v':)) (v - v') + where + res = sift eo v + +-- | The result of a sifting operation. Each sift either yields +-- a residual, or a new IMF. +data SiftResult v n a = SRResidual !(SVG.Vector v n a) + | SRIMF !(SVG.Vector v n a) !Int -- number of iterations + +-- | Iterated sifting process, used to produce either an IMF or a residual. +sift + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => EMDOpts a + -> SVG.Vector v (n + 2) a + -> SiftResult v (n + 2) a +sift EO{..} = go 1 + where + go !i !v = case sift' eoSplineEnd eoClampEnvelope v of + Nothing -> SRResidual v + Just !v' + | testCondition eoSiftCondition i v v' -> SRIMF v' i + | otherwise -> go (i + 1) v' + +-- | Single sift +sift' + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => SplineEnd + -> Bool + -> SVG.Vector v (n + 2) a + -> Maybe (SVG.Vector v (n + 2) a) +sift' se cl v = go <$> envelopes se cl v + where + go (mins, maxs) = SVG.zipWith3 (\x mi ma -> x - (mi + ma)/2) v mins maxs + +-- | Returns cubic splines of local minimums and maximums. Returns +-- 'Nothing' if there are not enough local minimum or maximums to create +-- the splines. +envelopes + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => SplineEnd + -> Bool + -> SVG.Vector v (n + 2) a + -> Maybe (SVG.Vector v (n + 2) a, SVG.Vector v (n + 2) a) +envelopes se cl xs = (,) <$> splineAgainst se mins' + <*> splineAgainst se maxs' + where + minMax = M.fromList [(minBound, SVG.head xs), (maxBound, SVG.last xs)] + (mins,maxs) = extrema xs + (mins', maxs') + | cl = (mins `M.union` minMax, maxs `M.union` minMax) + | otherwise = (mins, maxs) + +-- | Build a splined vector against a map of control points. +splineAgainst + :: (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => SplineEnd + -> M.Map (Finite n) a + -> Maybe (SVG.Vector v n a) +splineAgainst se = fmap go . makeSpline se . M.mapKeysMonotonic fromIntegral + where + go spline = SVG.generate (sampleSpline spline . fromIntegral) diff --git a/src/Numeric/EMD/Internal/Extrema.hs b/src/Numeric/EMD/Internal/Extrema.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/EMD/Internal/Extrema.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} + +module Numeric.EMD.Internal.Extrema ( + extrema + ) where + +import Control.Applicative.Backwards +import Control.Monad.Trans.Class +import Control.Monad.Trans.State +import Control.Monad.Trans.Writer +import Data.Coerce +import Data.Either +import Data.Finite +import Data.Foldable +import Data.Functor.Identity +import Data.Monoid +import GHC.TypeNats +import qualified Data.Map as M +import qualified Data.Set as S +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Sized as SVG + +data ExState n = ESStart [Finite n] + | ESDown [Finite n] + | ESUp [Finite n] + +esList + :: Functor f + => ([Finite n] -> f [Finite m]) + -> ExState n + -> f (ExState m) +esList f = \case + ESStart is -> ESStart <$> f is + ESDown is -> ESDown <$> f is + ESUp is -> ESUp <$> f is + +data ExOut n = ExMin (Finite n) + | ExMax (Finite n) + +exEither :: ExOut n -> Either (Finite n) (Finite n) +exEither = \case + ExMin i -> Left i + ExMax i -> Right i + +-- | Treats every item in a "plateu" as a local minimum or maximum. +extrema + :: forall v n a. (VG.Vector v a, KnownNat n, Fractional a, Ord a) + => SVG.Vector v (n + 1) a + -> (M.Map (Finite (n + 1)) a, M.Map (Finite (n + 1)) a) +extrema xs = (makeMap mins, makeMap maxs) + where + (mins,maxs) = partitionEithers . map exEither $ optima + dxs :: SVG.Vector v n a + dxs = SVG.tail xs - SVG.init xs + optima :: [ExOut n] + optima = flip appEndo [] + . execWriter + . flip execStateT (ESStart []) + $ SVG.imapM_ go dxs + where + go :: Finite n + -> a + -> StateT (ExState n) (Writer (Endo [ExOut n])) () + go i d = case compare d 0 of + LT -> do + get >>= \case + ESStart _ -> pure () + ESDown _ -> pure () + ESUp is -> lift $ do + forwards . traverse_ (Backwards . tell . Endo . (:) . ExMax) $ is + tell $ Endo (ExMax i:) + put $ ESDown [] + EQ -> modify $ over esList (i:) + GT -> do + get >>= \case + ESStart _ -> pure () + ESDown is -> lift $ do + forwards . traverse_ (Backwards . tell . Endo . (:) . ExMin) $ is + tell $ Endo (ExMin i:) + ESUp _ -> pure () + put $ ESUp [] + makeMap :: [Finite n] -> M.Map (Finite (n + 1)) a + makeMap = M.fromSet (xs `SVG.index`) + . S.fromAscList + . map weaken + +over :: ((a -> Identity b) -> s -> Identity t) + -> (a -> b) -> (s -> t) +over = coerce diff --git a/src/Numeric/EMD/Internal/Spline.hs b/src/Numeric/EMD/Internal/Spline.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/EMD/Internal/Spline.hs @@ -0,0 +1,185 @@ +{-# LANGUAGE ApplicativeDo #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Numeric.EMD.Internal.Spline +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Internal splining functionality exported for testing purposes only. +-- This will likely go away in future versions, so please do not depend on +-- this! +-- + +module Numeric.EMD.Internal.Spline ( + Spline, SplineEnd(..) + , makeSpline + , sampleSpline + ) where + +import Data.Finite +import Data.Proxy +import Data.Type.Equality +import GHC.TypeLits.Compare +import GHC.TypeNats +import Numeric.EMD.Internal.Tridiagonal +import qualified Data.Map as M +import qualified Data.Vector.Sized as SV + +-- | End condition for spline +data SplineEnd = SENotAKnot + | SENatural + deriving (Show, Eq, Ord) + +data SplineCoef a = SC { scAlpha :: !a -- ^ a + , scBeta :: !a -- ^ b + , scGamma0 :: !a -- ^ y_{i-1} + , scGamma1 :: !a -- ^ y_i + , scDelta :: !a -- ^ x_i - x_{i-1} + } + deriving Show + +-- | 1D Cubic spline +data Spline a = Spline { splineHead :: !(a, SplineCoef a) + , splineTail :: !(M.Map a (SplineCoef a)) + } + +runSplineCoef + :: Fractional a + => a + -> SplineCoef a + -> a + -> a +runSplineCoef x0 (SC α β γ0 γ1 δ) x = q * γ0 + + t * γ1 + + t * q * (q * α + t * β) + where + t = (x - x0) / δ + q = 1 - t + +-- | Sample a spline at a given point. +sampleSpline + :: (Fractional a, Ord a) + => Spline a + -> a + -> a +sampleSpline Spline{..} x = case x `M.lookupLE` splineTail of + Nothing -> + let (x0, sc) = splineHead + in runSplineCoef x0 sc x + Just (x0, sc) -> runSplineCoef x0 sc x + +-- | Build a cubic spline based on control points using given end +-- conditions (not-a-knot, or natural) +-- +-- +makeSpline + :: forall a. (Ord a, Fractional a) + => SplineEnd + -> M.Map a a -- ^ (x, y) + -> Maybe (Spline a) +makeSpline se ps = do + (xy0, ps') <- M.minViewWithKey ps + SV.withSizedList (M.toList ps') $ \(xsys :: SV.Vector n (a, a)) -> do + Refl <- Proxy @1 `isLE` Proxy @n + let xs, ys :: SV.Vector (n + 1) a + (xs, ys) = SV.unzip $ xy0 `SV.cons` xsys + dxs, dys :: SV.Vector n a + dxs = SV.tail xs - SV.init xs + rdxs :: SV.Vector n a + rdxs = recip dxs + rdxssq :: SV.Vector n a + rdxssq = rdxs * rdxs + dys = SV.tail ys - SV.init ys + dydxssq = dys * rdxssq + mainDiag :: SV.Vector (n - 1) a + mainDiag = SV.zipWith (\rdx0 rdx1 -> 2 * ( rdx0 + rdx1 )) + (SV.init rdxs) + (SV.tail rdxs) + lowerDiag :: SV.Vector (n - 1) a + lowerDiag = SV.take rdxs + upperDiag :: SV.Vector (n - 1) a + upperDiag = SV.tail rdxs + rhs :: SV.Vector (n - 1) a + rhs = SV.zipWith (\dydxsq0 dydxsq1 -> 3 * (dydxsq0 + dydxsq1)) + (SV.init dydxssq) + (SV.tail dydxssq) + EE{..} = case se of + SENotAKnot -> notAKnot rdxs rdxssq dydxssq + SENatural -> natural rdxs dydxssq + solution <- solveTridiagonal (lowerDiag `SV.snoc` eeLower1) + (eeMain0 `SV.cons` mainDiag `SV.snoc` eeMain1) + (eeUpper0 `SV.cons` upperDiag) + (eeRhs0 `SV.cons` rhs `SV.snoc` eeRhs1) + let as :: SV.Vector n a + as = SV.zipWith3 (\k dx dy -> k * dx - dy) (SV.init solution) dxs dys + bs :: SV.Vector n a + bs = SV.zipWith3 (\k dx dy -> - k * dx + dy) (SV.tail solution) dxs dys + coefs :: SV.Vector n (a, SplineCoef a) + coefs = SV.zipWith6 (\x α β γ0 γ1 δ -> (x, SC α β γ0 γ1 δ)) + (SV.init xs) as bs (SV.init ys) (SV.tail ys) dxs + + pure Spline + { splineHead = SV.head coefs + , splineTail = M.fromAscList . SV.toList . SV.tail $ coefs + } + +data EndEqn a = EE { eeMain0 :: !a + , eeUpper0 :: !a + , eeLower1 :: !a + , eeMain1 :: !a + , eeRhs0 :: !a + , eeRhs1 :: !a + } + +natural + :: (KnownNat n, Num a) + => SV.Vector (n + 1) a + -> SV.Vector (n + 1) a + -> EndEqn a +natural rdxs dydxssq = EE + { eeMain0 = 2 * (rdxs `SV.index` minBound) + , eeUpper0 = rdxs `SV.index` minBound + , eeLower1 = rdxs `SV.index` maxBound + , eeMain1 = 2 * (rdxs `SV.index` maxBound) + , eeRhs0 = 3 * (dydxssq `SV.index` minBound) + , eeRhs1 = 3 * (dydxssq `SV.index` maxBound) + } + +notAKnot + :: (KnownNat n, Num a) + => SV.Vector (n + 1) a + -> SV.Vector (n + 1) a + -> SV.Vector (n + 1) a + -> EndEqn a +notAKnot rdxs rdxssq dydxssq = EE + { eeMain0 = rdxssq `SV.index` minBound + rdx12Upper + , eeUpper0 = rdxssq `SV.index` minBound + + rdxssq `SV.index` shift minBound + + 2 * rdx12Upper + , eeLower1 = - (rdxssq `SV.index` weaken maxBound) + - (rdxssq `SV.index` maxBound) + - 2 * rdx12Lower + , eeMain1 = - rdxssq `SV.index` maxBound - rdx12Lower + , eeRhs0 = 2 * (dydxssq `SV.index` minBound) * (rdxs `SV.index` minBound) + + 3 * (dydxssq `SV.index` minBound) * (rdxs `SV.index` shift minBound) + + (dydxssq `SV.index` shift minBound) * (rdxs `SV.index` shift minBound) + , eeRhs1 = - (dydxssq `SV.index` weaken maxBound) * (rdxs `SV.index` weaken maxBound) + - 3 * (dydxssq `SV.index` maxBound) * (rdxs `SV.index` weaken maxBound) + - 2 * (dydxssq `SV.index` maxBound) * (rdxs `SV.index` maxBound) + } + where + rdx12Upper = rdxs `SV.index` minBound * rdxs `SV.index` shift minBound + rdx12Lower = rdxs `SV.index` maxBound * rdxs `SV.index` weaken maxBound diff --git a/src/Numeric/EMD/Internal/Tridiagonal.hs b/src/Numeric/EMD/Internal/Tridiagonal.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/EMD/Internal/Tridiagonal.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} + +module Numeric.EMD.Internal.Tridiagonal ( + solveTridiagonal + ) where + +import Control.Applicative.Backwards +import Control.Monad +import Control.Monad.ST +import Control.Monad.Trans.Class +import Control.Monad.Trans.Maybe +import Data.Finite +import Data.Foldable +import GHC.TypeNats +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Mutable.Sized as SMVG +import qualified Data.Vector.Generic.Sized as SVG + +-- | +-- +-- Will return 'Nothing' if the matrix is not invertible. This will happen +-- if: +-- +-- 1. The first item in the main diagonal is zero +-- 2. There is any i such that b_{i + 1} = a_i * c_i. That is, an item in +-- the main diagonal is equal to the product of the off-diagonal elements +-- a row above it +-- 3. Another mystery condition! +solveTridiagonal + :: forall v n a. (VG.Vector v a, KnownNat n, Fractional a, Eq a) + => SVG.Vector v (n + 1) a -- ^ a: Bottom diagonal of M + -> SVG.Vector v (n + 2) a -- ^ b: Main diagonal of M + -> SVG.Vector v (n + 1) a -- ^ c: Upper diagonal of M + -> SVG.Vector v (n + 2) a -- ^ y + -> Maybe (SVG.Vector v (n + 2) a) -- ^ x such that M x = y +solveTridiagonal as bs cs ds = runST $ runMaybeT $ do + guard $ SVG.head bs /= 0 + cs' <- makeCs + mxs <- lift $ SVG.thaw ds + makeDs cs' mxs + forwards . for_ (consecFinites @(n + 1)) $ \(i0, i1) -> Backwards $ do + x1 <- lift $ SMVG.read mxs i1 + let sbr = cs' `SVG.index` i0 * x1 + lift $ SMVG.modify mxs (subtract sbr) (weaken i0) + lift $ SVG.freeze mxs + where + makeCs :: MaybeT (ST s) (SVG.Vector v (n + 1) a) + makeCs = do + mcs <- lift $ SVG.thaw cs + lift $ SMVG.modify mcs (/ SVG.head bs) minBound + for_ (consecFinites @n) $ \(i0, i1) -> do + c0 <- lift $ SMVG.read mcs (weaken i0) + let dvr = bs `SVG.index` weaken i1 + - as `SVG.index` weaken i0 * c0 + guard $ dvr /= 0 + lift $ SMVG.modify mcs (/ dvr) i1 + lift $ SVG.freeze mcs + makeDs + :: SVG.Vector v (n + 1) a + -> SVG.MVector (VG.Mutable v) (n + 2) s a + -> MaybeT (ST s) () + makeDs cs' mds = do + lift $ SMVG.modify mds (/ SVG.head bs) minBound + for_ (consecFinites @(n + 1)) $ \(i0, i1) -> do + let c0 = cs' `SVG.index` i0 + d0 <- lift $ SMVG.read mds (weaken i0) + let sbr = as `SVG.index` i0 * d0 + dvr = bs `SVG.index` i1 + - as `SVG.index` i0 * c0 + guard $ dvr /= 0 + lift $ SMVG.modify mds ((/ dvr) . subtract sbr) i1 + +consecFinites :: KnownNat n => [(Finite n, Finite (n + 1))] +consecFinites = zip finites (tail finites) diff --git a/test-data/sintest.csv b/test-data/sintest.csv new file mode 100644 --- /dev/null +++ b/test-data/sintest.csv @@ -0,0 +1,41 @@ +0 +0.286388716728477 +0.5219840894569575 +0.7069551636589333 +0.8414709848078965 +0.9257005983773389 +0.9598130498407524 +0.9439773846716291 +0.8783626483434607 +0.7631378863297392 +0.5984721441039565 +0.3876642947461241 +0.1465325217623302 +-0.1059751637348173 +-0.3509107506327105 +-0.5697671704562872 +-0.7506351630848729 +-0.886684474333967 +-0.9712154996894535 +-0.9975286346372161 +-0.9589242746631385 +-0.8522846636228126 +-0.6888194388506665 +-0.4833200860508369 +-0.25057809092746 +-0.005384939184672421 +0.2374678834733893 +0.4631888913425887 +0.6569865987187891 +0.8065854562616487 +0.9097736600860038 +0.9668553426704856 +0.9781346364937249 +0.9439156740343529 +0.8645025877710003 +0.7401995101822982 +0.5713105737468775 +0.3581399109433696 +0.100991654250405 +-0.1998300638533851 +-0.5440211108893697 diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE TypeApplications #-} + +import Control.Monad +import Data.Maybe +import Numeric.EMD.Internal.Spline +import Test.HUnit +import qualified Data.Map as M +import qualified Data.Set as S + +main :: IO () +main = void . runTestTT $ TestList + [ "Sine spline" ~: splineTest + ] + +splineTest :: Assertion +splineTest = do + expected <- map (read @Double) . lines <$> readFile "test-data/sintest.csv" + roundOut expected @=? roundOut samples + where + roundOut :: [Double] -> [Double] + roundOut = map $ (/ 10e12) . fromInteger . round . (* 10e12) + spline :: Spline Double + spline = fromJust + . makeSpline SENotAKnot + . M.fromSet sin + . S.fromList + $ [0, 1, 2.5, 3.6, 5, 7, 8.1, 10] + samples :: [Double] + samples = sampleSpline spline . (/ 4) . fromInteger <$> [0..40] +