Fixing some scaling for swp2 output

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

swp2.py

@@ -176,7 +176,7 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
     ax1.set_title(title)
     breaks = len(full_dict['all_epochs']['plots'])
     if ax is None:
-        plt.savefig(title+'_'+str(breaks+1)+'_epochs.pdf')
+        plt.savefig(title+'_best_'+str(breaks+1)+'_epochs.pdf')
     # plot likelihood against nb of breakpoints
     if ax is None:
         fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
@@ -299,7 +299,15 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     # number of monomorphic sites
     S0 = L-S
     # print("SFS", SFS_stored)
-    # print("S", S, "L", L, "S0=", S0)
+    print("S", S, "L", L, "S0=", S0)
+
+    my_n = len(SFS_stored)*2
+    print("n=",my_n)
+    an = 1
+    for i in range(2, my_n):
+        an +=1.0/i
+    
+    print("an=", an, "theta_w", S/an, "theta_w_p_site", (S/an)/L)
     # compute Ln
     Ln = log_facto(S+S0) - log_facto(S0) + np.log(float(S0)/(S+S0)) * S0
     for xi in range(0, len(SFS_stored)):
@@ -413,6 +421,25 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         cumul = val+cumul
     prop = prop_cumul
 
+
+    # print("raw stairs", plots[3])
+
+
+    # ###########
+
+    # time = []
+    # for k in plots[0][0]:
+    #     k = int(k)
+    #     dt = 2.0/(k*(k-1))
+    #     time.append(2.0/(k*(k-1)))
+
+    # Ne = []
+    # for values in plots:
+    #     Ne.append(np.array(values[1])/(4*mu))
+    # print(time)
+    # print(Ne[3])
+
+    
     lines_fig2 = []
     for epoch, theta in best_epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
@@ -423,24 +450,33 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
             x += group[::-1]
             y += list(np.repeat(thetas[i], len(group)))
             if epoch == 0:
-                N0 = y[0]
+                # watterson theta
+                theta_w = y[0]
         if theta_scale :
             for i in range(len(y)):
                 y[i] = y[i]/N0
+        for i in range(len(y)):
+            y[i] = y[i]/(4*mu)
         x_2 = []
         T = 0
         for i in range(len(x)):
             x[i] = int(x[i])
         # compute the times as: theta_k / (k*(k-1))
         for i in range(0, len(x)):
-            T += y[i] / (x[i]*(x[i]-1))
+            T += y[i]*2 / (x[i]*(x[i]-1))
             x_2.append(T)
         # Save plotting (fig 2)
-        x_2 = [0]+x_2
-        y = [y[0]]+y
+        # x_2 = [0]+x_2
+        # y = [y[0]]+y
         # x2_plot, y2_plot = plot_straight_x_y(x_2, y)
         p2 = x_2, y
         lines_fig2.append(p2)
+    # print("breaks=", epoch, "scaled_theta", lines_fig2[10])
+    # print(lines_fig2[3][1][0]/(4*mu))
+    # print(np.array(lines_fig2[3][1])/lines_fig2[3][1][0])
+    # print("size list y=", len(lines_fig2[3][1]))
+    #exit(0)
+
     if input == None:
         saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
                         "prop":prop}
@@ -458,9 +494,9 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
     return saved_plots
 
 def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = None, n_ticks = 10, subset = None, theta_scale = False):
-    recent_limit_years = 500
+    recent_limit_years = 100
     # recent limit in coal. time
-    recent_limit = recent_limit_years/tgen*mu
+    recent_limit = recent_limit_years/tgen
     # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
     nb_epochs = len(plot_lines)
     # fig 2 & 3
@@ -480,9 +516,9 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
     #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     if swp2_lines:
         for k in range(len(swp2_lines[0])):
-            swp2_lines[0][k] = swp2_lines[0][k]/tgen*mu
+            swp2_lines[0][k] = swp2_lines[0][k]/tgen
         for k in range(len(swp2_lines[1])):
-            swp2_lines[1][k] = swp2_lines[1][k]*4*mu
+            swp2_lines[1][k] = swp2_lines[1][k]
         # x2_plot, y2_plot = plot_straight_x_y(swp2_lines[0],swp2_lines[1])
         x2_plot, y2_plot = swp2_lines[0], swp2_lines[1]
         p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black")
@@ -508,14 +544,14 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
 
                 # skip the base 0 points x_plot[0:3]
                 t_max_below_limit = 0
-                t_min_below_limit = 1
+                t_min_below_limit = recent_limit
                 recent_change = False
                 for t in x[1:]:
                     if t <= recent_limit:
                         recent_change = True
                         t_max_below_limit = max(t_max_below_limit, t)
                         t_min_below_limit = min(t_min_below_limit, t)
-                        Ne_max_below_limit = y[x.index(t_max_below_limit)]
+                        Ne_max_below_limit = y[min(x.index(t_max_below_limit)+1, len(y)-1)]
                         Ne_min_below_limit = y[x.index(t_min_below_limit)]
                 if recent_change:
                     print(f"\n{breaks} breaks ; This is below the recent limit of {recent_limit_years} years:\n",
@@ -547,6 +583,8 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
             lines_fig3.append(p3)
     # put the vertical line of the "recent" time limit
     ax3.axvline(x=recent_limit, linestyle="--")
+    ax3.axvline(x=recent_limit/2, linestyle="--", color="green")
+
     if theta_scale:
         xlabel = "Theta scaled by N0"
         ylabel = "Theta scaled by N0"
@@ -557,26 +595,36 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
         # if not ax, then use the plt syntax, not ax...
         plt.xlabel(xlabel, fontsize=fnt_size)
         plt.ylabel(ylabel, fontsize=fnt_size)
-        #plt.xlim(left=0)
+        plt.gca().set_xlim(0, recent_limit * 3)
+        if recent_change:
+            plt.ylim(Ne_min_below_limit/3, Ne_max_below_limit *3)
+        else:
+            plt.ylim(y2_plot[0]/3, y2_plot[0])
+        # plt.ylim(0, max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05)))
+        #plt.xlim(0, recent_limit * 3)
         #xlim_val = plt.gca().get_xlim()
-        #x_ticks = list(plt.xticks())[0]
-        plt.xlim(min(min_x,min(swp2_lines[0])), max(max(swp2_lines[0]), max_x))
-        x_ticks = list(plt.gca().get_xticks())
-        plt.gca().set_xticks(x_ticks)
+        x_ticks = list(plt.xticks())[0]
+        # plt.xlim(min(min_x,min(swp2_lines[0])), max(max(swp2_lines[0]), max_x))
+        # x_ticks = list(plt.gca().get_xticks())
+        # plt.gca().set_xticks(x_ticks)
         # plt.xticks(x_ticks)
         # plt.gca().set_xlim(xlim_val)
-        plt.gca().set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+        # plt.gca().set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+        plt.gca().set_xticklabels([f'{k:.1f}\n{k*tgen:.1f}' for k in x_ticks], fontsize = fnt_size*0.5)
+
         # rescale y to effective pop size
         # ylim_val = plt.gca().get_ylim()
-        plt.ylim(min(min_y,min(swp2_lines[1])), max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05)))        
-        y_ticks = list(plt.yticks())[0]
-        plt.gca().set_yticks(y_ticks)
+        # plt.ylim(min(min_y,min(swp2_lines[1])), max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05)))        
+        # y_ticks = list(plt.yticks())[0]
+        # plt.gca().set_yticks(y_ticks)
         # plt.gca().set_ylim(ylim_val)
-        plt.yticks(y_ticks)
-        plt.gca().set_yticklabels([f'{k/(4*mu):.0e}' for k in y_ticks], fontsize = fnt_size*0.5)
-        plt.title(title, fontsize=fnt_size)
-        plt.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
-        plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+        # plt.yticks(y_ticks)
+        # plt.gca().set_yticklabels([f'{k/(4*mu):.0e}' for k in y_ticks], fontsize = fnt_size*0.5)
+        # plt.title(title, fontsize=fnt_size)
+        # plt.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
+        # # plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+        plt.text(-0.13, -0.135, 'Gen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+
         plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
         plt.savefig(title+'_plotB_'+str(nb_epochs)+'_epochs.pdf')
         # close fig2 to save memory
@@ -594,16 +642,30 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
     ax3.set_xscale('log')
     ax3.set_yscale('log')
     # Scale the x-axis
-    x_ticks = list(ax3.get_xticks())
-    ax3.set_xticks(x_ticks)
-    ax3.set_xlim(min(min(x_ticks), min(swp2_lines[0])), max(max_x, max(swp2_lines[0])))
-    ax3.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+    # x_ticks = list(ax3.get_xticks())
+    # ax3.set_xticks(x_ticks)
+    # x_ticks = [i for i in range(0.1,max(max_x, max(swp2_lines[0]))), ]
+    # ax3.set_xticks(x_ticks)
+    ax3.set_xlim(0.1, max(max_x, max(swp2_lines[0])))
+    x_ticks = ax3.get_xticks()
+    # ax3.set_xlim(min(min(x_ticks), min(swp2_lines[0])), max(max_x, max(swp2_lines[0])))
+    # ax3.set_xlim(1, max(max_x, max(swp2_lines[0])))
+    # ax3.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+    # ax3.set_xticklabels([f'{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+    ax3.set_xticklabels([f'{k:.0e}\n{k*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+
     # rescale y to effective pop size
-    y_ticks = list(ax3.get_yticks())
-    ax3.set_yticks(y_ticks)
-    ax3.set_ylim(min(min(y_ticks), min(swp2_lines[1])), max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5)))
-    ax3.set_yticklabels([f'{k/(4*mu):.0e}' for k in y_ticks], fontsize = fnt_size*0.5)
-    plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+    # y_ticks = list(ax3.get_yticks())
+    # ax3.set_yticks(y_ticks)
+    # ax3.set_ylim(min(min(y_ticks), min(swp2_lines[1])), max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5)))
+    # ax3.set_ylim(1, max(max_y, max(swp2_lines[1])))
+    ax3.set_ylim(1, max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5)))
+
+    # ax3.set_yticklabels([f'{k/(4*mu):.0e}' for k in y_ticks], fontsize = fnt_size*0.5)
+    # plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+    # plt.text(-0.13, -0.135, 'Gen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+    plt.text(-0.13, -0.085, 'Gen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
+    
     plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
     if ax is None:
         # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
@@ -638,7 +700,7 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale =
         x,y = plot
         x_plot, y_plot = plot_straight_x_y(x,y)
         p, = ax1.plot(x_plot, y_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')
-
+        print("breaks=", breaks, "theta0", y[0])
         # add plot to the list of all plots to superimpose
         plots.append(p)
     x_ticks = x