From 9388f91dc3609a59c432b1e25aaed1c1d78b4b4c Mon Sep 17 00:00:00 2001
From: Haibao Tang <tanghaibao@gmail.com>
Date: Sun, 28 Apr 2024 20:49:21 -0700
Subject: [PATCH] Add diversity()

---
 jcvi/compara/pedigree.py | 14 ++++++--------
 jcvi/projects/jcvi.py    | 12 ++++++++++++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/jcvi/compara/pedigree.py b/jcvi/compara/pedigree.py
index b4592fb7..249e71b0 100644
--- a/jcvi/compara/pedigree.py
+++ b/jcvi/compara/pedigree.py
@@ -12,7 +12,7 @@
 import networkx as nx
 import numpy as np
 
-from ..apps.base import OptionParser, ActionDispatcher, logger, sh
+from ..apps.base import OptionParser, ActionDispatcher, logger
 from ..formats.base import BaseFile
 from ..graphics.base import set3_n
 
@@ -241,7 +241,7 @@ def pedigree(args):
     p.add_option("--ploidy", default=2, type="int", help="Ploidy")
     p.add_option("--N", default=10000, type="int", help="Number of samples")
     p.add_option("--title", default="", help="Title of the graph")
-    opts, args = p.parse_args(args)
+    opts, args, iopts = p.set_image_options(args)
 
     if len(args) != 1:
         sys.exit(not p.print_help())
@@ -254,12 +254,10 @@ def pedigree(args):
         print(v)
 
     G = ped.to_graph(inb, title=opts.title)
-    dotfile = f"{pedfile}.dot"
-    nx.nx_agraph.write_dot(G, dotfile)
-    pdf_file = dotfile + ".pdf"
-    file_format = pdf_file.split(".")[-1]
-    sh(f"dot -T{file_format} {dotfile} -o {pdf_file}")
-    logger.info("Pedigree graph written to `%s`", pdf_file)
+    A = nx.nx_agraph.to_agraph(G)
+    image_file = f"{pedfile}.{iopts.format}"
+    A.draw(image_file, prog="dot")
+    logger.info("Pedigree graph written to `%s`", image_file)
 
 
 def main():
diff --git a/jcvi/projects/jcvi.py b/jcvi/projects/jcvi.py
index f0a97999..9a828843 100644
--- a/jcvi/projects/jcvi.py
+++ b/jcvi/projects/jcvi.py
@@ -86,6 +86,18 @@ def genomebuild(args):
     savefig(image_name, dpi=iopts.dpi, iopts=iopts)
 
 
+def diversity(args):
+    """
+    %prog diversity pedigree.ped VAR?_srtd.wgs.regions.bed.gz
+
+    Plot diversity composite figure, including:
+    A. Pedigree
+    B. Depth distribution across genomes
+    """
+    p = OptionParser(diversity.__doc__)
+    _, args, iopts = p.set_image_options(args, figsize="14x7")
+
+
 def main():
 
     actions = (