src/Phylomemy/DOT.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 module Phylomemy.DOT where
   4
   5 -- import Debug.Pretty.Simple (pTraceShow, pTraceShowM)
   6 import Control.Applicative (Applicative (..))
   7 import Control.Monad (Monad (..), foldM_, forM_, mapM_, when, zipWithM_)
   8 import Control.Monad.Trans.Class qualified as MT
   9 import Control.Monad.Trans.Reader qualified as MT
  10 import Control.Monad.Trans.Writer.CPS qualified as MT
  11 import Data.Bool (otherwise)
  12 import Data.ByteString.Builder qualified as BS
  13 import Data.ByteString.Short qualified as BSh
  14 import Data.Eq (Eq (..))
  15 import Data.Foldable (foldMap', toList)
  16 import Data.Function (on, ($), (&), (.))
  17 import Data.Functor ((<&>))
  18 import Data.Int (Int)
  19 import Data.List qualified as List
  20 import Data.Map.Strict qualified as Map
  21 import Data.Maybe (Maybe (..), maybe)
  22 import Data.Monoid (Monoid (..))
  23 import Data.Ord (Ord (..))
  24 import Data.Semigroup (Semigroup (..))
  25 import Data.Set (Set)
  26 import Data.Set qualified as Set
  27 import Data.String (String)
  28 import Data.Text.Short qualified as TS
  29 import Data.Tree qualified as Tree
  30 import Data.Tuple (snd)
  31 import GHC.Real (floor)
  32 import Numeric (showFFloat)
  33 import Numeric.Probability
  34 import Text.Printf (printf)
  35 import Text.Show (Show (..))
  36 import Prelude
  37
  38 import Phylomemy.Indexation
  39 import Phylomemy.TemporalMatching
  40
  41 -- | @(`dotMaximalSpanningForest` msf)@
  42 -- returns a graph of the given `MaximalSpanningForest`
  43 -- in [DOT](https://graphviz.org/doc/info/lang.html) format.
  44 dotMaximalSpanningForest ::
  45   forall range cluster.
  46   Show range =>
  47   Show cluster =>
  48   Ord range =>
  49   Ord cluster =>
  50   ShowHuman range =>
  51   ShowHuman cluster =>
  52   MaximalSpanningForest range cluster ->
  53   BS.Builder
  54 dotMaximalSpanningForest msf = runDOT do
  55   let sortedMSF = msf & List.sortBy (compare `on` mstNodeRangeCluster . Tree.rootLabel)
  56   let rangeToMSTToClusters :: range :-> {-mstI-} Int :-> (Maybe Similarity, Set cluster) =
  57         let merge = Map.unionWith \(minSimil, x) (_minSimil, y) -> (minSimil, Set.union x y)
  58         in Map.unionsWith
  59             merge
  60             [ Map.fromListWith merge $
  61               case mst of
  62                 Tree.Node MSTNode{mstNodeRangeCluster = (rootR, rootC)} ts ->
  63                   (rootR, Map.singleton mstI (minSimil, Set.singleton rootC))
  64                     : [ (range, Map.singleton mstI (minSimil, Set.singleton cluster))
  65                       | MSTNode{mstNodeRangeCluster = (range, cluster)} <- ts & List.concatMap toList
  66                       ]
  67             | (mstI, mst) <- sortedMSF & List.zip [1 :: Int ..]
  68             , let minSimil = mstMinimalSimilarity mst
  69             ]
  70   let showSimilarity (s :: Similarity) = showFFloat (Just 2) (s & runProbability & fromRational @Double) ""
  71   dotComments [(BS.stringUtf8 $ show $ rangeToMSTToClusters & Map.map Map.keys)]
  72   -- pTraceShow ("num of nodes", Map.size nodeToBranch, "num of branches", Map.size msf) $
  73   dotLine "digraph g"
  74   dotBlock do
  75     dotLine "splines=\"ortho\""
  76     indexFrom1M_ (rangeToMSTToClusters & Map.toList) \(srcR, mstToClusters) srcRI -> do
  77       let srcRB = "r" <> BS.intDec srcRI
  78       dotLine $ "subgraph cluster_" <> srcRB
  79       dotBlock do
  80         dotComments ["Create a node for the range " <> srcRB]
  81         dotNode
  82           srcRB
  83           [ ("shape", "box")
  84           , ("label", builderQuotedString (showHuman srcR))
  85           , ("color", "gray")
  86           , ("style", "filled")
  87           , ("fillcolor", "gray")
  88           ]
  89         dotLine "color=gray"
  90         dotBlock do
  91           dotLine "rank=same"
  92           dotComments ["Create the cluster nodes within the range " <> srcRB]
  93           forM_ (mstToClusters & Map.toList) \(mstI, (minSimil, clusters)) -> do
  94             indexFrom1M_ (clusters & toList) \srcC srcCI -> do
  95               dotNodeCluster
  96                 srcRI
  97                 mstI
  98                 srcCI
  99                 [ ("label", builderQuotedString $ showHuman srcC <> "\nT" <> printf "%03d" mstI <> maybe "" (("\n" <>) . showSimilarity) minSimil)
 100                 , ("style", "filled")
 101                 , minSimil & maybe ("", "") (\s -> ("fillcolor", 1 + ((floor (runProbability s * 10)) `mod` 10) & BS.intDec))
 102                 , ("colorscheme", "ylorrd9")
 103                 , ("shape", "box")
 104                 ]
 105           dotComments ["Horizontally align the cluster nodes within the same range"]
 106           let row =
 107                 [ (mstI, clusterI)
 108                 | (mstI, (_minSimil, clusters)) <- mstToClusters & Map.toList
 109                 , clusterI <- [1 .. Set.size clusters]
 110                 ]
 111           case row of
 112             [] -> return ()
 113             c@(firstTI, firstCI) : cs -> do
 114               dotEdges
 115                 [srcRB, srcRB <> "t" <> BS.intDec firstTI <> "c" <> BS.intDec firstCI]
 116                 [ ("style", "invis")
 117                 ]
 118               cs & (`foldM_` c) \(srcTI, srcCI) dst@(dstTI, dstCI) -> do
 119                 dotEdgesCluster
 120                   [(srcRI, srcTI, srcCI), (srcRI, dstTI, dstCI)]
 121                   [ ("weight", "10")
 122                   , ("style", "invis")
 123                   ]
 124                 return dst
 125     indexFrom1M_ sortedMSF \mst mstI -> do
 126       dotComments ["Create the edges of the MST " <> BS.intDec mstI]
 127       -- pTraceShowM (mstI, List.length (Tree.flatten mst))
 128       let loop (Tree.Node MSTNode{mstNodeRangeCluster = src} dsts) = do
 129             forM_ dsts \dstNode@(Tree.Node MSTNode{mstNodeRangeCluster = dst, mstNodeSimilarity = simil} _) -> do
 130               -- let similB = BS.stringUtf8 $ showFFloat (Just 2) (simil & runProbability & fromRational @Double) ""
 131               let indexRangeCluster (r, c) =
 132                     ( 1 + Map.findIndex r rangeToMSTToClusters
 133                     , mstI
 134                     , 1 + Set.findIndex c (rangeToMSTToClusters Map.! r Map.! mstI & snd)
 135                     )
 136               dotEdgesCluster
 137                 [ indexRangeCluster src
 138                 , indexRangeCluster dst
 139                 ]
 140                 [ ("constraint", "false")
 141                 , ("color", (floor (runProbability simil * 10)) `mod` 10 & BS.intDec)
 142                 , ("colorscheme", "ylorrd9")
 143                 , -- , ("label", similB)
 144                   ("fontcolor", "blue")
 145                 , ("dir", "both")
 146                 , ("arrowhead", "dot")
 147                 , ("arrowtail", "dot")
 148                 ]
 149               loop dstNode
 150       loop mst
 151     dotRanges rangeToMSTToClusters
 152
 153 dotRanges :: range :-> a -> DOT
 154 dotRanges rangeTo = do
 155   dotComments ["Vertically align range nodes"]
 156   let rangeLinks =
 157         [ "r" <> BS.intDec srcRI
 158         | srcRI <- [1 .. Map.size rangeTo]
 159         ]
 160   when (1 < List.length rangeLinks) do
 161     dotEdges rangeLinks [("weight", "10"), ("style", "invis")]
 162
 163 dotNodeCluster :: Int -> Int -> Int -> [(BS.Builder, BS.Builder)] -> DOT
 164 dotNodeCluster r t c = dotNode ("r" <> BS.intDec r <> "t" <> BS.intDec t <> "c" <> BS.intDec c)
 165
 166 dotEdgesCluster :: [(Int, Int, Int)] -> [(BS.Builder, BS.Builder)] -> DOT
 167 dotEdgesCluster rtc =
 168   dotEdges
 169     [ ("r" <> BS.intDec r <> "t" <> BS.intDec t <> "c" <> BS.intDec c)
 170     | (r, t, c) <- rtc
 171     ]
 172
 173 -- Alternative to `Show` to get a more human-readable `String`.
 174 class ShowHuman a where
 175   showHuman :: a -> String
 176 instance ShowHuman (Set.Set Root) where
 177   showHuman a =
 178     mconcat (List.intersperse " & " (as <&> TS.unpack))
 179     where
 180       as = a & Set.toList <&> unNgram . rootLabel
 181 instance ShowHuman Int where
 182   showHuman = show
 183
 184 type DOT = MT.ReaderT BSh.ShortByteString (MT.Writer BS.Builder) ()
 185
 186 indexFrom1M_ :: Applicative m => [a] -> (a -> Int -> m b) -> m ()
 187 indexFrom1M_ xs f = zipWithM_ f xs [1 :: Int ..]
 188
 189 runDOT :: DOT -> BS.Builder
 190 runDOT = MT.execWriter . (`MT.runReaderT` {-indent-} "")
 191
 192 dotBlock :: DOT -> DOT
 193 dotBlock s = do
 194   dotLine "{"
 195   () <- MT.withReaderT ("  " <>) s
 196   dotLine "}"
 197
 198 dotLine :: BS.Builder -> DOT
 199 dotLine s = do
 200   indent <- MT.ask
 201   MT.lift $ MT.tell $ BS.shortByteString indent <> s <> "\n"
 202
 203 dotComments :: [BS.Builder] -> DOT
 204 dotComments = mapM_ \c -> dotLine $ "// " <> c
 205
 206 dotEdges :: [BS.Builder] -> [(BS.Builder, BS.Builder)] -> DOT
 207 dotEdges names as = dotLine $ mconcat (List.intersperse " -> " names) <> builderAttrs as
 208
 209 dotNode :: BS.Builder -> [(BS.Builder, BS.Builder)] -> DOT
 210 dotNode name as = dotLine $ name <> builderAttrs as
 211
 212 builderAttrs :: [(BS.Builder, BS.Builder)] -> BS.Builder
 213 builderAttrs as
 214   | List.null as = ""
 215   | otherwise = "[" <> mconcat (List.intersperse "," [k <> "=" <> v | (k, v) <- as, BS.toLazyByteString k /= ""]) <> "]"
 216
 217 builderQuotedString :: String -> BS.Builder
 218 builderQuotedString cs = BS.charUtf8 '"' <> foldMap' escape cs <> BS.charUtf8 '"'
 219   where
 220     escape '\\' = BS.charUtf8 '\\' <> BS.charUtf8 '\\'
 221     escape '\"' = BS.charUtf8 '\\' <> BS.charUtf8 '\"'
 222     escape c = BS.charUtf8 c