Add weightedMean

src/Streamly/Statistics.hs

@@ -17,13 +17,17 @@ 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
@@ -33,6 +37,19 @@ 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 +288,17 @@ 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)
+    => Fold m ((a, a), Maybe (a, a)) a
+weightedMean = Fold.lmap func sum
+
+    where
+
+    func ((w, v), Nothing) = (w * v, Nothing)
+    func ((w, v), Just (w1, v1)) = (w * v, Just $ w1 * v1)