From 0ffb02e5110f23f1a2e5b4fb061a866af873e5df Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Sun, 25 Feb 2024 03:32:01 +0100
Subject: [PATCH] Some tweaking of axis limits

---
 swp2.py | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/swp2.py b/swp2.py
index edabe79..207021c 100644
--- a/swp2.py
+++ b/swp2.py
@@ -483,8 +483,12 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
         if subset is not None:
             if breaks in subset:
                 masking_alpha = 0.75
+                autoscale = True
             else:
                 masking_alpha = 0
+                autoscale = False
+        ax2.set_autoscale_on(autoscale)
+        ax3.set_autoscale_on(autoscale)
         p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
         # Plotting (fig 3) which is the same but log scale for x
         p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
@@ -497,16 +501,24 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
         xlabel = "Theta scaled by N0"
         ylabel = "Theta scaled by N0"
     else:
-        xlabel = "t"
-        ylabel = r"$\theta$"
+        xlabel = "time"
+        ylabel = "Effective pop. size (Ne)"
     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)
+        xlim_val = plt.gca().get_xlim()
         x_ticks = list(plt.xticks())[0]
         plt.gca().set_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)
+        # rescale y to effective pop size
+        ylim_val = plt.gca().get_ylim()
+        y_ticks = list(plt.yticks())[0]
+        plt.gca().set_yticks(y_ticks)
+        plt.gca().set_ylim(ylim_val)
+        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)
@@ -522,14 +534,21 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax =
         ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
     ax3.set_xscale('log')
     ax3.set_yscale('log')
-    ax3.set_xlabel("time log scale", fontsize=fnt_size)
-    ax3.set_ylabel("theta", fontsize=fnt_size)
+    ax3.set_xlabel(xlabel, fontsize=fnt_size)
+    ax3.set_ylabel(ylabel, fontsize=fnt_size)
     ax3.set_title(title, fontsize=fnt_size)
     ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
+    xlim_val = ax3.get_xlim()
     x_ticks = list(ax3.get_xticks())
-    ax3.set_xlim(left=min(x_ticks))
     ax3.set_xticks(x_ticks)
+    ax3.set_xlim(xlim_val)
     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)
+    ylim_val = ax3.get_ylim()
+    # rescale y to effective pop size
+    y_ticks = list(ax3.get_yticks())
+    ax3.set_yticks(y_ticks)
+    ax3.set_ylim(ylim_val)
+    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.subplots_adjust(bottom=0.2)  # Adjust the value as needed
     if ax is None: