From 2dd51d45b01d8efa90f33018f785bd181f8a7d73 Mon Sep 17 00:00:00 2001
From: Haibao Tang <tanghaibao@gmail.com>
Date: Sun, 28 Apr 2024 14:57:36 -0700
Subject: [PATCH] Check pedigree contains valid nodes

---
 jcvi/compara/pedigree.py | 63 +++++++++++++++++++++++++++-------------
 1 file changed, 43 insertions(+), 20 deletions(-)

diff --git a/jcvi/compara/pedigree.py b/jcvi/compara/pedigree.py
index 62c9a305..b68bfefe 100644
--- a/jcvi/compara/pedigree.py
+++ b/jcvi/compara/pedigree.py
@@ -70,23 +70,42 @@ def __init__(self, pedfile: str):
                 mom = mom if mom != "0" else None
                 s = Sample(name, dad, mom)
                 self[s.name] = s
+        self._check()
 
-    def to_graph(self, inbreeding: Dict[str, SampleInbreeding]) -> nx.DiGraph:
+    def _check(self):
+        """
+        # Check if all nodes are assigned, including the roots
+        """
+        terminal_nodes = set()
+        for s in self:
+            dad, mom = self[s].dad, self[s].mom
+            if dad and dad not in self:
+                terminal_nodes.add(dad)
+            if mom and mom not in self:
+                terminal_nodes.add(mom)
+        for s in terminal_nodes:
+            logger.info("Adding %s to pedigree", s)
+            self[s] = Sample(s, None, None)
+        self.terminal_nodes = terminal_nodes
+
+    def to_graph(self, inbreeding_dict: Dict[str, SampleInbreeding]) -> nx.DiGraph:
         """
         Convert the pedigree to a graph.
         """
-        G = nx.DiGraph()
+        graph_styles = {"splines": "curved"}
+        edge_styles = {"color": "lightslategray", "arrowhead": "none"}
+        G = nx.DiGraph(**graph_styles)
         for s in self:
             dad, mom = self[s].dad, self[s].mom
             if dad:
-                G.add_edge(dad, s, color="lightslategray", arrowhead="none")
+                G.add_edge(dad, s, **edge_styles)
             if mom:
-                G.add_edge(mom, s, color="lightslategray", arrowhead="none")
+                G.add_edge(mom, s, **edge_styles)
         # Map colors to terminal nodes
         terminal_nodes = [s for s in self if self[s].is_terminal]
         colors = dict(zip(terminal_nodes, set3_n(len(terminal_nodes))))
         for s in self:
-            inb = inbreeding[s]
+            inb = inbreeding_dict[s]
             label = s
             if inb.mean_inbreeding > 0.01:
                 label += f"\n(F={inb.mean_inbreeding:.2f})"
@@ -98,16 +117,20 @@ def to_graph(self, inbreeding: Dict[str, SampleInbreeding]) -> nx.DiGraph:
                 fillcolor += ":white"
             else:
                 fillcolor = fillcolor.rsplit(";", 1)[0]
-            G._node[s]["label"] = label
-            G._node[s]["shape"] = "circle"
-            G._node[s]["fixedsize"] = "true"
-            G._node[s]["width"] = "0.6"
-            G._node[s]["height"] = "0.6"
-            G._node[s]["style"] = "wedged"
-            G._node[s]["fillcolor"] = fillcolor
-            G._node[s]["color"] = "none"
-            G._node[s]["fontsize"] = "10"
-            G._node[s]["fontname"] = "Helvetica"
+            node_styles = {
+                "color": "none",
+                "fillcolor": fillcolor,
+                "fixedsize": "true",
+                "fontname": "Helvetica",
+                "fontsize": "10",
+                "height": "0.6",
+                "label": label,
+                "shape": "circle",
+                "style": "wedged",
+                "width": "0.6",
+            }
+            for k, v in node_styles.items():
+                G._node[s][k] = v
         return G
 
 
@@ -206,13 +229,13 @@ def calculate_inbreeding(
     return results
 
 
-def inbreeding(args):
+def pedigree(args):
     """
-    %prog inbreeding pedfile
+    %prog pedigree pedfile
 
-    Calculate inbreeding coefficients from a pedigree file.
+    Plot pedigree and calculate pedigree coefficients from a pedigree file.
     """
-    p = OptionParser(inbreeding.__doc__)
+    p = OptionParser(pedigree.__doc__)
     p.add_option("--ploidy", default=2, type="int", help="Ploidy")
     p.add_option("--N", default=10000, type="int", help="Number of samples")
     opts, args = p.parse_args(args)
@@ -237,7 +260,7 @@ def inbreeding(args):
 
 
 def main():
-    actions = (("inbreeding", "calculate inbreeding coefficients"),)
+    actions = (("pedigree", "Plot pedigree and calculate inbreeding coefficients"),)
     p = ActionDispatcher(actions)
     p.dispatch(globals())