Update SFS plotting function

master
tforest 2024-02-21 22:11:24 +01:00
parent fed1a36d79
commit 44449033db
2 changed files with 28 additions and 13 deletions

View File

@ -200,14 +200,14 @@ def scatter(x, y, ylab=None, xlab=None, title=None):
plt.title(title) plt.title(title)
plt.show() plt.show()
def barplot(x=None, y=None, ylab=None, xlab=None, title=None): def barplot(x=None, y=None, ylab=None, xlab=None, title=None, label=None, xticks = None, width=1):
if x: if x:
x = list(x) x = list(x)
plt.xticks(x) plt.xticks(x)
plt.bar(x, y) plt.bar(x, y, width=width, label=label)
else: else:
x = list(range(len(y))) x = list(range(len(y)))
plt.bar(x, y) plt.bar(x, y, width=width, label=label)
plt.xticks(x) plt.xticks(x)
if ylab: if ylab:
plt.ylabel(ylab) plt.ylabel(ylab)
@ -215,6 +215,9 @@ def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
plt.xlabel(xlab) plt.xlabel(xlab)
if title: if title:
plt.title(title) plt.title(title)
if xticks:
plt.xticks(xticks)
plt.legend()
plt.show() plt.show()
def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None, def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None,

View File

@ -21,6 +21,7 @@ import gzip
import sys import sys
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from frst import customgraphics from frst import customgraphics
import numpy as np
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): strip = False, count_ext = False):
@ -192,7 +193,7 @@ def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = Fal
return SFS_values, count_pluriall return SFS_values, count_pluriall
def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False): def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False, ploidy = 2):
sfs_val = [] sfs_val = []
n = len(sfs.values()) n = len(sfs.values())
sum_sites = sum(list(sfs.values())) sum_sites = sum(list(sfs.values()))
@ -222,7 +223,7 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed =
#terminal case, same for folded or unfolded #terminal case, same for folded or unfolded
if transformed: if transformed:
last_bin = list(sfs.values())[n-1] * n/2 last_bin = list(sfs.values())[n-1] * n/ploidy
else: else:
last_bin = list(sfs.values())[n-1] last_bin = list(sfs.values())[n-1]
sfs_val[-1] = last_bin sfs_val[-1] = last_bin
@ -235,22 +236,33 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed =
#print(sum(sfs_val)) #print(sum(sfs_val))
#build the plot #build the plot
title = title+" (n="+str(len(sfs_val))+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
print("SFS =", sfs)
if folded: if folded:
xlab = "Minor allele frequency" xlab = "Minor allele frequency"
n_title = n
else:
# the spectrum is n-1 long when unfolded
n_title = n+1
title = title+" (n="+str(n_title)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
print("SFS =", sfs)
X_axis = list(sfs.keys())
if transformed: if transformed:
print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val) print("Transformed SFS ( n =",n_title, ") :", sfs_val)
#plt.axhline(y=1/n, color='r', linestyle='-') #plt.axhline(y=1/n, color='r', linestyle='-')
plt.bar([x+0.2 for x in list(sfs.keys())], [1/n]*n, color='r', linestyle='-', width = 0.4, label= "H0 Theoric constant")
else: else:
if normalized: if normalized:
# then plot a theoritical distribution as 1/i # then plot a theoritical distribution as 1/i
expected_y = [1/(2*x+1) for x in list(sfs.keys())] sum_expected = sum([(1/(i+1)) for i,x in enumerate(list(sfs.keys()))])
expected_y = [(1/(i+1))/sum_expected for i,x in enumerate(list(sfs.keys()))]
print(expected_y)
plt.bar([x+0.2 for x in list(sfs.keys())], expected_y, color='r', linestyle='-', width = 0.4, label= "H0 Theoric constant")
print(sum(expected_y)) print(sum(expected_y))
#plt.plot([x for x in list(sfs.keys())], expected_y, color='r', linestyle='-') customgraphics.barplot(x = [x-0.2 for x in X_axis], width=0.4, y= sfs_val, xlab = xlab, ylab = ylab, title = title, label = "H1 Observed spectrum", xticks =list(sfs.keys()) )
#print(expected_y)
customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title)
plt.show() plt.show()
if __name__ == "__main__": if __name__ == "__main__":