fix karyotype plot

customgraphics.py

@@ -152,12 +152,24 @@ def plot_matrix(mat, legend=None, color_scale_type="YlGn", cbarlabel = "qt", tit
     plt.show()
 
 def plot(x, y, outfile = None, outfolder = None, ylab=None, xlab=None,
-         title=None, label = None, show=True):
+         title=None, label = None, show=True, nb_subplots = None, subplot_init = False,
+         subplot_id = None):
+    if subplot_init:
+        # define a certain amount of subplots
+        fig, axs = plt.subplots(nb_subplots)
+    
     if x:
-        fig, = plt.plot(x, y)
+        if nb_subplots:
+            axs[subplot_id].plot(x, y)
+        else:
+            fig, = plt.plot(x, y)
     else:
         # x is optional
-        fig, = plt.plot(y)
+        if nb_subplots:
+            # define a certain amount of subplots
+            axs[subplot_id].plot(y)
+        else:
+            fig, = plt.plot(y)
     if label:
         # if legend
         fig.set_label(label)
@@ -195,37 +207,69 @@ def barplot(x, y, ylab=None, xlab=None, title=None):
     plt.show()
 
 def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None,
-                          outfolder = None, returned=False, show=True, label=True):
+                          outfolder = None, returned=False, show=True, label=True, step=1, nb_subplots = None,
+                          subplot_init = False, subplot_id = None, title = None):
     chr_name = list(vcf_entries.keys())[chr_id]
     if label:
         label = chr_name
+    if not title:
+        title = "Genotyped pos in chr "+str(chr_id+1)+":'"+chr_name+"'"
     chr_entries = vcf_entries[chr_name]
-    genotyped_pos = vcf_utils.genotyping_continuity_plot(chr_entries)
+    genotyped_pos = vcf_utils.genotyping_continuity_plot(chr_entries, step=step)
     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.",
+        plot(x=genotyped_pos[0], 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)
+             title = title,
+             outfile = outfile, outfolder = outfolder, show=show, label=label,
+             nb_subplots = nb_subplots, subplot_init = subplot_init, subplot_id = subplot_id)
 
-def plot_whole_karyotype(recent_variants, mem_clean = False):
+def plot_whole_karyotype(recent_variants, mem_clean = False, step = 1, show = True, min_chr_id = 0,
+                         max_chr_id = None, stacked = False, title = None):
     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)))
+    if max_chr_id :
+        nb_iter = max_chr_id
+    else:
+        nb_iter =  len(recent_variants)
+    if show :
+        iter_start = min_chr_id + 1
+        if not step :
+            step = round(len(recent_variants[list(recent_variants.keys())[min_chr_id]]) / step)
+        if stacked:
+            nb_subplots = nb_iter - min_chr_id
+            subplot_init = True
+        else:
+            nb_subplots = None
+            subplot_init = False
+        vcf_utils.customgraphics.plot_chrom_continuity(recent_variants, chr_id = min_chr_id, show = False, returned = False, step = step,
+                                                       nb_subplots = nb_subplots, subplot_init = subplot_init, subplot_id = min_chr_id)
+    else :
+        iter_start = 0
+    for chr in range(iter_start, nb_iter):
+        if show == False:
+            x, y = vcf_utils.customgraphics.plot_chrom_continuity(recent_variants, chr_id = chr, show = False, returned = True, step = step)
+            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)))
+        else:
+            # if show is enable, use a step
+            step = round(len(recent_variants[list(recent_variants.keys())[chr]]) / 1000)
+            vcf_utils.customgraphics.plot_chrom_continuity(recent_variants, chr_id = chr, show = False, returned = False, step = step, subplot_id = chr)
+        # last case
+    if show == True:
+        vcf_utils.customgraphics.plot_chrom_continuity(recent_variants, chr_id = nb_iter, show = True, returned = False, step = step, subplot_id = nb_iter, title = title)
     # maybe add a clean of recent_variants in extreme cases, before building the plots
-    return coords
+    if show == False:
+        return coords
 
 def plot_chrom_coverage(vcf_entries, chr_id):
     chr_name = list(vcf_entries.keys())[chr_id]

vcf_utils.py

@@ -170,21 +170,23 @@ def build_polymorph_coverage_matrix(entries, noGenotype, diploid=True, na_omit =
         mat = mat / row_sums[:, np.newaxis]
     return mat
 
-def genotyping_continuity_plot(vcf_entries, verbose=False):
+def genotyping_continuity_plot(vcf_entries,
+                                                            verbose=False,
+                                                            step = 1):
     last_pos = int(sorted(list(vcf_entries.keys()))[-1])
     x = 0
     y = 1
     coords = [[], []]
     print(last_pos, "sites to scan")
-    for k, pos in enumerate(range(last_pos)):
+    for k, pos in enumerate(range(0, last_pos, step)):
         if verbose:
             progress = round(k/int(last_pos))*100
             if progress % 10 == 0:
                 print(progress, "%")
         # if pos is genotyped
         if k in vcf_entries:
-            y+=1
-        x+=1
+            y+=1*step
+        x+=1*step
         coords[0].append(x)
         coords[1].append(y)
     return coords