Improving sfs plotting

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,

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__":