on fa9fa09 help sign in
357 lines 12 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Fold.Window
3-- Copyright : (c) 2020 Composewell Technologies
4-- License : Apache-2.0
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9-- Simple incremental statistical measures over a stream of data. All
10-- operations use numerically stable floating point arithmetic.
11--
12-- Measurements can be performed over the entire input stream or on a sliding
13-- window of fixed or variable size. Where possible, measures are computed
14-- online without buffering the input stream.
15--
16-- Currently there is no overflow detection.
17--
18-- For more advanced statistical measures see the @streamly-statistics@
19-- package.
21-- XXX A window fold can be driven either using the Ring.slidingWindow
22-- combinator or by zipping nthLast fold and last fold.
24module Streamly.Internal.Data.Fold.Window
25 (
26 -- * Incremental Folds
27 -- | Folds of type @Fold m (a, Maybe a) b@ are incremental sliding window
28 -- folds. An input of type @(a, Nothing)@ indicates that the input element
29 -- @a@ is being inserted in the window without ejecting an old value
30 -- increasing the window size by 1. An input of type @(a, Just a)@
31 -- indicates that the first element is being inserted in the window and the
32 -- second element is being removed from the window, the window size remains
33 -- the same. The window size can only increase and never decrease.
34 --
35 -- You can compute the statistics over the entire stream using sliding
36 -- window folds by keeping the second element of the input tuple as
37 -- @Nothing@.
38 --
39 windowLmap
40 , cumulative
42 , windowRollingMap
43 , windowRollingMapM
45 -- ** Sums
46 , windowLength
47 , windowSum
48 , windowSumInt
49 , windowPowerSum
50 , windowPowerSumFrac
52 -- ** Location
53 , windowMinimum
54 , windowMaximum
55 , windowRange
56 , windowMean
57 )
58where
60import Control.Monad.IO.Class (MonadIO (liftIO))
61import Data.Bifunctor(bimap)
62import Foreign.Storable (Storable, peek)
64import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
65import Streamly.Internal.Data.Tuple.Strict
66 (Tuple'(..), Tuple3Fused' (Tuple3Fused'))
68import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
70import qualified Streamly.Internal.Data.Fold.Type as Fold
71import qualified Streamly.Internal.Data.Ring as Ring
73import Prelude hiding (length, sum, minimum, maximum)
75-- $setup
76-- >>> import Data.Bifunctor(bimap)
77-- >>> import qualified Streamly.Data.Fold as Fold
78-- >>> import qualified Streamly.Internal.Data.Fold.Window as Fold
79-- >>> import qualified Streamly.Internal.Data.Ring as Ring
80-- >>> import qualified Streamly.Data.Stream as Stream
81-- >>> import Prelude hiding (length, sum, minimum, maximum)
83-------------------------------------------------------------------------------
84-- Utilities
85-------------------------------------------------------------------------------
87-- | Map a function on the incoming as well as outgoing element of a rolling
88-- window fold.
89--
90-- >>> lmap f = Fold.lmap (bimap f (f <$>))
91--
92{-# INLINE windowLmap #-}
93windowLmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
94windowLmap f = Fold.lmap (bimap f (f <$>))
96-- | Convert an incremental fold to a cumulative fold using the entire input
97-- stream as a single window.
98--
99-- >>> cumulative f = Fold.lmap (\x -> (x, Nothing)) f
100--
101{-# INLINE cumulative #-}
102cumulative :: Fold m (a, Maybe a) b -> Fold m a b
103cumulative = Fold.lmap (, Nothing)
105-- XXX Exchange the first two arguments of rollingMap or exchange the order in
106-- the fold input tuple.
108-- | Apply an effectful function on the latest and the oldest element of the
109-- window.
110{-# INLINE windowRollingMapM #-}
111windowRollingMapM :: Monad m =>
112 (Maybe a -> a -> m (Maybe b)) -> Fold m (a, Maybe a) (Maybe b)
113windowRollingMapM f = Fold.foldlM' f1 initial
115 where
117 initial = return Nothing
119 f1 _ (a, ma) = f ma a
121-- | Apply a pure function on the latest and the oldest element of the window.
122--
123-- >>> windowRollingMap f = Fold.windowRollingMapM (\x y -> return $ f x y)
124--
125{-# INLINE windowRollingMap #-}
126windowRollingMap :: Monad m =>
127 (Maybe a -> a -> Maybe b) -> Fold m (a, Maybe a) (Maybe b)
128windowRollingMap f = Fold.foldl' f1 initial
130 where
132 initial = Nothing
134 f1 _ (a, ma) = f ma a
136-------------------------------------------------------------------------------
137-- Sum
138-------------------------------------------------------------------------------
140-- XXX Overflow.
141--
142-- | The sum of all the elements in a rolling window. The input elements are
143-- required to be intergal numbers.
144--
145-- This was written in the hope that it would be a tiny bit faster than 'sum'
146-- for 'Integral' values. But turns out that 'sum' is 2% faster than this even
147-- for intergal values!
148--
149-- /Internal/
150--
151{-# INLINE windowSumInt #-}
152windowSumInt :: forall m a. (Monad m, Integral a) => Fold m (a, Maybe a) a
153windowSumInt = Fold step initial extract extract
155 where
157 initial = return $ Partial (0 :: a)
159 step s (a, ma) =
160 return
161 $ Partial
162 $ case ma of
163 Nothing -> s + a
164 Just old -> s + a - old
166 extract = return
168-- XXX Overflow.
169--
170-- | Sum of all the elements in a rolling window:
171--
172-- \(S = \sum_{i=1}^n x_{i}\)
173--
174-- This is the first power sum.
175--
176-- >>> sum = powerSum 1
177--
178-- Uses Kahan-Babuska-Neumaier style summation for numerical stability of
179-- floating precision arithmetic.
180--
181-- /Space/: \(\mathcal{O}(1)\)
182--
183-- /Time/: \(\mathcal{O}(n)\)
184--
185{-# INLINE windowSum #-}
186windowSum :: forall m a. (Monad m, Num a) => Fold m (a, Maybe a) a
187windowSum = Fold step initial extract extract
189 where
191 initial =
192 return
193 $ Partial
194 $ Tuple'
195 (0 :: a) -- running sum
196 (0 :: a) -- accumulated rounding error
198 step (Tuple' total err) (new, mOld) =
199 let incr =
200 case mOld of
201 -- XXX new may be large and err may be small we may lose it
202 Nothing -> new - err
203 -- XXX if (new - old) is large we may lose err
204 Just old -> (new - old) - err
205 -- total is large and incr may be small, we may round incr here but
206 -- we will accumulate the rounding error in err1 in the next step.
207 total1 = total + incr
208 -- Accumulate any rounding error in err1
209 -- XXX In the Nothing case above we may lose err, therefore we
210 -- should use ((total1 - total) - new) + err here.
211 -- Or even in the just case if (new - old) is large we may lose
212 -- err, so we should use ((total1 - total) + (old - new)) + err.
213 err1 = (total1 - total) - incr
214 in return $ Partial $ Tuple' total1 err1
216 extract (Tuple' total _) = return total
218-- | The number of elements in the rolling window.
219--
220-- This is the \(0\)th power sum.
221--
222-- >>> length = powerSum 0
223--
224{-# INLINE windowLength #-}
225windowLength :: (Monad m, Num b) => Fold m (a, Maybe a) b
226windowLength = Fold.foldl' step 0
228 where
230 step w (_, Nothing) = w + 1
231 step w _ = w
233-- | Sum of the \(k\)th power of all the elements in a rolling window:
234--
235-- \(S_k = \sum_{i=1}^n x_{i}^k\)
236--
237-- >>> powerSum k = lmap (^ k) sum
238--
239-- /Space/: \(\mathcal{O}(1)\)
240--
241-- /Time/: \(\mathcal{O}(n)\)
242{-# INLINE windowPowerSum #-}
243windowPowerSum :: (Monad m, Num a) => Int -> Fold m (a, Maybe a) a
244windowPowerSum k = windowLmap (^ k) windowSum
246-- | Like 'powerSum' but powers can be negative or fractional. This is slower
247-- than 'powerSum' for positive intergal powers.
248--
249-- >>> powerSumFrac p = lmap (** p) sum
250--
251{-# INLINE windowPowerSumFrac #-}
252windowPowerSumFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
253windowPowerSumFrac p = windowLmap (** p) windowSum
255-------------------------------------------------------------------------------
256-- Location
257-------------------------------------------------------------------------------
259-- XXX Remove MonadIO constraint
261-- | Determine the maximum and minimum in a rolling window.
262--
263-- If you want to compute the range of the entire stream @Fold.teeWith (,)
264-- Fold.maximum Fold.minimum@ would be much faster.
265--
266-- /Space/: \(\mathcal{O}(n)\) where @n@ is the window size.
267--
268-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
269--
270{-# INLINE windowRange #-}
271windowRange :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe (a, a))
272windowRange n = Fold step initial extract extract
274 where
276 -- XXX Use Ring unfold and then fold for composing maximum and minimum to
277 -- get the range.
279 initial =
280 if n <= 0
281 then error "range: window size must be > 0"
282 else
283 let f (a, b) = Partial $ Tuple3Fused' a b (0 :: Int)
284 in fmap f $ liftIO $ Ring.new n
286 step (Tuple3Fused' rb rh i) a = do
287 rh1 <- liftIO $ Ring.unsafeInsert rb rh a
288 return $ Partial $ Tuple3Fused' rb rh1 (i + 1)
290 -- XXX We need better Ring array APIs so that we can unfold the ring to a
291 -- stream and fold the stream using a fold of our choice.
292 --
293 -- We could just scan the stream to get a stream of ring buffers and then
294 -- map required folds over those, but we need to be careful that all those
295 -- rings refer to the same mutable ring, therefore, downstream needs to
296 -- process those strictly before it can change.
297 foldFunc i
298 | i < n = Ring.unsafeFoldRingM
299 | otherwise = Ring.unsafeFoldRingFullM
301 extract (Tuple3Fused' rb rh i) =
302 if i == 0
303 then return Nothing
304 else do
305 -- Here we use "ringStart" over "ringHead" as "ringHead" will be
306 -- uninitialized if the ring is not full.
307 -- Using "unsafeForeignPtrToPtr" here is safe as we touch the ring
308 -- again in "foldFunc".
309 x <- liftIO $ peek (unsafeForeignPtrToPtr (Ring.ringStart rb))
310 let accum (mn, mx) a = return (min mn a, max mx a)
311 fmap Just $ foldFunc i rh accum (x, x) rb
313-- | Find the minimum element in a rolling window.
314--
315-- This implementation traverses the entire window buffer to compute the
316-- minimum whenever we demand it. It performs better than the dequeue based
317-- implementation in @streamly-statistics@ package when the window size is
318-- small (< 30).
319--
320-- If you want to compute the minimum of the entire stream
321-- 'Streamly.Data.Fold.minimum' is much faster.
322--
323-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
324--
325{-# INLINE windowMinimum #-}
326windowMinimum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
327windowMinimum n = fmap (fmap fst) $ windowRange n
329-- | The maximum element in a rolling window.
330--
331-- See the performance related comments in 'minimum'.
332--
333-- If you want to compute the maximum of the entire stream 'Fold.maximum' would
334-- be much faster.
335--
336-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
337--
338{-# INLINE windowMaximum #-}
339windowMaximum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
340windowMaximum n = fmap (fmap snd) $ windowRange n
342-- | Arithmetic mean of elements in a sliding window:
343--
344-- \(\mu = \frac{\sum_{i=1}^n x_{i}}{n}\)
345--
346-- This is also known as the Simple Moving Average (SMA) when used in the
347-- sliding window and Cumulative Moving Avergae (CMA) when used on the entire
348-- stream.
349--
350-- >>> mean = Fold.teeWith (/) sum length
351--
352-- /Space/: \(\mathcal{O}(1)\)
353--
354-- /Time/: \(\mathcal{O}(n)\)
355{-# INLINE windowMean #-}
356windowMean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
357windowMean = Fold.teeWith (/) windowSum windowLength