diff --git a/sfs_tools.py b/sfs_tools.py index 139f260..20fff26 100755 --- a/sfs_tools.py +++ b/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()) diff --git a/vcf_utils.py b/vcf_utils.py index ccf2b46..049a6fa 100755 --- a/vcf_utils.py +++ b/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()