@@ -0,0 +1,263 @@
1+module Streamly.Internal.Data.Fold.Window
2+( lmap
3+, length
4+, whole
5+, sum
6+, sumInt
7+, minimum
8+, maximum
9+, range
10+)
11+where
12+
13+import Data.Bifunctor(bimap)
14+import Data.Function ((&))
15+import Data.Maybe (fromMaybe)
16+
17+import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
18+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
19+
20+import Prelude hiding (length, sum, minimum, maximum)
21+
22+import qualified Deque.Strict as DQ
23+import qualified Streamly.Data.Fold as Fold
24+
25+-- $setup
26+-- >>> import Data.Bifunctor(bimap)
27+-- >>> import qualified Streamly.Data.Fold as Fold
28+-- >>> import Prelude hiding (length, sum, minimum, maximum)
29+
30+-------------------------------------------------------------------------------
31+-- Utilities
32+-------------------------------------------------------------------------------
33+
34+-- | Map a function on the incoming as well as outgoing element of a rolling
35+-- window fold.
36+--
37+-- >>> lmap f = Fold.lmap (bimap f (f <$>))
38+--
39+{-# INLINE lmap #-}
40+lmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
41+lmap f = Fold.lmap (bimap f (f <$>))
42+
43+-- | Convert a rolling fold to a normal fold using the entire input stream as a
44+-- single window.
45+--
46+-- >>> whole f = Fold.lmap (\x -> (x, Nothing)) f
47+--
48+{-# INLINE whole #-}
49+whole :: Fold m (a, Maybe a) b -> Fold m a b
50+whole = Fold.lmap (, Nothing)
51+
52+-------------------------------------------------------------------------------
53+-- Sum
54+-------------------------------------------------------------------------------
55+
56+-- XXX Overflow.
57+--
58+-- | The sum of all the elements in a rolling window. The input elements are
59+-- required to be intergal numbers.
60+--
61+-- This was written in the hope that it would be a tiny bit faster than 'sum'
62+-- for 'Integral' values. But turns out that 'sum' is 2% faster than this even
63+-- for intergal values!
64+--
65+-- /Internal/
66+--
67+{-# INLINE sumInt #-}
68+sumInt :: forall m a. (Monad m, Integral a) => Fold m (a, Maybe a) a
69+sumInt = Fold step initial extract
70+
71+ where
72+
73+ initial = return $ Partial (0 :: a)
74+
75+ step s (a, ma) =
76+ return
77+ $ Partial
78+ $ case ma of
79+ Nothing -> s + a
80+ Just old -> s + a - old
81+
82+ extract = return
83+
84+-- XXX Overflow.
85+--
86+-- | Sum of all the elements in a rolling window:
87+--
88+-- Uses Kahan-Babuska-Neumaier style summation for numerical stability of
89+-- floating precision arithmetic.
90+--
91+-- /Space/: \(\mathcal{O}(1)\)
92+--
93+-- /Time/: \(\mathcal{O}(n)\)
94+--
95+{-# INLINE sum #-}
96+sum :: forall m a. (Monad m, Num a) => Fold m (a, Maybe a) a
97+sum = Fold step initial extract
98+
99+ where
100+
101+ initial =
102+ return
103+ $ Partial
104+ $ Tuple'
105+ (0 :: a) -- running sum
106+ (0 :: a) -- accumulated rounding error
107+
108+ step (Tuple' total err) (new, mOld) =
109+ let incr =
110+ case mOld of
111+ -- XXX new may be large and err may be small we may lose it
112+ Nothing -> new - err
113+ -- XXX if (new - old) is large we may lose err
114+ Just old -> (new - old) - err
115+ -- total is large and incr may be small, we may round incr here but
116+ -- we will accumulate the rounding error in err1 in the next step.
117+ total1 = total + incr
118+ -- Accumulate any rounding error in err1
119+ -- XXX In the Nothing case above we may lose err, therefore we
120+ -- should use ((total1 - total) - new) + err here.
121+ -- Or even in the just case if (new - old) is large we may lose
122+ -- err, so we should use ((total1 - total) + (old - new)) + err.
123+ err1 = (total1 - total) - incr
124+ in return $ Partial $ Tuple' total1 err1
125+
126+ extract (Tuple' total _) = return total
127+
128+-- | The number of elements in the rolling window.
129+--
130+{-# INLINE length #-}
131+length :: (Monad m, Num b) => Fold m (a, Maybe a) b
132+length = Fold.foldl' step 0
133+
134+ where
135+
136+ step w (_, Nothing) = w + 1
137+ step w _ = w
138+
139+-------------------------------------------------------------------------------
140+-- Location
141+-------------------------------------------------------------------------------
142+
143+-- Theoretically, we can approximate minimum in a rolling window by using a
144+-- 'powerMean' with sufficiently large negative power.
145+--
146+-- XXX If we need to know the minimum in the window only once in a while then
147+-- we can use linear search when it is extracted and not pay the cost all the
148+-- time.
149+--
150+-- | The minimum element in a rolling window.
151+--
152+-- If you want to compute the minimum of the entire stream Fold.minimum from streamly
153+-- package would be much faster.
154+--
155+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
156+--
157+{-# INLINE minimum #-}
158+minimum :: (Monad m, Ord a) => Fold m (a, Maybe a) a
159+minimum = Fold step initial extract
160+
161+ where
162+
163+ initial = return $ Partial $ Tuple3' (0 :: Int) (0 :: Int)
164+ (mempty :: DQ.Deque (Int, a))
165+
166+ step (Tuple3' i w q) (a, ma) =
167+ case ma of
168+ Nothing ->
169+ return $ Partial $ Tuple3' (i + 1) (w + 1)
170+ (headCheck i q (w + 1) & dqloop (i, a))
171+ Just _ ->
172+ return $ Partial $ Tuple3' (i + 1) w
173+ (headCheck i q w & dqloop (i,a))
174+
175+ {-# INLINE headCheck #-}
176+ headCheck i q w =
177+ case DQ.uncons q of
178+ Nothing -> q
179+ Just (ia', q') ->
180+ if fst ia' <= i - w
181+ then q'
182+ else q
183+
184+ dqloop ia q =
185+ case DQ.unsnoc q of
186+ Nothing -> DQ.snoc ia q
187+ -- XXX This can be improved for the case of `=`
188+ Just (ia', q') ->
189+ if snd ia <= snd ia'
190+ then dqloop ia q'
191+ else DQ.snoc ia q
192+
193+ extract (Tuple3' _ _ q) = return $ snd
194+ $ fromMaybe (0, error "min: Empty stream")
195+ $ DQ.head q
196+
197+-- Theoretically, we can approximate maximum in a rolling window by using a
198+-- 'powerMean' with sufficiently large positive power.
199+--
200+-- | The maximum element in a rolling window.
201+--
202+-- If you want to compute the maximum of the entire stream Fold.maximum from streamly
203+-- package would be much faster.
204+--
205+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
206+--
207+{-# INLINE maximum #-}
208+maximum :: (Monad m, Ord a) => Fold m (a, Maybe a) a
209+maximum = Fold step initial extract
210+
211+ where
212+
213+ initial = return $ Partial $ Tuple3' (0 :: Int) (0 :: Int)
214+ (mempty :: DQ.Deque (Int, a))
215+
216+ step (Tuple3' i w q) (a, ma) =
217+ case ma of
218+ Nothing ->
219+ return $ Partial $ Tuple3' (i + 1) (w + 1)
220+ (headCheck i q (w + 1) & dqloop (i, a))
221+ Just _ ->
222+ return $ Partial $ Tuple3' (i + 1) w
223+ (headCheck i q w & dqloop (i,a))
224+
225+ {-# INLINE headCheck #-}
226+ headCheck i q w =
227+ case DQ.uncons q of
228+ Nothing -> q
229+ Just (ia', q') ->
230+ if fst ia' <= i - w
231+ then q'
232+ else q
233+
234+ dqloop ia q =
235+ case DQ.unsnoc q of
236+ Nothing -> DQ.snoc ia q
237+ -- XXX This can be improved for the case of `=`
238+ Just (ia', q') ->
239+ if snd ia >= snd ia'
240+ then dqloop ia q'
241+ else DQ.snoc ia q
242+
243+ extract (Tuple3' _ _ q) =
244+ return
245+ $ snd
246+ $ fromMaybe (0, error "max: Empty stream")
247+ $ DQ.head q
248+
249+
250+-- | The difference between the maximum and minimum elements of a rolling window.
251+--
252+-- >>> range = Fold.teeWith (-) maximum minimum
253+--
254+-- If you want to compute the range of the entire stream @Fold.teeWith (-)
255+-- Fold.maximum Fold.min@ from the streamly package would be much faster.
256+--
257+-- /Space/: \(\mathcal{O}(n)\) where @n@ is the window size.
258+--
259+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
260+--
261+{-# INLINE range #-}
262+range :: (Monad m, Num a, Ord a) => Fold m (a, Maybe a) a
263+range = Fold.teeWith (-) maximum minimum