random-fu-0.3.0.1: src/Data/Random/Distribution/Binomial.hs
{-# LANGUAGE
MultiParamTypeClasses,
FlexibleInstances, FlexibleContexts,
UndecidableInstances,
BangPatterns
#-}
{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}
module Data.Random.Distribution.Binomial where
import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Beta
import Data.Random.Distribution.Uniform
import Data.Int
import Data.Word
import Numeric.SpecFunctions ( stirlingError )
import Numeric.SpecFunctions.Extra ( bd0 )
import Numeric ( log1p )
-- algorithm from Knuth's TAOCP, 3rd ed., p 136
-- specific choice of cutoff size taken from gsl source
-- note that although it's fast enough for large (eg, 2^10000)
-- @Integer@s, it's not accurate enough when using @Double@ as
-- the @b@ parameter.
integralBinomial :: (Integral a, Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => a -> b -> RVarT m a
integralBinomial = bin 0
where
bin :: (Integral a, Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => a -> a -> b -> RVarT m a
bin !k !t !p
| t > 10 = do
let a = 1 + t `div` 2
b = 1 + t - a
x <- betaT (fromIntegral a) (fromIntegral b)
if x >= p
then bin k (a - 1) (p / x)
else bin (k + a) (b - 1) ((p - x) / (1 - x))
| otherwise = count k t
where
count !k' 0 = return k'
count !k' n | n > 0 = do
x <- stdUniformT
count (if x < p then k' + 1 else k') (n-1)
count _ _ = error "integralBinomial: negative number of trials specified"
integralBinomialCDF :: (Integral a, Real b) => a -> b -> a -> Double
integralBinomialCDF t p x = sum $ map (integralBinomialPDF t p) $ [0 .. x]
-- | The probability of getting exactly k successes in n trials is
-- given by the probability mass function:
--
-- \[
-- f(k;n,p) = \Pr(X = k) = \binom n k p^k(1-p)^{n-k}
-- \]
--
-- Note that in `integralBinomialPDF` the parameters of the mass
-- function are given first and the range of the random variable
-- distributed according to the binomial distribution is given
-- last. That is, \(f(2;4,0.5)\) is calculated by @integralBinomialPDF 4 0.5 2@.
integralBinomialPDF :: (Integral a, Real b) => a -> b -> a -> Double
integralBinomialPDF t p x =
exp $ integralBinomialLogPdf t p x
-- | We use the method given in \"Fast and accurate computation of
-- binomial probabilities, Loader, C\",
-- <http://octave.1599824.n4.nabble.com/attachment/3829107/0/loader2000Fast.pdf>
integralBinomialLogPdf :: (Integral a, Real b) => a -> b -> a -> Double
integralBinomialLogPdf nI pR xI
| p == 0.0 && xI == 0 = 1.0
| p == 0.0 = 0.0
| p == 1.0 && xI == nI = 1.0
| p == 1.0 = 0.0
| xI == 0 = n * log (1-p)
| xI == nI = n * log p
| otherwise = lc - 0.5 * lf
where
n = fromIntegral nI
x = fromIntegral xI
p = realToFrac pR
lc = stirlingError n -
stirlingError x -
stirlingError (n - x) -
bd0 x (n * p) -
bd0 (n - x) (n * (1 - p))
lf = log (2 * pi) + log x + log1p (- x / n)
-- would it be valid to repeat the above computation using fractional @t@?
-- obviously something different would have to be done with @count@ as well...
{-# SPECIALIZE floatingBinomial :: Float -> Float -> RVar Float #-}
{-# SPECIALIZE floatingBinomial :: Float -> Double -> RVar Float #-}
{-# SPECIALIZE floatingBinomial :: Double -> Float -> RVar Double #-}
{-# SPECIALIZE floatingBinomial :: Double -> Double -> RVar Double #-}
floatingBinomial :: (RealFrac a, Distribution (Binomial b) Integer) => a -> b -> RVar a
floatingBinomial t p = fmap fromInteger (rvar (Binomial (truncate t) p))
floatingBinomialCDF :: (CDF (Binomial b) Integer, RealFrac a) => a -> b -> a -> Double
floatingBinomialCDF t p x = cdf (Binomial (truncate t :: Integer) p) (floor x)
floatingBinomialPDF :: (PDF (Binomial b) Integer, RealFrac a) => a -> b -> a -> Double
floatingBinomialPDF t p x = pdf (Binomial (truncate t :: Integer) p) (floor x)
floatingBinomialLogPDF :: (PDF (Binomial b) Integer, RealFrac a) => a -> b -> a -> Double
floatingBinomialLogPDF t p x = logPdf (Binomial (truncate t :: Integer) p) (floor x)
{-# SPECIALIZE binomial :: Int -> Float -> RVar Int #-}
{-# SPECIALIZE binomial :: Int -> Double -> RVar Int #-}
{-# SPECIALIZE binomial :: Integer -> Float -> RVar Integer #-}
{-# SPECIALIZE binomial :: Integer -> Double -> RVar Integer #-}
{-# SPECIALIZE binomial :: Float -> Float -> RVar Float #-}
{-# SPECIALIZE binomial :: Float -> Double -> RVar Float #-}
{-# SPECIALIZE binomial :: Double -> Float -> RVar Double #-}
{-# SPECIALIZE binomial :: Double -> Double -> RVar Double #-}
binomial :: Distribution (Binomial b) a => a -> b -> RVar a
binomial t p = rvar (Binomial t p)
{-# SPECIALIZE binomialT :: Int -> Float -> RVarT m Int #-}
{-# SPECIALIZE binomialT :: Int -> Double -> RVarT m Int #-}
{-# SPECIALIZE binomialT :: Integer -> Float -> RVarT m Integer #-}
{-# SPECIALIZE binomialT :: Integer -> Double -> RVarT m Integer #-}
{-# SPECIALIZE binomialT :: Float -> Float -> RVarT m Float #-}
{-# SPECIALIZE binomialT :: Float -> Double -> RVarT m Float #-}
{-# SPECIALIZE binomialT :: Double -> Float -> RVarT m Double #-}
{-# SPECIALIZE binomialT :: Double -> Double -> RVarT m Double #-}
binomialT :: Distribution (Binomial b) a => a -> b -> RVarT m a
binomialT t p = rvarT (Binomial t p)
data Binomial b a = Binomial a b
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Integer where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Integer) => CDF (Binomial b) Integer where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Integer) => PDF (Binomial b) Integer where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Int where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Int) => CDF (Binomial b) Int where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Int) => PDF (Binomial b) Int where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Int8 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Int8) => CDF (Binomial b) Int8 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Int8) => PDF (Binomial b) Int8 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Int16 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Int16) => CDF (Binomial b) Int16 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Int16) => PDF (Binomial b) Int16 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Int32 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Int32) => CDF (Binomial b) Int32 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Int32) => PDF (Binomial b) Int32 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Int64 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Int64) => CDF (Binomial b) Int64 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Int64) => PDF (Binomial b) Int64 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Word where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Word) => CDF (Binomial b) Word where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Word) => PDF (Binomial b) Word where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Word8 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Word8) => CDF (Binomial b) Word8 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Word8) => PDF (Binomial b) Word8 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Word16 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Word16) => CDF (Binomial b) Word16 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Word16) => PDF (Binomial b) Word16 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Word32 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Word32) => CDF (Binomial b) Word32 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Word32) => PDF (Binomial b) Word32 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance (Floating b, Ord b, Distribution Beta b, Distribution StdUniform b) => Distribution (Binomial b) Word64 where
rvarT (Binomial t p) = integralBinomial t p
instance (Real b, Distribution (Binomial b) Word64) => CDF (Binomial b) Word64 where
cdf (Binomial t p) = integralBinomialCDF t p
instance (Real b, Distribution (Binomial b) Word64) => PDF (Binomial b) Word64 where
pdf (Binomial t p) = integralBinomialPDF t p
logPdf (Binomial t p) = integralBinomialLogPdf t p
instance Distribution (Binomial b) Integer => Distribution (Binomial b) Float where
rvar (Binomial t p) = floatingBinomial t p
instance CDF (Binomial b) Integer => CDF (Binomial b) Float where
cdf (Binomial t p) = floatingBinomialCDF t p
instance PDF (Binomial b) Integer => PDF (Binomial b) Float where
pdf (Binomial t p) = floatingBinomialPDF t p
logPdf (Binomial t p) = floatingBinomialLogPDF t p
instance Distribution (Binomial b) Integer => Distribution (Binomial b) Double where
rvar (Binomial t p) = floatingBinomial t p
instance CDF (Binomial b) Integer => CDF (Binomial b) Double where
cdf (Binomial t p) = floatingBinomialCDF t p
instance PDF (Binomial b) Integer => PDF (Binomial b) Double where
pdf (Binomial t p) = floatingBinomialPDF t p
logPdf (Binomial t p) = floatingBinomialLogPDF t p