Fixing some scaling for swp2 output

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