1 files changed, 8 insertions, 7 deletions
diff --git a/bin/cluster.hs b/bin/cluster.hs
index ba9e9f1..e78e261 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -6,7 +6,7 @@ module Main where
 import Control.Monad
 import Data.Fixed (mod')
 import Data.Foldable (toList)
-import Data.List
+import Data.List hiding (group)
 import qualified Data.Map as M
 import GHC.Compact
 import Numeric.Log hiding (sum)
@@ -47,10 +47,10 @@ instance Foldable Tree where
       bftrav [] = mempty
       bftrav ((Tree a l r) : ts) = f a <> bftrav (ts <> [l, r])
 
-{-# INLINE pairs #-}
-pairs (a : b : rs) = (a, b) : pairs rs
-pairs [] = []
-pairs _ = error "unexpected number of columns, expecting pairs"
+{-# INLINE group #-}
+group (a : b : c: d: rs) = (a, b, c, d) : group rs
+group [] = []
+group s = error $ "unexpected number of columns, expecting 4:" <> show s
 
 -- Infinite trees from infinite lists
 -- NB: it's harder to partition a list so that it folds back to
@@ -59,6 +59,7 @@ pairs _ = error "unexpected number of columns, expecting pairs"
 treeFromList (x : xs) = Tree x (treeFromList lpart) (treeFromList rpart)
   where
     (lpart, rpart) = unzip $ pairs xs
+    pairs (a:b:rs) = (a, b) : pairs rs
 
 -- Constrain trees so leaves sum to node value
 normTree :: Tree Double -> Tree Double
@@ -87,9 +88,9 @@ model xs = do
       n = length (head xs) `div` 2
   pure (map (take k) $ take n ps, cls)
   where
-    likelihood ps cnts = product $ zipWith go ps (pairs cnts)
+    likelihood ps cnts = product $ zipWith go ps (group cnts)
       where
-        go p (c, d) = binom d c p
+        go p (c, d, e, f) = binom d c (min 1 (p * fromIntegral (1+e-f) / fromIntegral (1+f)))
 
 -- Tabulate list
 tabulate xs = M.elems $ M.fromListWith (+) [(c, 1) | c <- xs]