From 257c713458015717294546cd3e2ad4632ae50b55 Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Mon, 14 Feb 2022 10:20:04 +0100
Subject: [PATCH] update plot lib form multiple chrom coverage

---
 __init__.py       |  4 +---
 customgraphics.py | 56 ++++++++++++++++++++++++++++++++++++++++-------
 sfs_tools.py      |  3 +--
 vcf_utils.py      |  9 ++++++++
 4 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/__init__.py b/__init__.py
index 853f264..bd8595f 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,3 +1 @@
-from frst import sfs_tools
-from frst import customgraphics
-from frst import vcf_utils
+from frst import sfs_tools, customgraphics, vcf_utils, sfs_tools
diff --git a/customgraphics.py b/customgraphics.py
index 49a53a5..dadacff 100644
--- a/customgraphics.py
+++ b/customgraphics.py
@@ -9,8 +9,14 @@ FOREST Thomas (thomas.forest@college-de-france.fr)
 import matplotlib.pyplot as plt
 import matplotlib.ticker as ticker
 import numpy as np
+import gc
+import time
+import datetime
+import pandas as pd
+# custom libs
 from frst import vcf_utils
 
+
 def heatmap(data, row_labels=None, col_labels=None, ax=None,
             cbar_kw={}, cbarlabel="", **kwargs):
     """
@@ -145,8 +151,17 @@ def plot_matrix(mat, legend=None, color_scale_type="YlGn", cbarlabel = "qt", tit
     fig.tight_layout()
     plt.show()
 
-def plot(x, y, outfile = None, outfolder = None, ylab=None, xlab=None, title=None):
-    plt.plot(x, y)
+def plot(x, y, outfile = None, outfolder = None, ylab=None, xlab=None,
+         title=None, label = None, show=True):
+    if x:
+        fig, = plt.plot(x, y)
+    else:
+        # x is optional
+        fig, = plt.plot(y)
+    if label:
+        # if legend
+        fig.set_label(label)
+        plt.legend()
     if ylab:
         plt.ylabel(ylab)
     if xlab:
@@ -156,7 +171,8 @@ def plot(x, y, outfile = None, outfolder = None, ylab=None, xlab=None, title=Non
     if outfile:
         plt.savefig(outfile)
     else:
-        plt.show()
+        if show == True:
+            plt.show()
 
 def scatter(x, y, ylab=None, xlab=None, title=None):
     plt.scatter(x, y)
@@ -178,14 +194,38 @@ def barplot(x, y, ylab=None, xlab=None, title=None):
         plt.title(title)
     plt.show()
 
-def plot_chrom_continuity(vcf_entries, chr_id, outfile = None, outfolder = None):
+def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None,
+                          outfolder = None, returned=False, show=True, label=True):
     chr_name = list(vcf_entries.keys())[chr_id]
+    if label:
+        label = chr_name
     chr_entries = vcf_entries[chr_name]
     genotyped_pos = vcf_utils.genotyping_continuity_plot(chr_entries)
-    plot(genotyped_pos[0], genotyped_pos[1], ylab = "genotyped pos.",
-         xlab = "pos. in ref.",
-         title = "Genotyped pos in chr "+str(chr_id+1)+":'"+chr_name+"'",
-         outfile = outfile, outfolder = outfolder)
+    if returned:
+        # if we do not want to plot while executing
+        # useful for storing the x,y coords in a variable for ex.
+        return genotyped_pos
+    else:
+        # to plot on the fly
+        plot(x, y=genotyped_pos[1], ylab = "genotyped pos.",
+             xlab = "pos. in ref.",
+             title = "Genotyped pos in chr "+str(chr_id+1)+":'"+chr_name+"'",
+             outfile = outfile, outfolder = outfolder, show=show, label=label)
+
+def plot_whole_karyotype(recent_variants, mem_clean = False):
+    coords = []
+    for chr in range(len(recent_variants)):
+        x, y = vcf_utils.customgraphics.plot_chrom_continuity(recent_variants, chr_id = chr, show = False, returned = True)
+        coords.append([x, y])
+        if mem_clean:
+            start = time.time()
+            del x
+            del y
+            gc.collect()
+            end = time.time()
+            print("Cleaned mem. in", str(datetime.timedelta(seconds=end - start)))
+    # maybe add a clean of recent_variants in extreme cases, before building the plots
+    return coords
 
 def plot_chrom_coverage(vcf_entries, chr_id):
     chr_name = list(vcf_entries.keys())[chr_id]
diff --git a/sfs_tools.py b/sfs_tools.py
index 2126d71..d06fec4 100755
--- a/sfs_tools.py
+++ b/sfs_tools.py
@@ -19,8 +19,7 @@ import matplotlib.pyplot as plt
 def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False):
 
     """
-    Multiplication de deux nombres entiers.
-    Cette fonction ne sert pas à grand chose.
+    Generates a Site Frequency Spectrum from a gzipped VCF file format.
 
     Parameters
     ----------
diff --git a/vcf_utils.py b/vcf_utils.py
index 24e7213..2ebe2cd 100755
--- a/vcf_utils.py
+++ b/vcf_utils.py
@@ -18,6 +18,9 @@ from frst import customgraphics
 import json 
 import time
 import datetime
+import gc
+import pandas as pd
+
 
 def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
     start = time.time()
@@ -199,6 +202,12 @@ def compute_coverage(vcf_entries, verbose=False):
         coords[1].append(y)
     return coords
 
+def free(obj):
+    """ Free the object and call the garbage collector explicitely
+    """
+    del obj
+    gc.collect()
+
 if __name__ == "__main__":
     # check args
     if len(sys.argv) !=2: