diff --git a/vcf_to_sfs.py b/vcf_to_sfs.py index 33403c5..64614d8 100755 --- a/vcf_to_sfs.py +++ b/vcf_to_sfs.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 """ +FOREST Thomas (thomas.forest@college-de-france.fr) + Caution : At the moment for gzipped files only. ARGS @@ -26,6 +28,8 @@ if diploid and not folded: # initiate SFS_values with a zeros dict SFS_values = dict.fromkeys(range(n),0) +polycount = 0 + with gzip.open(sys.argv[1], "rb") as inputgz: line = inputgz.readline() genotypes = [] @@ -34,23 +38,60 @@ with gzip.open(sys.argv[1], "rb") as inputgz: line = line.decode('utf-8').strip() # every snp line, not comment or header if not line.startswith("##") and not line.startswith("#"): + FIELDS = line.split("\t") + # REF is col 4 of VCF + 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:] snp_genotypes = [] allele_counts = {} + allele_counts_list = [] + # SKIP the SNP if : + # 1 : missing + # 2 : deletion among REF + # 3 : deletion among ALT + if "./.:." in line \ + or len(ALT[0]) > 1 \ + or len(REF[0]) > 1: + line = inputgz.readline() + continue for sample in SAMPLES: - # for UNPHASED data - smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.'] - + 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 + # if set(snp_genotypes) > 2: + # polyallelic = set(snp_genotypes) + # else: + # polyallelic = False + polyallelic = len(ALT) + ##print(REF, ALT, snp_genotypes) # skip if all individuals have the same genotype if len(set(snp_genotypes)) == 1: 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 : - SFS_values[min(allele_counts.keys())-1] += 1 + #allele_counts_list = list(allele_counts.keys()) + ##print("ALC", allele_counts_list, "POLY", polyallelic, ALT) + # for al in range(polyallelic-1): + # SFS_values[min(allele_counts_list)-1] += 1/len(ALT) + # allele_counts_list.remove(min(allele_counts_list)) + if len(ALT) == 1: + SFS_values[min(allele_counts_list)-1] += 1 + else: + for al in range(polyallelic-1): + SFS_values[min(allele_counts_list)-1] += 1/len(ALT) + allele_counts_list.remove(min(allele_counts_list)) + polycount += 1 line = inputgz.readline() print(SFS_values) +print(polycount)