use binomial

1 files changed, 12 insertions, 5 deletions

diff --git a/bin/cluster.hs b/bin/cluster.hs
index 3eee95a..42e25dc 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -10,19 +10,26 @@ import qualified Data.Map as M
 import GHC.Compact
 import Numeric.Log hiding (sum)
 import Options.Applicative
-import PPL
+import PPL hiding (binom)
 import System.Random (mkStdGen, setStdGen)
 
 cumsum = scanl1 (+)
 
 first f xs = snd . head . filter (f . fst) $ zip xs [0 ..]
 
+stirling 0 = 0
 stirling n = n * log n - n
 
-pois lambda (fromIntegral -> k) = lambda' ** k' * Exp (negate lambda) / Exp (stirling k)
+logFromInt = logFrom . fromIntegral
+
+logFrom = Exp . log
+
+binom :: Int -> Int -> Double -> Log Double
+binom 0 _ _ = 1
+binom n 0 p = logFrom (1 - p) ** logFromInt n
+binom n@(logFromInt -> n') k@(logFromInt -> k') p@(logFrom -> p') = choose * p' ** k' * logFrom (1 - p) ** (n' - k')
   where
-    lambda' = Exp $ log lambda
-    k' = Exp $ log k
+    choose = product [(n' + 1 - logFromInt i) / logFromInt i | i <- [1 .. k]]
 
 -- (infinite) binary trees
 data Tree a = Tree a (Tree a) (Tree a)
@@ -76,7 +83,7 @@ model xs = do
   where
     likelihood ps cnts = product $ zipWith go ps (pairs cnts)
       where
-        go p (c, d) = pois (fromIntegral d * p) c
+        go p (c, d) = binom d c p
 
 -- Tabulate list
 tabulate xs = M.elems $ M.fromListWith (+) [(c, 1) | c <- xs]