mwc-probability

commit e16b6078340bc2fb5f840a0edf75971368cd560b
parent aaeb8f599c1aa4c70c2971244ef4edd63e779344
Author: Jared Tobin <jared@jtobin.ca>
Date:   Fri, 10 Apr 2015 14:31:33 +1200

Add categorical generator.

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

2 files changed, 13 insertions(+), 4 deletions(-)
diff --git a/mwc-probability.cabal b/mwc-probability.cabal
@@ -1,5 +1,5 @@
 name:                mwc-probability
-version:             0.3.2.0
+version:             0.3.3.0
 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
@@ -5,6 +5,7 @@ module System.Random.MWC.Probability (
   , Prob(..)
   , uniform
   , uniformR
+  , discreteUniform
   , categorical
   , standard
   , normal
@@ -66,11 +67,11 @@ uniformR :: (PrimMonad m, Variate a) => (a, a) -> Prob m a
 uniformR r = Prob $ QMWC.uniformR r
 {-# INLINE uniformR #-}
 
-categorical :: PrimMonad m => [a] -> Prob m a
-categorical cs = do
+discreteUniform :: PrimMonad m => [a] -> Prob m a
+discreteUniform cs = do
   j <- uniformR (0, length cs - 1)
   return $ cs !! j
-{-# INLINE categorical #-}
+{-# INLINE discreteUniform #-}
 
 standard :: PrimMonad m => Prob m Double
 standard = Prob MWC.Dist.standard
@@ -149,3 +150,11 @@ poisson l = Prob $ genFromTable table where
   table = tablePoisson l
 {-# INLINE poisson #-}
 
+categorical :: PrimMonad m => [Double] -> Prob m Int
+categorical ps = do
+  xs <- multinomial 1 ps
+  case xs of
+    [x] -> return x
+    _   -> error "categorical: invalid return value"
+{-# INLINE categorical #-}
+