on 4eaeb10 help sign in
388 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 lmap
40 , cumulative
42 , rollingMap
43 , rollingMapM
45 -- ** Sums
46 , length
47 , sum
48 , sumInt
49 , powerSum
50 , powerSumFrac
52 -- ** Location
53 , minimum
54 , maximum
55 , range
56 , mean
58 -- ** Distribution
59 , frequency
60 )
61where
63import Control.Monad.IO.Class (MonadIO (liftIO))
64import Data.Bifunctor(bimap)
65import Foreign.Storable (Storable, peek)
67import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
68import Streamly.Internal.Data.Tuple.Strict
69 (Tuple'(..), Tuple3Fused' (Tuple3Fused'))
71import qualified Data.Map as Map
72import qualified Streamly.Internal.Data.Fold.Type as Fold
73import qualified Streamly.Internal.Data.Ring.Unboxed as Ring
75import Prelude hiding (length, sum, minimum, maximum)
77-- $setup
78-- >>> import Data.Bifunctor(bimap)
79-- >>> import qualified Streamly.Data.Fold as Fold
80-- >>> import qualified Streamly.Internal.Data.Fold.Window as FoldW
81-- >>> import qualified Streamly.Internal.Data.Ring.Unboxed as Ring
82-- >>> import qualified Streamly.Data.Stream as Stream
83-- >>> import Prelude hiding (length, sum, minimum, maximum)
85-------------------------------------------------------------------------------
86-- Utilities
87-------------------------------------------------------------------------------
89-- | Map a function on the incoming as well as outgoing element of a rolling
90-- window fold.
91--
92-- >>> lmap f = Fold.lmap (bimap f (f <$>))
93--
94{-# INLINE lmap #-}
95lmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
96lmap f = Fold.lmap (bimap f (f <$>))
98-- | Convert an incremental fold to a cumulative fold using the entire input
99-- stream as a single window.
100--
101-- >>> cumulative f = Fold.lmap (\x -> (x, Nothing)) f
102--
103{-# INLINE cumulative #-}
104cumulative :: Fold m (a, Maybe a) b -> Fold m a b
105cumulative = Fold.lmap (, Nothing)
107-- XXX Exchange the first two arguments of rollingMap or exchange the order in
108-- the fold input tuple.
110-- | Apply an effectful function on the latest and the oldest element of the
111-- window.
112{-# INLINE rollingMapM #-}
113rollingMapM :: Monad m =>
114 (Maybe a -> a -> m (Maybe b)) -> Fold m (a, Maybe a) (Maybe b)
115rollingMapM f = Fold.foldlM' f1 initial
117 where
119 initial = return Nothing
121 f1 _ (a, ma) = f ma a
123-- | Apply a pure function on the latest and the oldest element of the window.
124--
125-- >>> rollingMap f = FoldW.rollingMapM (\x y -> return $ f x y)
126--
127{-# INLINE rollingMap #-}
128rollingMap :: Monad m =>
129 (Maybe a -> a -> Maybe b) -> Fold m (a, Maybe a) (Maybe b)
130rollingMap f = Fold.foldl' f1 initial
132 where
134 initial = Nothing
136 f1 _ (a, ma) = f ma a
138-------------------------------------------------------------------------------
139-- Sum
140-------------------------------------------------------------------------------
142-- XXX Overflow.
143--
144-- | The sum of all the elements in a rolling window. The input elements are
145-- required to be intergal numbers.
146--
147-- This was written in the hope that it would be a tiny bit faster than 'sum'
148-- for 'Integral' values. But turns out that 'sum' is 2% faster than this even
149-- for intergal values!
150--
151-- /Internal/
152--
153{-# INLINE sumInt #-}
154sumInt :: forall m a. (Monad m, Integral a) => Fold m (a, Maybe a) a
155sumInt = Fold step initial extract
157 where
159 initial = return $ Partial (0 :: a)
161 step s (a, ma) =
162 return
163 $ Partial
164 $ case ma of
165 Nothing -> s + a
166 Just old -> s + a - old
168 extract = return
170-- XXX Overflow.
171--
172-- | Sum of all the elements in a rolling window:
173--
174-- \(S = \sum_{i=1}^n x_{i}\)
175--
176-- This is the first power sum.
177--
178-- >>> sum = powerSum 1
179--
180-- Uses Kahan-Babuska-Neumaier style summation for numerical stability of
181-- floating precision arithmetic.
182--
183-- /Space/: \(\mathcal{O}(1)\)
184--
185-- /Time/: \(\mathcal{O}(n)\)
186--
187{-# INLINE sum #-}
188sum :: forall m a. (Monad m, Num a) => Fold m (a, Maybe a) a
189sum = Fold step initial extract
191 where
193 initial =
194 return
195 $ Partial
196 $ Tuple'
197 (0 :: a) -- running sum
198 (0 :: a) -- accumulated rounding error
200 step (Tuple' total err) (new, mOld) =
201 let incr =
202 case mOld of
203 -- XXX new may be large and err may be small we may lose it
204 Nothing -> new - err
205 -- XXX if (new - old) is large we may lose err
206 Just old -> (new - old) - err
207 -- total is large and incr may be small, we may round incr here but
208 -- we will accumulate the rounding error in err1 in the next step.
209 total1 = total + incr
210 -- Accumulate any rounding error in err1
211 -- XXX In the Nothing case above we may lose err, therefore we
212 -- should use ((total1 - total) - new) + err here.
213 -- Or even in the just case if (new - old) is large we may lose
214 -- err, so we should use ((total1 - total) + (old - new)) + err.
215 err1 = (total1 - total) - incr
216 in return $ Partial $ Tuple' total1 err1
218 extract (Tuple' total _) = return total
220-- | The number of elements in the rolling window.
221--
222-- This is the \(0\)th power sum.
223--
224-- >>> length = powerSum 0
225--
226{-# INLINE length #-}
227length :: (Monad m, Num b) => Fold m (a, Maybe a) b
228length = Fold.foldl' step 0
230 where
232 step w (_, Nothing) = w + 1
233 step w _ = w
235-- | Sum of the \(k\)th power of all the elements in a rolling window:
236--
237-- \(S_k = \sum_{i=1}^n x_{i}^k\)
238--
239-- >>> powerSum k = lmap (^ k) sum
240--
241-- /Space/: \(\mathcal{O}(1)\)
242--
243-- /Time/: \(\mathcal{O}(n)\)
244{-# INLINE powerSum #-}
245powerSum :: (Monad m, Num a) => Int -> Fold m (a, Maybe a) a
246powerSum k = lmap (^ k) sum
248-- | Like 'powerSum' but powers can be negative or fractional. This is slower
249-- than 'powerSum' for positive intergal powers.
250--
251-- >>> powerSumFrac p = lmap (** p) sum
252--
253{-# INLINE powerSumFrac #-}
254powerSumFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
255powerSumFrac p = lmap (** p) sum
257-------------------------------------------------------------------------------
258-- Location
259-------------------------------------------------------------------------------
261-- XXX Remove MonadIO constraint
263-- | Determine the maximum and minimum in a rolling window.
264--
265-- If you want to compute the range of the entire stream @Fold.teeWith (,)
266-- Fold.maximum Fold.minimum@ would be much faster.
267--
268-- /Space/: \(\mathcal{O}(n)\) where @n@ is the window size.
269--
270-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
271--
272{-# INLINE range #-}
273range :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe (a, a))
274range n = Fold step initial extract
276 where
278 -- XXX Use Ring unfold and then fold for composing maximum and minimum to
279 -- get the range.
281 initial =
282 if n <= 0
283 then error "range: window size must be > 0"
284 else
285 let f (a, b) = Partial $ Tuple3Fused' a b (0 :: Int)
286 in fmap f $ liftIO $ Ring.new n
288 step (Tuple3Fused' rb rh i) a = do
289 rh1 <- liftIO $ Ring.unsafeInsert rb rh a
290 return $ Partial $ Tuple3Fused' rb rh1 (i + 1)
292 -- XXX We need better Ring array APIs so that we can unfold the ring to a
293 -- stream and fold the stream using a fold of our choice.
294 --
295 -- We could just scan the stream to get a stream of ring buffers and then
296 -- map required folds over those, but we need to be careful that all those
297 -- rings refer to the same mutable ring, therefore, downstream needs to
298 -- process those strictly before it can change.
299 foldFunc i
300 | i < n = Ring.unsafeFoldRingM
301 | otherwise = Ring.unsafeFoldRingFullM
303 extract (Tuple3Fused' rb rh i) =
304 if i == 0
305 then return Nothing
306 else do
307 x <- liftIO $ peek rh
308 let accum (mn, mx) a = return (min mn a, max mx a)
309 fmap Just $ foldFunc i rh accum (x, x) rb
311-- | Find the minimum element in a rolling window.
312--
313-- This implementation traverses the entire window buffer to compute the
314-- minimum whenever we demand it. It performs better than the dequeue based
315-- implementation in @streamly-statistics@ package when the window size is
316-- small (< 30).
317--
318-- If you want to compute the minimum of the entire stream
319-- 'Streamly.Data.Fold.minimum' is much faster.
320--
321-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
322--
323{-# INLINE minimum #-}
324minimum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
325minimum n = fmap (fmap fst) $ range n
327-- | The maximum element in a rolling window.
328--
329-- See the performance related comments in 'minimum'.
330--
331-- If you want to compute the maximum of the entire stream 'Fold.maximum' would
332-- be much faster.
333--
334-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
335--
336{-# INLINE maximum #-}
337maximum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
338maximum n = fmap (fmap snd) $ range n
340-- | Arithmetic mean of elements in a sliding window:
341--
342-- \(\mu = \frac{\sum_{i=1}^n x_{i}}{n}\)
343--
344-- This is also known as the Simple Moving Average (SMA) when used in the
345-- sliding window and Cumulative Moving Avergae (CMA) when used on the entire
346-- stream.
347--
348-- >>> mean = Fold.teeWith (/) sum length
349--
350-- /Space/: \(\mathcal{O}(1)\)
351--
352-- /Time/: \(\mathcal{O}(n)\)
353{-# INLINE mean #-}
354mean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
355mean = Fold.teeWith (/) sum length
357-------------------------------------------------------------------------------
358-- Distribution
359-------------------------------------------------------------------------------
361-- XXX We can use a Windowed classifyWith operation, that will allow us to
362-- express windowed frequency, mode, histograms etc idiomatically.
364-- | Count the frequency of elements in a sliding window.
365--
366-- >>> input = Stream.fromList [1,1,3,4,4::Int]
367-- >>> f = Ring.slidingWindow 4 FoldW.frequency
368-- >>> Stream.fold f input
369-- fromList [(1,1),(3,1),(4,2)]
370--
371-- /Pre-release/
372--
373{-# INLINE frequency #-}
374frequency :: (Monad m, Ord a) => Fold m (a, Maybe a) (Map.Map a Int)
375frequency = Fold.foldl' step Map.empty
377 where
379 decrement v =
380 if v == 1
381 then Nothing
382 else Just (v - 1)
384 step refCountMap (new, mOld) =
385 let m1 = Map.insertWith (+) new 1 refCountMap
386 in case mOld of
387 Just k -> Map.update decrement k m1
388 Nothing -> m1