------------------------------------------------------------------------
-- |
-- Module      :  Math.Statistics
-- Copyright   :  (c) Johan Tibell 2008
-- License     :  BSD3-style (see LICENSE)
--
-- Maintainer  :  me@willsewell.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Common statistical functions.
--
------------------------------------------------------------------------

module Math.Statistics
    ( mean
    , median
    , stddev
    , variance
    ) where

import Data.List (sort)

-- | Numerically stable mean.
mean :: Floating a => [a] -> a
mean :: [a] -> a
mean = a -> Int -> [a] -> a
forall a. Floating a => a -> Int -> [a] -> a
go a
0 Int
0
    where
      go :: Floating a => a -> Int -> [a] -> a
      go :: a -> Int -> [a] -> a
go a
m Int
_ []     = a
m
      go a
m Int
n (a
x:[a]
xs) = a -> Int -> [a] -> a
forall a. Floating a => a -> Int -> [a] -> a
go (a
m a -> a -> a
forall a. Num a => a -> a -> a
+ (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m) a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

-- | Median.
median :: (Floating a, Ord a) => [a] -> a
median :: [a] -> a
median [a]
xs
    | Int -> Bool
forall a. Integral a => a -> Bool
odd Int
n     = [a] -> a
forall a. [a] -> a
head ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [a]
xs'
    | Bool
otherwise = [a] -> a
forall a. Floating a => [a] -> a
mean ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
2 ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
i [a]
xs'
    where
      i :: Int
i   = ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      xs' :: [a]
xs' = [a] -> [a]
forall a. Ord a => [a] -> [a]
sort [a]
xs
      n :: Int
n   = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs

-- | Standard deviation.
stddev :: Floating a => [a] -> a
stddev :: [a] -> a
stddev [a]
xs = a -> a
forall a. Floating a => a -> a
sqrt (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ [a] -> a
forall a. Floating a => [a] -> a
variance [a]
xs

-- | Numerically stable sample variance.
variance :: Floating a => [a] -> a
variance :: [a] -> a
variance [a]
xs = a -> Int -> a -> [a] -> a
forall a. Floating a => a -> Int -> a -> [a] -> a
go a
0 Int
0 a
0 [a]
xs a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    where
      go :: Floating a => a -> Int -> a -> [a] -> a
      go :: a -> Int -> a -> [a] -> a
go a
_ Int
_ a
s [] = a
s
      go a
m Int
n a
s (a
x:[a]
xs') = a -> Int -> a -> [a] -> a
forall a. Floating a => a -> Int -> a -> [a] -> a
go a
nm (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
delta a -> a -> a
forall a. Num a => a -> a -> a
* (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
nm)) [a]
xs'
         where
           delta :: a
delta = a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m
           nm :: a
nm = a
m a -> a -> a
forall a. Num a => a -> a -> a
+ a
delta a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)