Update SFS plotting function

customgraphics.py

@@ -200,14 +200,14 @@ 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):
+def barplot(x=None, y=None, ylab=None, xlab=None, title=None, label=None, xticks = None, width=1):
     if x:
         x = list(x)
         plt.xticks(x)
-        plt.bar(x, y)
+        plt.bar(x, y, width=width, label=label)
     else:
         x = list(range(len(y)))
-        plt.bar(x, y)
+        plt.bar(x, y, width=width, label=label)
         plt.xticks(x)
     if ylab:
         plt.ylabel(ylab)
@@ -215,6 +215,9 @@ def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
         plt.xlabel(xlab)
     if title:
         plt.title(title)
+    if xticks:
+        plt.xticks(xticks)
+    plt.legend()
     plt.show()
 
 def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None,

sfs_tools.py

@@ -21,6 +21,7 @@ import gzip
 import sys
 import matplotlib.pyplot as plt
 from frst import customgraphics
+import numpy as np
 
 def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False,
                  strip = False, count_ext = False):
@@ -192,7 +193,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):
+def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False, ploidy = 2):
     sfs_val = []
     n = len(sfs.values())
     sum_sites = sum(list(sfs.values()))
@@ -222,7 +223,7 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
             
     #terminal case, same for folded or unfolded
     if transformed:
-        last_bin = list(sfs.values())[n-1] * n/2
+        last_bin = list(sfs.values())[n-1] * n/ploidy
     else:
         last_bin = list(sfs.values())[n-1]
     sfs_val[-1] = last_bin
@@ -235,22 +236,33 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
         
         #print(sum(sfs_val))
     #build the plot
-    title = title+" (n="+str(len(sfs_val))+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
-    print("SFS =", sfs)
     if folded:
         xlab = "Minor allele frequency"
+        n_title = n
+    else:
+        # the spectrum is n-1 long when unfolded
+        n_title = n+1
+    
+    title = title+" (n="+str(n_title)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
+    print("SFS =", sfs)
+
+    X_axis = list(sfs.keys()) 
+    
+
     if transformed:
-        print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val)
+        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")
+
     else:
         if normalized:
             # then plot a theoritical distribution as 1/i
-            expected_y = [1/(2*x+1) for x in list(sfs.keys())]
+            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")
             print(sum(expected_y))
-            #plt.plot([x for x in list(sfs.keys())], expected_y, color='r', linestyle='-')
+    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()) )
-            #print(expected_y)
-            
-    customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title)
     plt.show()
 
 if __name__ == "__main__":

swp2.py

@@ -3,13 +3,11 @@ import os
 import numpy as np
 import math
 import json
-import io
+
-from scipy.special import gammaln
-from matplotlib.backends.backend_pdf import PdfPages
-from matplotlib.ticker import MaxNLocator
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes
-from matplotlib.ticker import MultipleLocator
 def log_facto(k):
+    """
+    Using the Stirling's approximation
+    """
     k = int(k)
     if k > 1e6:
         return k * np.log(k) - k + np.log(2*math.pi*k)/2
@@ -18,15 +16,7 @@ def log_facto(k):
         val += np.log(i)
     return val
 
-def log_facto_1(k):
+def parse_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False):
-    startf = 1 # start of factorial sequence
-    stopf  = int(k+1) # end of of factorial sequence
-
-    q = gammaln(range(startf+1, stopf+1)) # n! = G(n+1)
-
-    return q[-1]
-
-def return_x_y_from_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()
@@ -119,108 +109,9 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_
         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
 
-def return_x_y_from_stwp_theta_file_as_is(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()
-        L = float(line.split()[2])
-        rands = swp_file.readline()
-        line = swp_file.readline()
-        # skip empty lines before SFS
-        while line == "\n":
-            line = swp_file.readline()
-        sfs = np.array(line.split()).astype(float)
-        # Process lines until the end of the file
-        while line:
-            # check at each line
-            if line.startswith("dim") :
-                dim = int(line.split()[1])
-                if dim == breaks+1:
-                    likelihood = line.split()[5]
-                    groups = line.split()[6:6+dim]
-                    theta_site = line.split()[6+dim:6+dim+1+dim]
-                elif dim < breaks+1:
-                    line = swp_file.readline()
-                    continue
-                elif dim > breaks+1:
-                    break
-                    #return 0,0,0
-            # Read the next line
-            line = swp_file.readline()
-    #### END of parsing
-    # quit this file if the number of dimensions is incorrect
-    if dim < breaks+1:
-        return 0,0
-    # get n, the last bin of the last group
-    # revert the list of groups as the most recent times correspond
-    # to the closest and last leafs of the coal. tree.
-    groups = groups[::-1]
-    theta_site = theta_site[::-1]
-
-    thetas = {}
-
-    for i in range(len(groups)):
-        groups[i] = groups[i].split(',')
-        # print(groups[i], len(groups[i]))
-        thetas[i] = [float(theta_site[i]), groups[i], likelihood]
-    return thetas, sfs
-
-def plot_k_epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title", theta_scale = True):
-    scenari = {}
-    cpt = 0
-    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,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
-                                                             tgen = tgen,
-                                     mu = mu, relative_theta_scale = theta_scale)
-            if x == 0 or y == 0:
-                continue
-            cpt +=1
-            scenari[likelihood] = x,y
-    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.")
-    # sort starting by the smallest -log(Likelihood)
-    print(scenari)
-    best10_scenari = (sorted(list(scenari.keys())))[:10]
-    print("10 greatest Likelihoods", best10_scenari)
-    greatest_likelihood = best10_scenari[0]
-    x, y = scenari[greatest_likelihood]
-    my_dpi = 300
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    plt.plot(x, y, 'r-', lw=2, label = 'Lik='+greatest_likelihood)
-    #plt.yscale('log')
-    plt.xscale('log')
-    plt.grid(True,which="both", linestyle='--', alpha = 0.3)
-
-    for scenario in best10_scenari[1:]:
-        x,y = scenari[scenario]
-        #print("\n----  Lik:",scenario,"\n\nt=", x,"\n\nN=",y, "\n\n")
-        plt.plot(x, y, '--', lw=1, label = 'Lik='+scenario)
-    if theta_scale:
-        plt.xlabel("Coal. time")
-        plt.ylabel("Pop. size scaled by N0")
-        recent_scale_lower_bound = y[0] * 0.01
-        recent_scale_upper_bound = y[0] * 0.1
-        plt.axvline(x=recent_scale_lower_bound)
-        plt.axvline(x=recent_scale_upper_bound)
-    else:
-        # years
-        plt.xlabel("Time (years)")
-        plt.ylabel("Individuals (N)")
-    plt.legend(loc='upper right')
-    plt.title(title)
-    plt.savefig(title+'_b'+str(breaks)+'.pdf')
-
 def plot_straight_x_y(x,y):
     x_1 = [x[0]]
     y_1 = []
@@ -233,7 +124,8 @@ def plot_straight_x_y(x,y):
     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):
+def plot_all_epochs_thetafolder_old(folder_path, mu, tgen, title = "Title",
+    theta_scale = True, ax = None, input = None, output = None):
     #scenari = {}
     cpt = 0
     epochs = {}
@@ -241,7 +133,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
         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
@@ -251,7 +143,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
                     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:
@@ -384,8 +276,195 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     # return plots
     return ax[0], ax[1]
 
+def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
+    theta_scale = True, ax = None, input = None, output = None):
+    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)
+    plot_handles = []
+    best_plot = full_dict['all_epochs']['best']
+    p0, = ax1.plot(best_plot[0], best_plot[1], 'o', linestyle = "-",
+    alpha=1, lw=2, label = str(best_plot[2])+' brks | Lik='+best_plot[3])
+    plot_handles.append(p0)
+    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 = "--",
+        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:
+        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
+    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(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], 'o', linestyle = "dotted", lw=2)
+    ax2.axhline(y=full_dict['best_Ln'], linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(full_dict['best_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 = full_dict['AIC_Brks']
+    ax3.plot(AIC[0], AIC[1], 'o', linestyle = "dotted", lw=2)
+    ax3.axhline(y=full_dict['best_AIC'], linestyle = "-.", color = "red",
+    label = "Min. AIC = "+str(round(full_dict['best_AIC'], 2)))
+    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')
+    # 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)
+    plots['plots'] = []
+    for k, plot_Lk in enumerate(best10_plots[1:]):
+        plot_Lk = str(plot_Lk)
+        plots['plots'].append([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 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.
     """
@@ -396,7 +475,7 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
             for k in range(breaks_max):
-                thetas,sfs = return_x_y_from_stwp_theta_file_as_is(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:
@@ -478,9 +557,16 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         # x2_plot, y2_plot = plot_straight_x_y(x_2, y)
         p2 = x_2, y
         lines_fig2.append(p2)
-
+    if input == None:
-    saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
+        saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
-                    "prop":prop}
+                        "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:
@@ -580,159 +666,6 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
     # return plots
     return ax
 
-def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 10, ax = None, n_ticks = 10):
-    """
-    Use theta values as is to do basic plots.
-    """
-    cpt = 0
-    epochs = {}
-    len_sfs = 0
-    for file_name in os.listdir(folder_path):
-        cpt +=1
-        if os.path.isfile(os.path.join(folder_path, file_name)):
-            for k in range(breaks_max):
-                thetas,sfs = return_x_y_from_stwp_theta_file_as_is(folder_path+file_name, breaks = k,
-                                                                 tgen = tgen,
-                                                                 mu = mu, relative_theta_scale = theta_scale)
-                if thetas == 0:
-                    continue
-                if len(thetas)-1 != k:
-                    continue
-                if k not in epochs.keys():
-                    epochs[k] = {}
-                likelihood = str(eval(thetas[k][2]))
-                epochs[k][likelihood] = thetas
-                #epochs[k] = thetas
-    print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(k)+"\n*******\n")
-    print(cpt, "theta file(s) have been scanned.")
-    # multiple fig
-    if ax is None:
-        # intialize figure 1
-        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[0, 1]
-        plt.subplots_adjust(wspace=0.3, hspace=0.3)
-    plots = []
-    best_epochs = {}
-    for epoch in epochs:
-        likelihoods = []
-        for key in epochs[epoch].keys():
-            likelihoods.append(key)
-        likelihoods.sort()
-        minLogLn = str(likelihoods[0])
-        best_epochs[epoch] = epochs[epoch][minLogLn]
-    for epoch, theta in best_epochs.items():
-        groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
-        x = []
-        y = []
-        thetas = np.array(list(theta.values()), dtype=object)[:, 0]
-        for i,group in enumerate(groups):
-            x += group[::-1]
-            y += list(np.repeat(thetas[i], len(group)))
-            if epoch == 0:
-                N0 = y[0]
-                # compute the proportion of information used at each bin of the SFS
-                sum_theta_i = 0
-                for i in range(2, len(y)+2):
-                    sum_theta_i+=y[i-2] / (i-1)
-                prop = []
-                for k in range(2, len(y)+2):
-                    prop.append(y[k-2] / (k - 1) / sum_theta_i)
-                prop = prop[::-1]
-                # print(prop, "\n", sum(prop))
-        # normalise to N0 (N0 of epoch1)
-        x_ticks = ax1.get_xticks()
-        for i in range(len(y)):
-            y[i] = y[i]/N0
-        # plot
-        x_plot, y_plot = plot_straight_x_y(x, y)
-        #plt.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        p, = ax1.plot(x_plot, y_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        # add plot to the list of all plots to superimpose
-        plots.append(p)
-    #print(prop, "\n", sum(prop))
-    #ax.legend(handles=[p0]+plots)
-    ax1.set_xlabel("# bin", fontsize=fnt_size)
-    # Set the x-axis locator to reduce the number of ticks to 10
-    ax1.set_ylabel("theta", fontsize=fnt_size)
-    ax1.set_title(title, fontsize=fnt_size)
-    ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
-    ax1.set_xticks(x_ticks)
-    if len(prop) >= 18:
-        ax1.locator_params(nbins=n_ticks)
-    # new scale of ticks if too many values
-    cumul = 0
-    prop_cumul = []
-    for val in prop:
-        prop_cumul.append(val+cumul)
-        cumul = val+cumul
-    ax1.set_xticklabels([f'{x[k]}\n{val:.2f}' for k, val in enumerate(prop_cumul)])
-    if ax is None:
-        plt.savefig(title+'_raw'+str(k)+'.pdf')
-    # fig 2 & 3
-    if ax is None:
-        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)
-    else:
-        # plt.rcParams['font.size'] = fnt_size
-        # place of plots on the grid
-        ax2 = ax[1,0]
-        ax3 = ax[1,1]
-    lines_fig2 = []
-    lines_fig3 = []
-    #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    for epoch, theta in best_epochs.items():
-        groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
-        x = []
-        y = []
-        thetas = np.array(list(theta.values()), dtype=object)[:, 0]
-        for i,group in enumerate(groups):
-            x += group[::-1]
-            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
-        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))
-            x_2.append(T)
-        # Plotting (fig 2)
-        x_2 = [0]+x_2
-        y = [y[0]]+y
-        x2_plot, y2_plot = plot_straight_x_y(x_2, y)
-        p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(epoch)+' 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(epoch)+' 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)
-    if ax is None:
-        plt.savefig(title+'_plot2_'+str(k)+'.pdf')
-    ax3.set_xscale('log')
-    ax3.set_yscale('log')
-    ax3.set_xlabel("log Relative 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)
-    if ax is None:
-        plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
-        plt.clf()
-    # return plots
-    return ax
-
 def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True):
     my_dpi = 300
     # # Add some extra space for the second axis at the bottom
@@ -751,7 +684,9 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale =
     # # plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = None)
     # # plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = None)
     # # plt.clf()
-    # save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, output = title+"_plotdata.json")
+    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")
+
     with open(title+"_plotdata.json", 'r') as json_file:
         loaded_data = json.load(json_file)
     # plot page 1 of summary
@@ -767,7 +702,7 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale =
 
     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])
+    ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2])
     fig1.savefig(title+'_combined_p1.pdf')
     fig2.savefig(title+'_combined_p2.pdf')
     plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],