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.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:
x = list(x)
plt.xticks(x)
plt.bar(x, y)
plt.bar(x, y, width=width, label=label)
else:
x = list(range(len(y)))
plt.bar(x, y)
plt.bar(x, y, width=width, label=label)
plt.xticks(x)
if ylab:
plt.ylabel(ylab)
@ -215,6 +215,9 @@ def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
plt.xlabel(xlab)
if title:
plt.title(title)
if xticks:
plt.xticks(xticks)
plt.legend()
plt.show()
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 matplotlib.pyplot as plt
from frst import customgraphics
import numpy as np
def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = 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
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 = []
n = len(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
if transformed:
last_bin = list(sfs.values())[n-1] * n/2
last_bin = list(sfs.values())[n-1] * n/ploidy
else:
last_bin = list(sfs.values())[n-1]
sfs_val[-1] = last_bin
@ -235,22 +236,33 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed =
#print(sum(sfs_val))
#build the plot
title = title+" (n="+str(len(sfs_val))+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
print("SFS =", sfs)
if folded:
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:
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.bar([x+0.2 for x in list(sfs.keys())], [1/n]*n, color='r', linestyle='-', width = 0.4, label= "H0 Theoric constant")
else:
if normalized:
# 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))
#plt.plot([x for x in list(sfs.keys())], expected_y, color='r', linestyle='-')
#print(expected_y)
customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title)
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()) )
plt.show()
if __name__ == "__main__":