mwc-probability

commit ecc640f1d73d140ad6ea7f794e63dd74e5c6daeb
parent 6237bcac11ee776cb6528634efea35c5d4ac5e4d
Author: Marco Zocca <marco.zocca@recordunion.com>
Date:   Tue, 30 Jan 2018 17:34:35 +0100

add negative binomial

Diffstat:
M
CHANGELOG
 | 
3
+++
M
mwc-probability.cabal
 | 
2
+-
M
src/System/Random/MWC/Probability.hs
 | 
10
++++++++++

3 files changed, 14 insertions(+), 1 deletion(-)
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
@@ -87,6 +87,7 @@ module System.Random.MWC.Probability (
   , categorical
   , bernoulli
   , binomial
+  , negativeBinomial
   , multinomial
   , poisson  
 
@@ -295,6 +296,15 @@ 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`.
+--
+-- 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