improve gestion of already parsed vcf for other functions

sfs_tools.py

@@ -102,6 +102,67 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
         print("Pluriallelic sites =", count_pluriall)
     return SFS_values, count_pluriall
 
+
+def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = False, verbose = False):
+
+    """
+    Generates a Site Frequency Spectrum from a gzipped VCF file format.
+
+    Parameters
+    ----------
+    n : int
+        Nb of individuals in sample.
+    vcf_file : str
+        SNPs in VCF file format.
+
+        Used to generate a Site Frequency Spectrum (SFS) from a VCF.
+
+    Returns
+    -------
+    dict
+        Site Frequency Spectrum (SFS)
+
+
+    """
+    
+    if diploid and not folded:
+        n *= 2
+    # initiate SFS_values with a zeros dict
+    SFS_values = dict.fromkeys(range(n),0)
+    count_pluriall = 0
+
+    for CHROM in vcf_dict:
+        for SNP in vcf_dict[CHROM]:
+            snp_genotypes = []
+            allele_counts = {}
+            allele_counts_list = []
+            print(CHROM, SNP)
+            for sample in vcf_dict[CHROM][SNP]["SAMPLES"]:
+                if not phased:
+                    # for UNPHASED data
+                    smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']
+                else:
+                    # for PHASED
+                    smpl_genotype = [int(a) for a in sample.split(':')[0].split('|') if a != '.']
+                nb_alleles = set(smpl_genotype)
+                snp_genotypes += smpl_genotype
+            # skip if all individuals have the same genotype
+            if len(set(snp_genotypes)) == 1:
+                continue
+            for k in set(snp_genotypes):
+                allele_counts[snp_genotypes.count(k)] = k
+                allele_counts_list.append(snp_genotypes.count(k))
+            SFS_values[min(allele_counts_list)-1] += 1
+        # sum pluriall counts for this CHR to the rest
+        count_pluriall += vcf_dict[CHROM]['NB_PLURIALL']
+                    
+    if verbose:
+        print("SFS=", SFS_values)
+    print("Pluriallelic sites =", count_pluriall)
+    
+    return SFS_values, count_pluriall
+
+
 def barplot_sfs(sfs, folded=True, title = "Barplot"):
     sfs_val = []
     n = len(sfs.values())

vcf_utils.py

@@ -105,6 +105,7 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                         entries = {
                             'POS':POS,
                             'CHR':CHROM,
+                            # removed to save some space
                             #'FIELDS':FIELDS,
                             'REF':REF,
                             'ALT':ALT,
@@ -113,13 +114,16 @@ def parse_vcf(vcf_file, phased=False, stop_at=None, chr_starts_with="*"):
                             'SAMPLES':SAMPLES,
                             'QUALITY':QUALITY,
                             'GENOTYPE':GENOTYPE,
-                            'LIKELIHOOD':LIKELIHOOD
+                            'LIKELIHOOD':LIKELIHOOD,
+                            'NB_PLURIALL':pluriall_counts
                         }
                         if CHROM.startswith(chr_starts_with):
                         # keep if chr name starts with filter
                         # default : *, every chr is kept
                             if CHROM not in chrom:
-                                 chrom[CHROM] = {}
+                                # resent when changing chrom
+                                pluriall_counts = 0
+                                chrom[CHROM] = {}
                             chrom[CHROM][POS] = entries
                 byte_line = inputgz.readline()
     end = time.time()