update sfs plot and vcf parsing

__init__.py

@@ -1 +1,2 @@
-from frst import sfs_tools, customgraphics, vcf_utils, sfs_tools, stats_sfs
+from frst import sfs_tools, customgraphics, vcf_utils, sfs_tools, stats_sfs, dependences
+

bam_utils.py

@@ -0,0 +1,17 @@
+from frst.dependences import pybam
+
+def parse_bam(bam_file):
+    cov = {}
+    cov_val = 0
+    for alignment in pybam.read(bam_file):
+        #pos 1 based (see Pybam doc)
+        start_pos = alignment.sam_pos1
+        end_pos = alignment.sam_pos1 + alignment.sam_block_size
+        if start_pos not in cov:
+            cov[start_pos] = 0
+        if end_pos not in cov:
+            cov[end_pos] = 0
+        # add a weight on the start and lower on the end
+        cov[start_pos] = cov_val +1
+        cov[end_pos] = cov_val -1
+    return cov

customgraphics.py

@@ -200,8 +200,12 @@ def scatter(x, y, ylab=None, xlab=None, title=None):
         plt.title(title)
     plt.show()
 
-def barplot(x, y, ylab=None, xlab=None, title=None):
-    plt.bar(x, y)
+def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
+    if x:
+        plt.bar(x, y)
+    else:
+        x = list(range(len(y)))
+        plt.bar(x, y)
     if ylab:
         plt.ylabel(ylab)
     if xlab:

dependences/__init__.py

@@ -0,0 +1 @@
+from frst.dependences import pybam

sfs_tools.py

@@ -10,13 +10,19 @@ ARGS
 
 standalone usage : vcf_to_sfs.py VCF.gz nb_indiv
 
+TODO
+_____
+Externalize sfs transforms in a function
+Rectify SFS comp in parsed funct.
+
 """
 
 import gzip
 import sys
 import matplotlib.pyplot as plt
 
-def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False):
+def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False,
+                 strip = False, count_ext = False):
 
     """
     Generates a Site Frequency Spectrum from a gzipped VCF file format.
@@ -41,7 +47,13 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
     if diploid and not folded:
         n *= 2
     # initiate SFS_values with a zeros dict
-    SFS_values = dict.fromkeys(range(n),0)
+    # if strip:
+    #     # "[1" removes the 0 bin
+    #     # "n-1]" crop the last bin (n or n/2 for folded)
+    #     SFS_dim = [1, n-1]
+    # else:
+    SFS_dim = [0, n+1]
+    SFS_values = dict.fromkeys(range(SFS_dim[1]),0)
     count_pluriall = 0
     with gzip.open(vcf_file, "rb") as inputgz:
         line = inputgz.readline()
@@ -81,13 +93,20 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
                     nb_alleles = set(smpl_genotype)
                     snp_genotypes += smpl_genotype
                 # skip if all individuals have the same genotype
-                if len(set(snp_genotypes)) == 1:
-                    line = inputgz.readline()
-                    continue
+                # if len(set(snp_genotypes)) == 1:
+                #     if folded or (folded == False and snp_genotypes.count(1) == 0) :
+                #         line = inputgz.readline()
+                #         continue         
                 for k in set(snp_genotypes):
                     allele_counts[snp_genotypes.count(k)] = k
                     allele_counts_list.append(snp_genotypes.count(k))
-                if folded and len(ALT) >= 2:
+                #print(allele_counts_list)
+                if  len(set(snp_genotypes)) == 1 or allele_counts_list[0] == allele_counts_list[1]:
+                    # If only heterozygous sites 0/1; skip the site (equivalent to n bin or n/2 bin for folded)
+                    # skip if all individuals have the same genotype
+                    line = inputgz.readline()
+                    continue     
+                if len(ALT) >= 2:
                     #pass
                     count_pluriall +=1
                     # TODO - work in progress
@@ -95,10 +114,19 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
                     #     SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
                     #     allele_counts_list.remove(min(allele_counts_list))
                 else:
-                    SFS_values[min(allele_counts_list)-1] += 1
+                    if folded:
+                        SFS_values[min(allele_counts_list)-SFS_dim[0]] += 1
+                    else :
+                        # if unfolded, count the Ones (ALT allele)
+                        #print(snp_genotypes, snp_genotypes.count(1))
+                        SFS_values[snp_genotypes.count(1)-SFS_dim[0]] += 1
+            # all the parsing is done, change line
             line = inputgz.readline()
             if verbose:
                 print("SFS=", SFS_values)
+        if strip:
+            del SFS_values[0]
+            del SFS_values[n]
         print("Pluriallelic sites =", count_pluriall)
     return SFS_values, count_pluriall
 
@@ -163,20 +191,39 @@ def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = Fal
     return SFS_values, count_pluriall
 
 
-def barplot_sfs(sfs, folded=True, title = "Barplot"):
+def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False):
     sfs_val = []
     n = len(sfs.values())
     for k in range(1, n):
         ksi = list(sfs.values())[k-1]
         # k+1 because k starts from 0
-        if folded:
-            sfs_val.append(ksi * k * (n - k))
+        # if folded:
+        #     # ?check if 2*n or not?
+        #     sfs_val.append(ksi * k * (2*n - k))
+        # else:
+        #     if transformed:
+        #         sfs_val.append(ksi * k)
+        #     else:
+        #         sfs_val.append(ksi)
+        if transformed:
+            if folded:
+                sfs_val.append(ksi * k * (2*n - k))
+            else:
+                sfs_val.append(ksi * k)
         else:
-            sfs_val.append(ksi * k)
+             sfs_val.append(ksi)
+            
     #terminal case, same for folded or unfolded
-    sfs_val.append(list(sfs.values())[n-1] * n)
+    if transformed:
+        sfs_val.append(list(sfs.values())[n-1] * n)
+    else:
+         sfs_val.append(list(sfs.values())[n-1])
     #build the plot
     title = title+" [folded="+str(folded)+"]"
+    if ylab:
+        plt.ylabel(ylab)
+    if xlab:
+        plt.xlabel(xlab)
     plt.title(title)
     plt.bar([i+1 for i in sfs.keys()], sfs_val)
     plt.show()
@@ -189,5 +236,5 @@ if __name__ == "__main__":
 
     # PARAM : Nb of indiv
     n = int(sys.argv[2])
-    sfs = sfs_from_vcf(n, sys.argv[1], folded = True, diploid = True, phased = False)
+    sfs = sfs_from_vcf(n, sys.argv[1], folded = True, diploid = True, phased = False, strip = True)
     print(sfs)

stats_sfs.py

@@ -1,4 +1,3 @@
-from frst import vcf_to_sfs
 
 import math
 

vcf_utils.py

@@ -214,6 +214,25 @@ def free(obj):
     del obj
     gc.collect()
 
+def random_indiv_from_vcfs(vcf_files_list):
+    """
+    Either for two VCFs of two indiv or two samplenames from the same VCF.
+    """
+    if len(vcf_files_list)>1:
+        # two invid from two VCFs
+        for vcf in vcf_files_list:
+            with open(vcf) as vcf_stream:
+                for line in vcf_stream:
+                    if line.startswith('#'):
+                        continue
+                    genotype  = line.split("\t")[-1].strip()
+                    if genotype.split(":")[1] != '0':
+                        genotype_list = genotype.split(':')[0].split("/")
+                        print(genotype_list)
+    else:
+        # two samplename to randomly pick
+        pass
+    
 if __name__ == "__main__":
     # check args
     if len(sys.argv) !=2: