1 files changed, 25 insertions, 2 deletions
diff --git a/bin/cluster.hs b/bin/cluster.hs
index 7984bfd..553f6e0 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -2,6 +2,7 @@
 
 module Main where
 
+import Control.Monad
 import Data.Fixed (mod')
 import Data.Foldable (toList)
 import Data.List
@@ -78,6 +79,25 @@ model xs = do
         pairs [] = []
         pairs _ = error "unexpected number of columns, expecting count/depth pairs"
 
+-- Tabulate list
+tabulate xs = M.elems $ M.fromListWith (+) [(c, 1) | c <- xs]
+
+-- Draws a phylogeny to DOT format
+drawGraph :: FilePath -> [[Double]] -> [Int] -> IO ()
+drawGraph path ps cl = writeFile path $ "digraph{" <> edges <> "}"
+  where
+    edges = concatMap calcEdges [1 .. length tab - 1]
+    calcEdges idx = fmt (parent idx) <> "->" <> fmt idx <> ";"
+    fmt i = "\"" <> intercalate "," (map fmtDbl $ ps' !! i) <> " " <> show (tab !! i) <> "\""
+    tab = tabulate cl
+    parent = flip div 2 . (\x -> x - 1)
+    fmtDbl = show . (/ 10) . fromIntegral . round . (* 1000)
+    ps' = transpose $ map norm ps
+
+    -- Normalise to purity of first node
+    norm :: [Double] -> [Double]
+    norm (x : xs) = x : map (/ x) xs
+
 -- Command line args
 data Options = Options
   { seed :: Int,
@@ -85,7 +105,8 @@ data Options = Options
     mhfrac :: Double,
     input :: FilePath,
     propsPath :: FilePath,
-    clusterPath :: FilePath
+    clusterPath :: FilePath,
+    dotPath :: FilePath
   }
 
 main = run =<< execParser opts
@@ -101,6 +122,7 @@ main = run =<< execParser opts
         <*> argument str (metavar "INPUT")
         <*> argument str (metavar "PROPS")
         <*> argument str (metavar "TREE")
+        <*> strOption (long "dot" <> short 'd' <> help "draw graph of phylogeny in dot format" <> value "" <> metavar "PATH")
 
     probability = eitherReader $ \arg -> case reads arg of
       [(r, "")] -> if r <= 1 && r > 0 then Right r else Left "mhfrac not a valid probability"
@@ -114,8 +136,9 @@ run opts = do
   ((ps, cl), _) <- foldl1' (\a c -> if mml a < mml c then a else c) . take (nsamples opts) <$> mh (mhfrac opts) (model parsed)
   writeFile (propsPath opts) . unlines $ map (intercalate "," . map show) ps
   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
       where
-        tab = M.elems $ M.fromListWith (+) [(c, 1) | c <- cl]
+        tab = tabulate cl
         sum' f = sum . map f