diff --git a/sfs_tools.py b/sfs_tools.py index 225273a..22d576e 100755 --- a/sfs_tools.py +++ b/sfs_tools.py @@ -192,10 +192,10 @@ def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = Fal return SFS_values, count_pluriall -def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = False): +def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False): sfs_val = [] n = len(sfs.values()) - print("n =", n) + sum_sites = sum(list(sfs.values())) for k, ksi in sfs.items(): #ksi = list(sfs.values())[k-1] # k+1 because k starts from 0 @@ -209,24 +209,30 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = # sfs_val.append(ksi) if transformed: if folded: - #sfs_val.append(ksi * k * (2*n - k)) - sfs_val.append(((k*(2*n - k)) / (2*n))*ksi) + val = ((k*(2*n - k)) / (2*n))*(ksi) else: - sfs_val.append(ksi * k) + val = ksi * k else: - sfs_val.append(ksi) + val = ksi + sfs_val.append(val) #terminal case, same for folded or unfolded if transformed: - sfs_val[-1] = list(sfs.values())[n-1] * n + last_bin = list(sfs.values())[n-1] * n/2 else: - sfs_val[-1] = list(sfs.values())[n-1] + last_bin = list(sfs.values())[n-1] + sfs_val[-1] = last_bin + if normalized: + ylab = "Fraction of SNPs" + sum_val = sum(sfs_val) + for k, sfs_bin in enumerate(sfs_val): + sfs_val[k] = sfs_bin / sum_val #build the plot - title = title+" [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]" + title = title+" (n="+str(len(sfs_val)+1)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]" print("SFS =", sfs) if transformed: - print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val) - customgraphics.barplot(x = sfs.keys(), y= sfs_val, xlab = xlab, ylab = ylab, title = title) + print("Transformed SFS ( n =",len(sfs_val)+1, ") :", sfs_val) + customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title) plt.show() if __name__ == "__main__":