Add weightedMean

src/Streamly/Statistics.hs

@@ -17,22 +17,40 @@ module Streamly.Statistics
     , sumInt
     , mean
     , welfordMean
+    , weightedMean
     )
 where
 
 import Control.Monad.IO.Class (MonadIO(..))
 import Data.Bifunctor(bimap)
 import Data.Function ((&))
 import Data.Maybe (fromMaybe)
+import Foreign.Ptr (castPtr, plusPtr)
+import Foreign.Storable
+
 import Streamly.Data.Fold.Tee(Tee(..), toFold)
 import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
-import Streamly.Internal.Data.Tuple.Strict
+import Streamly.Internal.Data.Ring.Foreign (slidingWindow)
+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
 
 import qualified Deque.Strict as DQ
 import qualified Streamly.Data.Fold as Fold
 
 import Prelude hiding (sum, min, max)
 
+instance Storable (Double, Double) where
+    sizeOf _ = 2 * sizeOf (undefined :: Double)
+    alignment _ = 16
+    peek p = do
+        let q = castPtr p
+        a <- peek q
+        b <- peek (q `plusPtr` sizeOf a)
+        return (a, b)
+    poke p (a, b) = do
+        let q = castPtr p
+        poke q a
+        poke (q `plusPtr` sizeOf a) b
+
 -- XXX Make the following more numerically stable. Try to extend welfordMean
 -- XXX method.
 -- XXX     - stdDev
@@ -271,3 +289,12 @@ welfordMean = Fold step initial extract
 {-# INLINE geometricMean #-}
 geometricMean :: forall m a. (MonadIO m, Floating a) => Fold m (a, Maybe a) a
 geometricMean = exp <$> Fold.lmap (bimap log (log <$>)) mean
+
+-- | @weightedMean@ computes the weighted mean of the sample data.
+--
+-- The weights should add up to 1. It uses Kahan-Babuska-Neumaier summation.
+--
+{-# INLINE weightedMean #-}
+weightedMean :: forall m a. (MonadIO m, Fractional a, Storable a)
+    => Int -> Fold m (a, a) a
+weightedMean n = Fold.lmap (uncurry (*)) (slidingWindow n sum)