Improve plotting with swp2 control curve

swp2.py

@@ -126,10 +126,11 @@ def plot_straight_x_y(x,y):
 
 def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
     theta_scale = True, ax = None, input = None, output = None):
-    my_dpi = 300
+    my_dpi = 500
+    L = full_dict["L"]
     if ax is None:
         # intialize figure
-        my_dpi = 300
+        #my_dpi = 300
         fnt_size = 18
         # plt.rcParams['font.size'] = fnt_size
         fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
@@ -139,16 +140,16 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
         ax1 = ax[1][0,0]
     ax1.set_yscale('log')
     ax1.set_xscale('log')
-    ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
     plot_handles = []
     best_plot = full_dict['all_epochs']['best']
-    p0, = ax1.plot(best_plot[0], best_plot[1], 'o', linestyle = "-",
+    p0, = ax1.plot(best_plot[0], best_plot[1], linestyle = "-",
     alpha=1, lw=2, label = str(best_plot[2])+' brks | Lik='+best_plot[3])
     plot_handles.append(p0)
+    #ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
     for k, plot_Lk in enumerate(full_dict['all_epochs']['plots']):
         plot_Lk = str(full_dict['all_epochs']['plots'][k][3])
         # plt.rcParams['font.size'] = fnt_size
-        p, = ax1.plot(full_dict['all_epochs']['plots'][k][0], full_dict['all_epochs']['plots'][k][1], 'o', linestyle = "--",
+        p, = ax1.plot(full_dict['all_epochs']['plots'][k][0], full_dict['all_epochs']['plots'][k][1], linestyle = "-",
         alpha=1/(k+1), lw=1.5, label = str(full_dict['all_epochs']['plots'][k][2])+' brks | Lik='+plot_Lk)
         plot_handles.append(p)
     if theta_scale:
@@ -160,12 +161,20 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
         # ax1.axvline(x=recent_scale_upper_bound)
     else:
         # years
-        plt.set_xlabel("Time (years)", fontsize=fnt_size)
-        plt.set_ylabel("Individuals (N)", fontsize=fnt_size)
+        if ax is not None:
+            plt.set_xlabel("Time (years)", fontsize=fnt_size)
+            plt.set_ylabel("Effective pop. size (Ne)", fontsize=fnt_size)
+        else:
+            plt.xlabel("Time (years)", fontsize=fnt_size)
+            plt.ylabel("Effective pop. size (Ne)", fontsize=fnt_size)
+    # x_ticks = ax1.get_xticks()
+    # ax1.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
+    # ax1.set_xticklabels([f'{k}\n{k/(mu)}\n{k/(mu)*tgen}' for k in x_ticks], fontsize = fnt_size*0.8)
     # plt.rcParams['font.size'] = fnt_size
     # print(fnt_size, "rcParam font.size=", plt.rcParams['font.size'])
     ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5)
     ax1.set_title(title)
+    breaks = len(full_dict['all_epochs']['plots'])
     if ax is None:
         plt.savefig(title+'_b'+str(breaks)+'.pdf')
     # plot likelihood against nb of breakpoints
@@ -203,8 +212,9 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
     ax3.set_title(title+" AIC")
     if ax is None:
         plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
-    # return plots
-    return ax[0], ax[1]
+    else:
+        # return plots
+        return ax[0], ax[1]
 
 def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, input = None, output = None):
     #scenari = {}
@@ -248,10 +258,11 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
             # do something with the theta without bp and skip the plotting
             N0 = y[0]
             #continue
-        for i in range(len(y)):
-            # divide by N0
-            y[i] = y[i]/N0
-            x[i] = x[i]/N0
+        if theta_scale:
+            for i in range(len(y)):
+                # divide by N0
+                y[i] = y[i]/N0
+                x[i] = x[i]/N0
         top_plots[greatest_likelihood] = x,y,epoch
     plots_likelihoods = list(top_plots.keys())
     for i in range(len(plots_likelihoods)):
@@ -294,9 +305,10 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC
     # 'plots' dict keys: 'best', {epochs}('0', '1',...)
     if input == None:
-        saved_plots = {"S":S, "S0":S0, "L":L, "all_epochs":plots, "Ln_Brks":Ln_Brks,
-                        "AIC_Brks":AIC_Brks, "best_Ln":best_Ln,
-                        "best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb}
+        saved_plots = {"S":S, "S0":S0, "L":L, "mu":mu, "tgen":tgen,
+                       "all_epochs":plots, "Ln_Brks":Ln_Brks,
+                       "AIC_Brks":AIC_Brks, "best_Ln":best_Ln,
+                       "best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb}
     else:
         # if the dict has to be loaded from input
         with open(input, 'r') as json_file:
@@ -304,6 +316,8 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
             saved_plots["S"] = S
             saved_plots["S0"] = S0
             saved_plots["L"] = L
+            saved_plots["mu"] = mu
+            saved_plots["tgen"] = tgen
             saved_plots["all_epochs"] = plots
             saved_plots["Ln_Brks"] = Ln_Brks
             saved_plots["AIC_Brks"] = AIC_Brks
@@ -369,9 +383,10 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
                 for k in range(2, len(y)+2):
                     prop.append(y[k-2] / (k - 1) / sum_theta_i)
                 prop = prop[::-1]
-        # normalise to N0 (N0 of epoch1)
-        for i in range(len(y)):
-            y[i] = y[i]/N0
+        if theta_scale :
+            # normalise to N0 (N0 of epoch1)
+            for i in range(len(y)):
+                y[i] = y[i]/N0
         # x_plot, y_plot = plot_straight_x_y(x, y)
         p = x, y
         # add plot to the list of all plots to superimpose
@@ -394,8 +409,9 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
             y += list(np.repeat(thetas[i], len(group)))
             if epoch == 0:
                 N0 = y[0]
-        for i in range(len(y)):
-            y[i] = y[i]/N0
+        if theta_scale :
+            for i in range(len(y)):
+                y[i] = y[i]/N0
         x_2 = []
         T = 0
         for i in range(len(x)):
@@ -426,10 +442,10 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         json.dump(saved_plots, json_file)
     return saved_plots
 
-def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset = None):
+def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = None, n_ticks = 10, subset = None, theta_scale = False):
     # fig 2 & 3
     if ax is None:
-        my_dpi = 300
+        my_dpi = 500
         fnt_size = 18
         fig2, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
         fig3, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
@@ -442,37 +458,79 @@ def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset =
     lines_fig2 = []
     lines_fig3 = []
     #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
+        for k in range(len(swp2_lines[1])):
+            swp2_lines[1][k] = swp2_lines[1][k]*4*mu
+        # plot_lines = [[swp2_lines[0], swp2_lines[1]]]+plot_lines 
+
+        x2_plot, y2_plot = plot_straight_x_y(swp2_lines[0],swp2_lines[1])
+        p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2')
+        lines_fig2.append(p2)
+        # Plotting (fig 3) which is the same but log scale for x
+        p3, = ax3.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2')
+        lines_fig3.append(p3)
     nb_breaks = len(plot_lines)
     for breaks, plot in enumerate(plot_lines):
         if subset is not None:
             if breaks not in subset :
                 # skip if not in subset
-                if max(subset) > nb_breaks and breaks == nb_breaks-1:
+                if max(subset) > nb_breaks and breaks == nb_breaks:
                     pass
                 else:
                     continue
         x,y=plot
+        # y = [k/(4*mu) for k in y]
+        # x = [k/(mu)*tgen for k in x]
         x2_plot, y2_plot = plot_straight_x_y(x,y)
         p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')
         lines_fig2.append(p2)
         # Plotting (fig 3) which is the same but log scale for x
         p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')
         lines_fig3.append(p3)
-    ax2.set_xlabel("Relative scale", fontsize=fnt_size)
-    ax2.set_ylabel("theta", fontsize=fnt_size)
-    ax2.set_title(title, fontsize=fnt_size)
-    ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
+
+    ax3.axvline(x=500/tgen*mu, linestyle="--")
+    if theta_scale:
+        xlabel = "Theta scaled by N0"
+        ylabel = "Theta scaled by N0"
+    else:
+        xlabel = "Theta scale"
+        ylabel = "Theta"
     if ax is None:
+        # 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)
+        x_ticks = list(plt.xticks())[0]
+        plt.gca().set_xticks(x_ticks)
+        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.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.subplots_adjust(bottom=0.2)  # Adjust the value as needed
         # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
         plt.savefig(title+'_plot2_'+str(len(plot_lines))+'.pdf')
         # close fig2 to save memory
         plt.close(fig2)
+    else:
+        # when ax subplotting is used
+        ax2.set_xlabel(xlabel, fontsize=fnt_size)
+        ax2.set_ylabel(ylabel, fontsize=fnt_size)
+        ax2.set_title(title, fontsize=fnt_size)
+        ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
     ax3.set_xscale('log')
     ax3.set_yscale('log')
-    ax3.set_xlabel("log Relative scale", fontsize=fnt_size)
+    ax3.set_xlabel("time log scale", fontsize=fnt_size)
     ax3.set_ylabel("theta", fontsize=fnt_size)
     ax3.set_title(title, fontsize=fnt_size)
     ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
+    x_ticks = list(ax3.get_xticks())
+    ax3.set_xlim(left=min(x_ticks))
+    ax3.set_xticks(x_ticks)
+    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)
+    plt.text(-0.13, -0.135, 'Coal. time\nGen. 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)
         plt.savefig(title+'_plot3_'+str(len(plot_lines))+'_log.pdf')
@@ -480,7 +538,7 @@ def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset =
         plt.close(fig3)
     return ax
 
-def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
+def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = False, subset = None):
     # multiple fig
     if ax is None:
         # intialize figure 1
@@ -494,7 +552,6 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
         ax1 = ax[0, 0]
         plt.subplots_adjust(wspace=0.3, hspace=0.3)
     plots = []
-
     for epoch, plot in enumerate(plot_lines):
         x,y = plot
         x_plot, y_plot = plot_straight_x_y(x,y)
@@ -527,40 +584,103 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
     # return plots
     return ax
 
-def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True, selected_breaks = []):
+def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = False, selected_breaks = []):
     my_dpi = 300
-
-    save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, output = title+"_plotdata.json")
-    save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, input = title+"_plotdata.json", output = title+"_plotdata.json")
+    saved_plots_dict = save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, output = title+"_plotdata.json")
+    nb_of_epochs = len(saved_plots_dict["all_epochs"]["plots"])
+    print(nb_of_epochs)
+    best_epoch = saved_plots_dict["best_epoch_by_AIC"]
+    save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = nb_of_epochs, input = title+"_plotdata.json", output = title+"_plotdata.json")
 
     with open(title+"_plotdata.json", 'r') as json_file:
         loaded_data = json.load(json_file)
-    # plot page 1 of summary
-    fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
-    # fig1.tight_layout()
-    # Adjust absolute space between the top and bottom rows
-    fig1.subplots_adjust(hspace=0.35)  # Adjust this value based on your requirement
-    # plot page 2 of summary
-    fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
-    # fig2.tight_layout()
-    ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
-                            prop = loaded_data['prop'], title = title, ax = ax1)
-    ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
-                                prop = loaded_data['prop'], title = title, ax = ax1, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks)
-    ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
-                            prop = loaded_data['prop'], title = title, ax = ax2)
-    ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2])
-    fig1.savefig(title+'_combined_p1.pdf')
-    print("Wrote", title+'_combined_p1.pdf')
-    fig2.savefig(title+'_combined_p2.pdf')
-    print("Wrote", title+'_combined_p2.pdf')
+
+    # START OF COMBINED PLOT CODE
+        
+    # # plot page 1 of summary
+    # fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
+    # # fig1.tight_layout()
+    # # Adjust absolute space between the top and bottom rows
+    # fig1.subplots_adjust(hspace=0.35)  # Adjust this value based on your requirement
+    # # plot page 2 of summary
+    # fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
+    # # fig2.tight_layout()
+    # ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
+    #                         prop = loaded_data['prop'], title = title, ax = ax1)
+    # ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
+    #                             prop = loaded_data['prop'], title = title, ax = ax1, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks)
+    # ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
+    #                         prop = loaded_data['prop'], title = title, ax = ax2)
+    # ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2])
+    
+    # fig1.savefig(title+'_combined_p1.pdf')
+    # print("Wrote", title+'_combined_p1.pdf')
+    # fig2.savefig(title+'_combined_p2.pdf')
+    # print("Wrote", title+'_combined_p2.pdf')
+
+    # END OF COMBINED PLOT CODE
+
+
+    # Start of Parsing real swp2 output
+    folder_splitted = folder_path.split("/")
+    swp2_summary = "/".join(folder_splitted[:-2])+'/'+folder_splitted[-3]+".final.summary"
+    swp2_vals = parse_stairwayplot_output_summary(stwplt_out = swp2_summary)
+    swp2_x, swp2_y = swp2_vals[0], swp2_vals[1]
+    # End of Parsing real swp2 output
     plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
                             prop = loaded_data['prop'], title = title, ax = None)
-    plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
-                            prop = loaded_data['prop'], title = title, ax = None)
+    plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], mu = mu, tgen = tgen, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks,
+    # plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], subset=list(range(0,3))+[loaded_data['best_epoch_by_AIC']]+selected_breaks,
+                            prop = loaded_data['prop'], title = title, swp2_lines = [swp2_x, swp2_y], ax = None)
+    plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = None)
+
+    # plt.close(fig1)
+    # plt.close(fig2)
+def parse_stairwayplot_output_summary(stwplt_out, xlim = None, ylim = None, title = "default title", plot = False):
+    #col 5
+    year = []
+    # col 6
+    ne_median = []
+    ne_2_5 = []
+    ne_97_5 = []
+    ne_12_5 = []
+    # col 10
+    ne_87_5 = []
+    with open(stwplt_out, "r") as stwplt_stream:
+        for line in stwplt_stream:
+            ## Line format
+            # mutation_per_site	n_estimation	theta_per_site_median	theta_per_site_2.5%	theta_per_site_97.5%	year	Ne_median	Ne_2.5%	Ne_97.5%	Ne_12.5%	Ne_87.5%
+            if not line.startswith("mutation_per_site"):
+                #not header
+                values = line.strip().split()
+                year.append(float(values[5]))
+                ne_median.append(float(values[6]))
+                ne_2_5.append(float(values[7]))
+                ne_97_5.append(float(values[8]))
+                ne_12_5.append(float(values[9]))
+                ne_87_5.append(float(values[10]))
 
-    plt.close(fig1)
-    plt.close(fig2)
+    vals = [year, ne_median, ne_2_5, ne_97_5, ne_12_5, ne_87_5]
+    if plot :
+        # plot parsed data
+        label = ["Ne median", "Ne 2.5%",	"Ne 97.5%",	"Ne 12.5%",	"Ne 87.5%"]
+        for i in range(1, 5):
+            fig, = plt.plot(year, vals[i], '--', alpha = 0.4)
+            fig.set_label(label[i])
+        #     # last plot is median
+        fig, = plt.plot(year, ne_median, 'r-', lw=2)
+        fig.set_label(label[0])
+        plt.legend()
+        plt.ylabel("Individuals (Ne)")
+        plt.xlabel("Time (years)")
+        if xlim:
+            plt.xlim(xlim)
+        if ylim:
+            plt.ylim(ylim)
+        plt.title(title)
+        plt.show()
+        plt.close()
+    return vals
 
 if __name__ == "__main__":