Save all epochs in JSON

swp2.py

         val += np.log(i)
     return val
 
-def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False):
+def parse_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False):
     with open(stwp_theta_file, "r") as swp_file:
         # Read the first line
         line = swp_file.readline()
         x.append(list(t.values())[time])
         x.append(list(t.values())[time])
     x.append(list(t.values())[len(t.values())-1])
-    # if relative_theta_scale:
-    #     # rescale
-    #     #N0 = y[0]
-    #     # for i in range(len(y)):
-    #     #     # divide by N0
-    #     #     y[i] = y[i]/N0
-    #     #     x[i] = x[i]/N0
-    return x,y,likelihood,thetas,sfs,L
 
+    return x,y,likelihood,thetas,sfs,L
 
 def plot_k_epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title", theta_scale = True):
     scenari = {}
     for file_name in os.listdir(folder_path):
         if os.path.isfile(os.path.join(folder_path, file_name)):
             # Perform actions on each file
-            x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
+            x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                              tgen = tgen,
                                      mu = mu, relative_theta_scale = theta_scale)
             if x == 0 or y == 0:
     x_1.append(x[-1])
     return x_1, y_1
 
+def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title",
+    theta_scale = True, ax = None, input = None, output = None):
+    #scenari = {}
+    cpt = 0
+    epochs = {}
+    for file_name in os.listdir(folder_path):
+        breaks = 0
+        cpt +=1
+        if os.path.isfile(os.path.join(folder_path, file_name)):
+            x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                                                             tgen = tgen,
+                                                             mu = mu, relative_theta_scale = theta_scale)
+            SFS_stored = sfs
+            L_stored = L
+            while not (x == 0 and y == 0):
+                if breaks not in epochs.keys():
+                    epochs[breaks] = {}
+                epochs[breaks][likelihood] = x,y
+                breaks += 1
+                x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                                                                 tgen = tgen,
+                                                                  mu = mu, relative_theta_scale = theta_scale)
+            if x == 0:
+                # last break did not work, then breaks = breaks-1
+                breaks -= 1
+    print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n")
+    print(cpt, "theta file(s) have been scanned.")
+    my_dpi = 300
+    if ax is None:
+        # intialize figure
+        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)
+    else:
+        fnt_size = 12
+        # plt.rcParams['font.size'] = fnt_size
+        ax1 = ax[1][0,0]
+    ax1.set_yscale('log')
+    ax1.set_xscale('log')
+    ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
+    brkpt_lik = []
+    top_plots = {}
+    for epoch, scenari in epochs.items():
+        # sort starting by the smallest -log(Likelihood)
+        best10_scenari = (sorted(list(scenari.keys())))[:10]
+        greatest_likelihood = best10_scenari[0]
+        # store the tuple breakpoints and likelihood for later plot
+        brkpt_lik.append((epoch, greatest_likelihood))
+        x, y = scenari[greatest_likelihood]
+        #without breakpoint
+        if epoch == 0:
+            # 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
+        top_plots[greatest_likelihood] = x,y,epoch
+    plots_likelihoods = list(top_plots.keys())
+    for i in range(len(plots_likelihoods)):
+        plots_likelihoods[i] = float(plots_likelihoods[i])
+    best10_plots = sorted(plots_likelihoods)[:10]
+    top_plot_lik = str(best10_plots[0])
+    plot_handles = []
+    # plt.rcParams['font.size'] = fnt_size
+    p0, = ax1.plot(top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], 'o', linestyle = "-",
+    alpha=1, lw=2, label = str(top_plots[top_plot_lik][2])+' brks | Lik='+top_plot_lik)
+    plot_handles.append(p0)
+    for k, plot_Lk in enumerate(best10_plots[1:]):
+        plot_Lk = str(plot_Lk)
+        # plt.rcParams['font.size'] = fnt_size
+        p, = ax1.plot(top_plots[plot_Lk][0], top_plots[plot_Lk][1], 'o', linestyle = "--",
+        alpha=1/(k+1), lw=1.5, label = str(top_plots[plot_Lk][2])+' brks | Lik='+plot_Lk)
+        plot_handles.append(p)
+    if theta_scale:
+        ax1.set_xlabel("Coal. time", fontsize=fnt_size)
+        ax1.set_ylabel("Pop. size scaled by N0", fontsize=fnt_size)
+        # recent_scale_lower_bound = 0.01
+        # recent_scale_upper_bound = 0.1
+        # ax1.axvline(x=recent_scale_lower_bound)
+        # 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)
+    # 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)
+    if ax is None:
+        plt.savefig(title+'_b'+str(breaks)+'.pdf')
+    # plot likelihood against nb of breakpoints
+    # best possible likelihood from SFS
+    # Segregating sites
+    S = sum(SFS_stored)
+    # Number of kept sites from which the SFS is computed
+    L = L_stored
+    # number of monomorphic sites
+    S0 = L-S
+    # print("SFS", SFS_stored)
+    # print("S", S, "L", L, "S0=", S0)
+    # 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)):
+        p_i = SFS_stored[xi] / float(S+S0)
+        Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi])
+    # basic plot likelihood
+    if ax is None:
+        fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        # plt.rcParams['font.size'] = fnt_size
+    else:
+        #plt.rcParams['font.size'] = fnt_size
+        ax2 = ax[0][0,1]
+    ax2.plot(np.array(brkpt_lik)[:, 0], np.array(brkpt_lik)[:, 1].astype(float), 'o', linestyle = "dotted", lw=2)
+    ax2.axhline(y=-Ln, linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(-Ln, 2)))
+    ax2.set_yscale('log')
+    ax2.set_xlabel("# breakpoints", fontsize=fnt_size)
+    ax2.set_ylabel("$-\log\mathcal{L}$", fontsize=fnt_size)
+    ax2.legend(loc='best', fontsize = fnt_size*0.5)
+    ax2.set_title(title+" Likelihood gain from # breakpoints")
+    if ax is None:
+        plt.savefig(title+'_Breakpts_Likelihood.pdf')
+    # AIC
+    if ax is None:
+        fig, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        # plt.rcParams['font.size'] = '18'
+    else:
+        #plt.rcParams['font.size'] = fnt_size
+        ax3 = ax[1][0,1]
+    AIC = []
+    for brk in np.array(brkpt_lik)[:, 0]:
+        brk = int(brk)
+        AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float))
+    ax3.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2)
+    # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1
+    AIC_ln = 2*(len(brkpt_lik)+1) - 2*Ln
+    ax3.axhline(y=AIC_ln, linestyle = "-.", color = "red",
+    label = "Min. AIC = "+str(round(AIC_ln, 2)))
+    selected_brks_nb = AIC.index(min(AIC))
+    ax3.set_yscale('log')
+    ax3.set_xlabel("# breakpoints", fontsize=fnt_size)
+    ax3.set_ylabel("AIC")
+    ax3.legend(loc='best', fontsize = fnt_size*0.5)
+    ax3.set_title(title+" AIC")
+    if ax is None:
+        plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
+    print("S", S)
+    # 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 = {}
+    cpt = 0
+    epochs = {}
+    plots = {}
+    # store ['best'], and [0] for epoch 0 etc...
+    for file_name in os.listdir(folder_path):
+        breaks = 0
+        cpt +=1
+        if os.path.isfile(os.path.join(folder_path, file_name)):
+            x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                                                             tgen = tgen,
+                                                             mu = mu, relative_theta_scale = theta_scale)
+            SFS_stored = sfs
+            L_stored = L
+            while not (x == 0 and y == 0):
+                if breaks not in epochs.keys():
+                    epochs[breaks] = {}
+                epochs[breaks][likelihood] = x,y
+                breaks += 1
+                x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                                                                 tgen = tgen,
+                                                                  mu = mu, relative_theta_scale = theta_scale)
+            if x == 0:
+                # last break did not work, then breaks = breaks-1
+                breaks -= 1
+    print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n")
+    print(cpt, "theta file(s) have been scanned.")
+
+    brkpt_lik = []
+    top_plots = {}
+    for epoch, scenari in epochs.items():
+        # sort starting by the smallest -log(Likelihood)
+        best10_scenari = (sorted(list(scenari.keys())))[:10]
+        greatest_likelihood = best10_scenari[0]
+        # store the tuple breakpoints and likelihood for later plot
+        brkpt_lik.append((epoch, greatest_likelihood))
+        x, y = scenari[greatest_likelihood]
+        #without breakpoint
+        if epoch == 0:
+            # 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
+        top_plots[greatest_likelihood] = x,y,epoch
+    plots_likelihoods = list(top_plots.keys())
+    for i in range(len(plots_likelihoods)):
+        plots_likelihoods[i] = float(plots_likelihoods[i])
+    best10_plots = sorted(plots_likelihoods)[:10]
+    top_plot_lik = str(best10_plots[0])
+    # store x,y,brks,likelihood
+    plots['best'] = (top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], str(top_plots[top_plot_lik][2]), top_plot_lik)
+    for k, plot_Lk in enumerate(best10_plots[1:]):
+        plot_Lk = str(plot_Lk)
+        plots[str(top_plots[plot_Lk][2])] = (top_plots[plot_Lk][0], top_plots[plot_Lk][1], str(top_plots[plot_Lk][2]), plot_Lk)
+
+    # plot likelihood against nb of breakpoints
+    # best possible likelihood from SFS
+    # Segregating sites
+    S = sum(SFS_stored)
+    # Number of kept sites from which the SFS is computed
+    L = L_stored
+    # number of monomorphic sites
+    S0 = L-S
+    # print("SFS", SFS_stored)
+    # print("S", S, "L", L, "S0=", S0)
+    # 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)):
+        p_i = SFS_stored[xi] / float(S+S0)
+        Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi])
+    # basic plot likelihood
+    Ln_Brks = [list(np.array(brkpt_lik)[:, 0]), list(np.array(brkpt_lik)[:, 1].astype(float))]
+    best_Ln = -Ln
+    AIC = []
+    for brk in np.array(brkpt_lik)[:, 0]:
+        brk = int(brk)
+        AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float))
+    AIC_Brks = [list(np.array(brkpt_lik)[:, 0]), AIC]
+    # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1
+    AIC_ln = 2*(len(brkpt_lik)+1) - 2*Ln
+    best_AIC = AIC_ln
+
+    # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC
+    # 'plots' dict keys: 'best', {epochs}('0', '1',...)
+    if input == None:
+        saved_plots = {"all_epochs":plots, "Ln_Brks":Ln_Brks,
+                        "AIC_Brks":AIC_Brks, "best_Ln":best_Ln,
+                        "best_AIC":best_AIC}
+    else:
+        # if the dict has to be loaded from input
+        with open(input, 'r') as json_file:
+            saved_plots = json.load(json_file)
+            saved_plots["all_epochs"] = plots
+            saved_plots["Ln_Brks"] = Ln_Brks
+            saved_plots["AIC_Brks"] = AIC_Brks
+            saved_plots["best_Ln"] = best_Ln
+            saved_plots["best_AIC"] = best_AIC
+    if output == None:
+        output = title+"_plotdata.json"
+    with open(output, 'w') as json_file:
+        json.dump(saved_plots, json_file)
+    return saved_plots
+
 def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, ax = None):
     #scenari = {}
     cpt = 0
         breaks = 0
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
-            x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
+            x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                              tgen = tgen,
                                                              mu = mu, relative_theta_scale = theta_scale)
             SFS_stored = sfs
                     epochs[breaks] = {}
                 epochs[breaks][likelihood] = x,y
                 breaks += 1
-                x,y,likelihood,theta,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                                  tgen = tgen,
                                                                   mu = mu, relative_theta_scale = theta_scale)
             if x == 0:
     return ax[0], ax[1]
 
 def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
-    breaks_max = 10, output = None):
+    breaks_max = 10, input = None, output = None):
     """
     Save theta values as is to do basic plots.
     """
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
             for k in range(breaks_max):
-                x,y,likelihood,thetas,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = k,
+                x,y,likelihood,thetas,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = k,
                                                                  tgen = tgen,
                                                                  mu = mu, relative_theta_scale = theta_scale)
                 if thetas == 0:
         # x2_plot, y2_plot = plot_straight_x_y(x_2, y)
         p2 = x_2, y
         lines_fig2.append(p2)
-
-    saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
-                    "prop":prop}
+    if input == None:
+        saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
+                        "prop":prop}
+    else:
+        # if the dict has to be loaded from input
+        with open(input, 'r') as json_file:
+            saved_plots = json.load(json_file)
+        saved_plots["raw_stairs"] = plots
+        saved_plots["scaled_stairs"] = lines_fig2
+        saved_plots["prop"] = prop
     if output == None:
         output = title+"_plotdata.json"
     with open(output, 'w') as json_file:
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
             for k in range(breaks_max):
-                x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = k,
+                x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = k,
                                                                  tgen = tgen,
                                                                  mu = mu, relative_theta_scale = theta_scale)
                 if thetas == 0:
     ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
                             prop = loaded_data['prop'], title = title, ax = ax1)
     ax1, ax2 = plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = [ax1, ax2])
+    save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, input = title+"_plotdata.json", output = title+"_plotdata.json")
     fig1.savefig(title+'_combined_p1.pdf')
     fig2.savefig(title+'_combined_p2.pdf')
     plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],