mwc-probability

commit 06aed8bb02c63a28e754399d8b4eef9352677840
parent f08cb4f9e10b68226fcaae2f3effb954f0aabd61
Author: Marco Zocca <ocramz@users.noreply.github.com>
Date:   Tue, 30 Jan 2018 17:48:26 +0100

Merge pull request #10 from ocramz/master

Add negative binomial, add References section
Diffstat:
M
CHANGELOG
 | 
3
+++
M
mwc-probability.cabal
 | 
2
+-
M
src/System/Random/MWC/Probability.hs
 | 
18
++++++++++++++++--

3 files changed, 20 insertions(+), 3 deletions(-)
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,5 +1,8 @@
 	# Changelog
 
+	- 2.0.2 (2018-01-30)
+	* Add negative binomial distribution
+
 	- 2.0.1 (2018-01-30)
 	* Add Normal-Gamma and Pareto distributions
 
diff --git a/mwc-probability.cabal b/mwc-probability.cabal
@@ -1,5 +1,5 @@
 name:                mwc-probability
-version:             2.0.1
+version:             2.0.2
 homepage:            http://github.com/jtobin/mwc-probability
 license:             MIT
 license-file:        LICENSE
diff --git a/src/System/Random/MWC/Probability.hs b/src/System/Random/MWC/Probability.hs
@@ -53,6 +53,10 @@
 --
 -- which will be reused throughout all examples.
 -- Note: creating a random generator is an expensive operation, so it should be only performed once in the code (usually in the top-level IO action, e.g `main`).
+--
+-- == References
+--
+-- 1) L.Devroye, Non-Uniform Random Variate Generation, Springer-Verlag, New York, 1986. (Made freely available by the author: http://www.nrbook.com/devroye )
 
 
 module System.Random.MWC.Probability (
@@ -87,6 +91,7 @@ module System.Random.MWC.Probability (
   , categorical
   , bernoulli
   , binomial
+  , negativeBinomial
   , multinomial
   , poisson  
 
@@ -295,6 +300,16 @@ binomial :: PrimMonad m => Int -> Double -> Prob m Int
 binomial n p = fmap (length . filter id) $ replicateM n (bernoulli p)
 {-# INLINABLE binomial #-}
 
+-- | The negative binomial distribution with `n` trials each with "success" probability `p`.
+-- Example X.1.5 in [1].
+--
+-- Note: `n` must be larger than 1 and `p` included between 0 and 1.
+negativeBinomial :: (PrimMonad m, Integral a) => a -> Double -> Prob m Int
+negativeBinomial n p = do
+  y <- gamma (fromIntegral n) ((1-p) / p)
+  poisson y
+{-# INLINABLE negativeBinomial #-}
+
 -- | The multinomial distribution.
 multinomial :: (Foldable f, PrimMonad m) => Int -> f Double -> Prob m [Int]
 multinomial n ps = do
@@ -336,8 +351,7 @@ categorical ps = do
 
 
 -- | The Zipf-Mandelbrot distribution, generated with the rejection
--- sampling algorithm X.6.1 shown in
--- L.Devroye, Non-Uniform Random Variate Generation.
+-- sampling algorithm X.6.1 shown in [1].
 --
 -- The parameter should be positive, but values close to 1 should be
 -- avoided as they are very computationally intensive. The following