From 25d7ef0858acc8199a2f005da8d323b1c3ba87c7 Mon Sep 17 00:00:00 2001 From: tforest Date: Thu, 29 Feb 2024 00:59:09 +0100 Subject: [PATCH] Improving sfs plotting --- customgraphics.py | 7 +++-- sfs_tools.py | 70 +++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 68 insertions(+), 9 deletions(-) diff --git a/customgraphics.py b/customgraphics.py index 82503a1..0b5cb27 100644 --- a/customgraphics.py +++ b/customgraphics.py @@ -200,11 +200,11 @@ 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, label=None, xticks = None, width=1): +def barplot(x=None, y=None, ylab=None, xlab=None, title=None, label=None, xticks = None, width=1, plot = True): if x: x = list(x) plt.xticks(x) - plt.bar(x, y, width=width, label=label) + plt.bar(x, y, width=width, label=label, color="tab:blue") else: x = list(range(len(y))) plt.bar(x, y, width=width, label=label) @@ -218,7 +218,8 @@ def barplot(x=None, y=None, ylab=None, xlab=None, title=None, label=None, xticks if xticks: plt.xticks(xticks) plt.legend() - plt.show() + if plot: + plt.show() def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None, outfolder = None, returned=False, show=True, label=True, step=1, nb_subplots = None, diff --git a/sfs_tools.py b/sfs_tools.py index 3cfe5bf..481f115 100755 --- a/sfs_tools.py +++ b/sfs_tools.py @@ -23,6 +23,54 @@ import matplotlib.pyplot as plt from frst import customgraphics import numpy as np +def parse_sfs(sfs_file): + """ + Parse a Site Frequency Spectrum (SFS) file and return a masked spectrum. + + This function reads an SFS file, extracts the spectrum data, and applies a mask to it. + The mask excludes specific bins from the spectrum, resulting in a masked SFS. + + Parameters: + - sfs_file (str): The path to the SFS file to be parsed, in dadi's .fs format. + + Returns: + - masked_spectrum (list): A masked SFS as a list of integers. + + Raises: + - FileNotFoundError: If the specified SFS file is not found. + - ValueError: If there are inconsistencies in the file format or data. + + Note: The actual structure of the SFS file is based on dadi's fs format. + """ + try: + with open(sfs_file, 'r') as file: + # Read the first line which contains information about the file + num_individuals, mode, species_name = file.readline().strip().split() + num_individuals = int(num_individuals) + # Read the spectrum data + spectrum_data = list(map(int, file.readline().strip().split())) + # Check if the number of bins in the spectrum matches the expected number + if len(spectrum_data) != num_individuals: + raise ValueError("Error: Number of bins in the spectrum doesn't match the expected number of individuals.") + # Read the mask data + mask_data = list(map(int, file.readline().strip().split())) + + # Check if the size of the mask matches the number of bins in the spectrum + if len(mask_data) != num_individuals: + raise ValueError("Error: Size of the mask doesn't match the number of bins in the spectrum.") + # Apply the mask to the spectrum + masked_spectrum = [spectrum_data[i] for i in range(num_individuals) if not mask_data[i]] + # Error handling + except FileNotFoundError: + print(f"Error: File not found - {sfs_file}") + except ValueError as ve: + print(f"Error: {ve}") + except Exception as e: + print(f"Error: {e}") + # final return of SFS as a list + return masked_spectrum + + def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False, strip = False, count_ext = False): @@ -193,7 +241,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, ploidy = 2): +def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False, ploidy = 2, output = None): sfs_val = [] n = len(sfs.values()) sum_sites = sum(list(sfs.values())) @@ -242,7 +290,8 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = else: # the spectrum is n-1 long when unfolded n_title = n+1 - + original_title = title + # reformat title and add infos title = title+" (n="+str(n_title)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]" print("SFS =", sfs) @@ -252,7 +301,7 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = if transformed: 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") + plt.bar([x+0.2 for x in list(sfs.keys())], [1/n]*n, fill=False, hatch="///", linestyle='-', width = 0.4, label= "H0 Theoric constant") else: if normalized: @@ -260,10 +309,19 @@ def barplot_sfs(sfs, xlab, ylab, folded=True, title = "Barplot", transformed = 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") + plt.bar([x+0.2 for x in list(sfs.keys())], expected_y, fill=False, hatch="///", linestyle='-', width = 0.4, label= "H0 Theoric constant") print(sum(expected_y)) - 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 output is not None: + # if write in a file, don't open the window dynamically + plot = False + else: + plot = True + 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()), plot = plot ) + if output: + plt.savefig(f"{output}/{original_title}_SFS.pdf") + else: + plt.show() + plt.close() if __name__ == "__main__":