From b8e4d958878c6db4bc1e8dd4f694d65330f2de89 Mon Sep 17 00:00:00 2001
From: Haibao Tang <tanghaibao@gmail.com>
Date: Fri, 12 Apr 2024 13:04:44 -0700
Subject: [PATCH] Add plot

---
 jcvi/apps/pedigree.py | 43 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/jcvi/apps/pedigree.py b/jcvi/apps/pedigree.py
index 3dc0a716..170e0ae7 100644
--- a/jcvi/apps/pedigree.py
+++ b/jcvi/apps/pedigree.py
@@ -9,9 +9,10 @@
 from random import sample
 from typing import Optional
 
+import networkx as nx
 import numpy as np
 
-from ..apps.base import OptionParser, ActionDispatcher
+from ..apps.base import OptionParser, ActionDispatcher, logger, sh
 from ..formats.base import BaseFile
 
 
@@ -54,6 +55,19 @@ def __init__(self, pedfile: str):
                 s = Sample(name, dad, mom)
                 self[s.name] = s
 
+    def to_graph(self) -> nx.DiGraph:
+        """
+        Convert the pedigree to a graph.
+        """
+        G = nx.DiGraph()
+        for s in self:
+            dad, mom = self[s].dad, self[s].mom
+            if dad:
+                G.add_edge(dad, s)
+            if mom:
+                G.add_edge(mom, s)
+        return G
+
 
 class GenotypeCollection(dict):
     """
@@ -155,8 +169,33 @@ def inbreeding(args):
         print(f"{s}\t{mean_inbreeding:.4f}\t{std_inbreeding:.4f}\t{dosage}")
 
 
+def plot(args):
+    """
+    %prog plot pedfile
+
+    Plot the pedigree with graphviz.
+    """
+    p = OptionParser(plot.__doc__)
+    _, args = p.parse_args(args)
+
+    if len(args) != 1:
+        sys.exit(not p.print_help())
+
+    (pedfile,) = args
+    ped = Pedigree(pedfile)
+    G = ped.to_graph()
+    dotfile = f"{pedfile}.dot"
+    nx.drawing.nx_pydot.write_dot(G, dotfile)
+    pdf_file = dotfile + ".pdf"
+    sh(f"dot -Tpdf {dotfile} -o {pdf_file}")
+    logger.info("Pedigree graph written to `%s`", pdf_file)
+
+
 def main():
-    actions = (("inbreeding", "calculate inbreeding coefficients"),)
+    actions = (
+        ("inbreeding", "calculate inbreeding coefficients"),
+        ("plot", "plot the pedigree with graphviz"),
+    )
     p = ActionDispatcher(actions)
     p.dispatch(globals())