Binomial.hs

module Binomial
  ( fact
  , comb
  , prob
  , probs
  , weights
  , expected
  , mean
  , squaredDev
  , variance
  , stdDev
  , toStdNormal
  , fromStdNormal
  , example100coins
  ) where

import Data.Ratio (Ratio, (%))
import Data.Foldable (foldl')
import Text.Printf (printf)

data Distrib a = Distrib
  { success :: Ratio a
  , samples :: a
  } deriving (Eq, Show, Read)

ratioToDouble :: (Integral a) => Ratio a -> Double
ratioToDouble = fromRational . toRational

doubleToRatio :: (Integral a) => Double -> Ratio a
doubleToRatio = fromRational . toRational

fact :: (Integral a) => a -> a
fact n = foldl' (*) 1 [1..n]

comb :: (Integral a) => Distrib a -> a -> Ratio a
comb dist k = fact (samples dist) % (fact k * fact (samples dist - k))

prob :: (Integral a) => Distrib a -> a -> Ratio a
prob dist k = correctFactor * succFactor * failFactor
  where
    correctFactor = comb dist k
    succFactor = success dist ^ k
    failFactor = (1 - success dist) ^ (samples dist - k)

probs :: (Integral a) => Distrib a -> [Ratio a]
probs dist = fmap makeProb [0 .. samples dist]
  where
    makeProb k = prob dist k

weights :: (Integral a) => Distrib a -> [Ratio a]
weights dist = fmap makeWeight [0 .. samples dist]
  where
    makeWeight k = k % 1 * prob dist k

expected :: (Integral a) => [Ratio a] -> Ratio a
expected = foldl' (+) 0

mean :: (Integral a) => Distrib a -> Ratio a
mean dist = expected (weights dist)

squaredDev :: (Integral a) => Distrib a -> [Ratio a]
squaredDev dist = fmap makeSquare (weights dist)
  where
    makeSquare p = (p - mean dist) ^ 2

variance :: (Integral a) => Distrib a -> Ratio a
variance dist = expected (squaredDev dist)

stdDev :: (Integral a) => Distrib a -> Double
stdDev dist = sqrt (ratioToDouble (variance dist))

toStdNormal :: (Integral a) => Distrib a -> a -> Double
toStdNormal dist k = ratioToDouble (k % 1 - mean dist) / stdDev dist

fromStdNormal :: (Integral a) => Distrib a -> Double -> a
fromStdNormal dist x = round (doubleToRatio (x * stdDev dist) + mean dist)

-- Example
example100coins :: IO ()
example100coins  = do
  let toDouble p = fromRational p :: Double
  let dist = Distrib { success = 0.5, samples = 100 }
  printf "success = %.3f\n" (toDouble (success dist))
  printf "samples = %i\n" (samples dist)
  printf "mean = %.3f\n" (toDouble (mean dist))
  printf "variance = %.3f\n" (toDouble (variance dist))
  printf "std dev = %.3f\n" (stdDev dist)