From 8d4d36a1f9bdf157bf16b34861ebf130786ddcba Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Tue, 15 Feb 2022 18:02:06 +0100
Subject: [PATCH] better gestion of memory while parsing vcf

---
 customgraphics.py |  7 ++++---
 vcf_utils.py      | 23 +++++++++++++----------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/customgraphics.py b/customgraphics.py
index 78743c1..727414f 100644
--- a/customgraphics.py
+++ b/customgraphics.py
@@ -237,8 +237,8 @@ def plot_whole_karyotype(recent_variants, mem_clean = False, step = 1, show = Tr
         nb_iter =  len(recent_variants) -1
     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 step == "auto" :
+            step = round(len(recent_variants[list(recent_variants.keys())[min_chr_id]]) / 1000)
         if stacked:
             nb_subplots = nb_iter - min_chr_id
             subplot_init = True
@@ -262,7 +262,8 @@ def plot_whole_karyotype(recent_variants, mem_clean = False, step = 1, show = Tr
                 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)
+            if step == "auto":
+                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:
diff --git a/vcf_utils.py b/vcf_utils.py
index 466e5a1..2894a5a 100755
--- a/vcf_utils.py
+++ b/vcf_utils.py
@@ -43,6 +43,7 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                 #  # every snp line, not comment or header
                 if not line.startswith("##") and not line.startswith("#"):
                     FIELDS = line.split("\t")
+                    # when line is parsed, delete it to save some memory
                     CHROM = FIELDS[0]
                     POS = int(FIELDS[1])
                     if stop_at:
@@ -52,8 +53,8 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                     REF = FIELDS[3].split(",")
                     # ALT is col 5 of VCF
                     ALT = FIELDS[4].split(",")
-                    FORMAT = line.split("\t")[8:9]
-                    SAMPLES = line.split("\t")[9:]
+                    FORMAT = FIELDS[8:9]
+                    SAMPLES = FIELDS[9:]
                     QUALITY = float(FIELDS[5])
                     INFO = FIELDS[7]
                     INFOS = {}
@@ -68,7 +69,7 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                     # 1 : missing
                     # 2 : deletion among REF
                     # 3 : deletion among ALT
-                    if "./.:." in line \
+                    if "./.:." in SAMPLES \
                        or len(ALT[0]) > 1 \
                        or len(REF[0]) > 1:
                         # sites that are not kept
@@ -104,7 +105,7 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                         entries = {
                             'POS':POS,
                             'CHR':CHROM,
-                            'FIELDS':FIELDS,
+                            #'FIELDS':FIELDS,
                             'REF':REF,
                             'ALT':ALT,
                             'FORMAT':FORMAT,
@@ -173,20 +174,22 @@ def build_polymorph_coverage_matrix(entries, noGenotype, diploid=True, na_omit =
 def genotyping_continuity_plot(vcf_entries,
                                                             verbose=False,
                                                             step = 1):
-    last_pos = int(sorted(list(vcf_entries.keys()))[-1])
+    genotyped_pos = sorted(list(vcf_entries.keys()))
+    last_pos = genotyped_pos[-1]
     x = 0
     y = 1
     coords = [[], []]
-    print(last_pos, "sites to scan")
-    for k, pos in enumerate(range(0, last_pos, step)):
+    print("Chr. len. =", last_pos, "bp \t ; nb. SNPs =", len(genotyped_pos[::step]))
+    for k, pos in enumerate(genotyped_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*step
-        x+=1*step
+        # if k in vcf_entries:
+        #     y=k*step
+        y+=1*step
+        x=pos
         coords[0].append(x)
         coords[1].append(y)
     return coords