Adjust x-axis ticks for large SFSs

swp2.py

@@ -4,6 +4,7 @@ import numpy as np
 import math
 from scipy.special import gammaln
 from matplotlib.backends.backend_pdf import PdfPages
+from matplotlib.ticker import MaxNLocator
 
 def log_facto(k):
     k = int(k)
@@ -320,7 +321,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     plt.title(title)
     plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
 
-def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 5):
+def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 12):
     """
     Use theta values as is to do basic plots.
     """
@@ -338,10 +339,17 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
                 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.")
-
     # intialize figure 1
     my_dpi = 300
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    # plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    # multiple fig
+    fig, ax = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    # Add some extra space for the second axis at the bottom
+    fig.subplots_adjust(bottom=0.15)
+    twin = ax.twiny()
+    # Offset the right spine of twin2
+    # twin.spines.right.set_position(("axes", 1.2))
+    plots = []
     for epoch, theta in epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
         x = []
@@ -352,21 +360,38 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
             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)
         for i in range(len(y)):
             y[i] = y[i]/N0
-        # compute the proportion of information used at each bin of the SFS
-        sum_theta_i = 0
-        for i in range(2, len(y)-1):
-            sum_theta_i+=y[i] / (i-1)
-        prop = []
-        for k in range(2, len(y)-1):
-            prop.append(y[k] / (k - 1) / sum_theta_i)
-        # plot 
-        plt.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        plt.xlabel("# breaks")
-        plt.ylabel("theta")
-        plt.legend(loc='upper right')
-        plt.savefig(title+'_test'+str(k)+'.pdf')
+        # plot
+        #plt.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        p, = ax.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        # add plot to the list of all plots to superimpose
+        plots.append(p)
+    # virtual line to get the second x axis for proportions
+    p0, = twin.plot(prop, y, alpha = 0, label="Proportion")
+    # Move twinned axis ticks and label from top to bottom
+    twin.xaxis.set_ticks_position("bottom")
+    twin.xaxis.set_label_position("bottom")
+    # Offset the twin axis below the host
+    twin.spines["bottom"].set_position(("axes", -0.1))
+    #ax.legend(handles=[p0]+plots)
+    ax.set_xlabel("# breaks")
+    # Set the x-axis locator to reduce the number of ticks = 10
+    ax.xaxis.set_major_locator(MaxNLocator(nbins=10))
+    ax.set_ylabel("theta")
+    # twin.set_ylabel("Proportion")
+    plt.legend(handles=plots, loc='upper right')
+    plt.savefig(title+'_raw'+str(k)+'.pdf')
     # fig 2 & 3
     plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     for epoch, theta in epochs.items():
@@ -394,14 +419,14 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         plt.xlabel("# breaks")
         plt.ylabel("theta")
         plt.legend(loc='upper right')
-        plt.savefig(title+'_test'+str(k)+'.pdf')
+        plt.savefig(title+'_plot2_'+str(k)+'.pdf')
         # Plotting (fig 3) which is the same but log scale for x
         plt.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
         plt.xscale('log')
         plt.xlabel("# breaks")
         plt.ylabel("theta")
         plt.legend(loc='upper right')
-        plt.savefig(title+'_test'+str(k)+'_log.pdf')
+        plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
 
 def save_multi_image(filename):
     pp = PdfPages(filename)