emd-0.1.1.0: src/Numeric/EMD/Internal/Spline.hs
{-# 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)
--
-- <https://en.wikipedia.org/wiki/Spline_interpolation>
makeSpline
:: forall a. (Ord a, Fractional a)
=> SplineEnd
-> M.Map a a -- ^ (x, y)
-> Maybe (Spline a)
makeSpline se ps = SV.withSizedList (M.toList ps) $ \(xsys :: SV.Vector n (a, a)) -> do
Refl <- Proxy @1 `isLE` Proxy @n
Refl <- Proxy @2 `isLE` Proxy @n
let xs, ys :: SV.Vector n a
(xs, ys) = SV.unzip xsys
dxs, dys :: SV.Vector (n - 1) a
dxs = SV.tail xs - SV.init xs
rdxs :: SV.Vector (n - 1) a
rdxs = recip dxs
rdxssq :: SV.Vector (n - 1) a
rdxssq = rdxs * rdxs
dys = SV.tail ys - SV.init ys
dydxssq = dys * rdxssq
mainDiag :: SV.Vector (n - 2) a
mainDiag = SV.zipWith (\rdx0 rdx1 -> 2 * ( rdx0 + rdx1 ))
(SV.init rdxs)
(SV.tail rdxs)
lowerDiag :: SV.Vector (n - 2) a
lowerDiag = SV.take rdxs
upperDiag :: SV.Vector (n - 2) a
upperDiag = SV.tail rdxs
rhs :: SV.Vector (n - 2) 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 - 1) a
as = SV.zipWith3 (\k dx dy -> k * dx - dy) (SV.init solution) dxs dys
bs :: SV.Vector (n - 1) a
bs = SV.zipWith3 (\k dx dy -> - k * dx + dy) (SV.tail solution) dxs dys
coefs :: SV.Vector (n - 1) (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