fix bionomial and improve stirling approximation

1 files changed, 6 insertions, 8 deletions

diff --git a/bin/cluster.hs b/bin/cluster.hs
index 67b6e02..7da6810 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -17,19 +17,17 @@ cumsum = scanl1 (+)
 
 first f xs = snd . head . filter (f . fst) $ zip xs [0 ..]
 
-stirling 0 = 0
-stirling n = n * log n - n
+stirling 0 = 1
+stirling 1 = 1
+stirling n = Exp $ n * (log n - 1) + log (sqrt (2 * pi)) + log n / 2
 
 logFromInt = logFrom . fromIntegral
 
 logFrom = Exp . log
 
 binom :: Int -> Int -> Double -> Log Double
-binom n@(logFromInt -> n') k@(logFromInt -> k') p
-  | n == 0 = 1
-  | k == 0 = logFrom (1 - p) ** logFromInt n
-  | k == n = logFrom p ** logFromInt n
-  | otherwise = n `choose` k * logFrom p ** k' * logFrom (1 - p) ** (n' - k')
+binom n@(logFromInt -> n') k@(logFromInt -> k') p =
+  n `choose` k * logFrom p ** k' * logFrom (1 - p) ** (n' - k')
   where
     choose (fromIntegral -> n) (fromIntegral -> k) = stirling n / stirling k / stirling (n - k)
 
@@ -146,7 +144,7 @@ run opts = do
   writeFile (clusterPath opts) . unlines $ map show cl
   when (dotPath opts /= "") $ drawGraph (dotPath opts) ps cl
   where
-    mml ((ps, cl), lik) = sum' (sum' (log . (+ 1))) ps + sum' (log . (+ 1)) tab - sum' stirling tab - ln lik
+    mml ((ps, cl), lik) = sum' (sum' (log . (+ 1))) ps + sum' (log . (+ 1)) tab - sum' (ln . stirling) tab - ln lik
       where
         tab = tabulate cl
         sum' f = sum . map f