Response to comments

haskell · May 12, 2024 · bb03f54 · bb03f54
1 parent b57a2eb
commit bb03f54
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 33 deletions.
diff --git a/System/Random/MWC/Distributions.hs b/System/Random/MWC/Distributions.hs
@@ -376,20 +376,18 @@ binomial :: forall g m . StatefulGen g m
          -> g                 -- ^ Generator
          -> m Int
 {-# INLINE binomial #-}
-binomial nTrials prob gen =
-  if prob == 0.0
-  then return 0
-  else if prob == 1.0
-       then return nTrials
-       else do let (p', flipped) = if prob > 0.5
-                                   then (1.0 - prob, True)
-                                   else (prob, False)
-               ix <- if fromIntegral nTrials * p' < bInvThreshold
-                     then binomialInv nTrials p' gen
-                     else binomialTPE nTrials p' gen
-               if flipped
-                 then return $ nTrials - ix
-                 else return ix
+binomial nTrials prob gen
+  | nTrials <= 0             = pkgError "binomial" "number of trials must be positive"
+  | prob < 0.0 || prob > 1.0 = pkgError "binomial" "probability must be >= 0 and <= 1"
+  | prob == 0.0 = return 0
+  | prob == 1.0 = return nTrials
+  | otherwise = do let (p', flipped) = if prob > 0.5 then (1.0 - prob, True) else (prob, False)
+                   ix <- if fromIntegral nTrials * p' < bInvThreshold
+                         then binomialInv nTrials p' gen
+                         else binomialTPE nTrials p' gen
+                   if flipped
+                     then return $ nTrials - ix
+                     else return ix
 
   where
     binomialTPE n p g =
@@ -424,16 +422,14 @@ binomial nTrials prob gen =
 
           -- Acceptance / rejection comparison
           step5 :: Int -> Double -> m Int
-          step5 ix v = if var <= accept
-                       then if p > 0
-                            then return ix
-                            else return $ n - ix
-                       else hh
-                        where
-                          var = log v
-                          accept = logFactorial bigM + logFactorial (n - bigM) -
-                                   logFactorial ix - logFactorial (n - ix) +
-                                   fromIntegral (ix - bigM) * log (p / q)
+          step5 ix v
+            | var <= accept = return $ if p > 0 then ix else n - ix
+            | otherwise     = hh
+            where
+              var = log v
+              accept = logFactorial bigM + logFactorial (n - bigM) -
+                       logFactorial ix - logFactorial (n - ix) +
+                       fromIntegral (ix - bigM) * log (p / q)
 
           h :: Double -> Double -> m Int
           h u v | -- Triangular region

diff --git a/mwc-random.cabal b/mwc-random.cabal
@@ -1,7 +1,7 @@
 cabal-version:  3.0
 build-type:     Simple
 name:           mwc-random
-version:        0.15.0.3
+version:        0.15.1.0
 license:        BSD-2-Clause
 license-file:   LICENSE
 copyright:      2009, 2010, 2011 Bryan O'Sullivan

diff --git a/tests/props.hs b/tests/props.hs
@@ -1,5 +1,4 @@
 import Control.Monad
-import Control.Monad.Reader
 import Data.Word
 import qualified Data.Vector.Unboxed as U
 
@@ -133,17 +132,17 @@ betaBinomial a b n g = do
 nSamples :: Int
 nSamples = 10000
 
-a, b :: Double
-a = 600.0
-b = 400.0
+alpha, delta :: Double
+alpha = 600.0
+delta = 400.0
 
-n :: Int
-n = 10
+nTrials :: Int
+nTrials = 10
 
 prop_betaBinomialMean :: IO ()
 prop_betaBinomialMean = do
   g <- create
-  ss <- replicateM nSamples $ betaBinomial 600 400 10 g
+  ss <- replicateM nSamples $ betaBinomial alpha delta nTrials g
   let m = fromIntegral (sum ss) / fromIntegral nSamples
-  let x1 = fromIntegral n * a / (a + b)
+  let x1 = fromIntegral nTrials * alpha / (alpha + delta)
   assertBool ("Mean is " ++ show x1 ++ " but estimated as " ++ show m) (abs (m - x1) < 0.001)