22: Aligned K factors

pymule/mpl_axes_aligner.py

+# Stolen from
+# https://github.com/ryutok/mpl_axes_aligner
+
+
+def _calc_rorg(org, ival, fval):
+    rorg = (org - ival) / (fval - ival)
+    if rorg < 0:
+        rorg = 0
+        ival = org
+    elif rorg > 1:
+        rorg = 1
+        fval = org
+    return rorg, ival, fval
+
+
+def _calc_pos(org1, org2, lim1, lim2):
+    ival1, fval1 = lim1
+    ival2, fval2 = lim2
+
+    rorg1, ival1, fval1 = _calc_rorg(org1, ival1, fval1)
+    rorg2, ival2, fval2 = _calc_rorg(org2, ival2, fval2)
+    pos = (rorg1 + rorg2) / 2
+
+    if pos == 0 or pos == 1:
+        raise ValueError("When pos=None, at least one origin should be "
+                         "within the initial plotting range.")
+
+    return pos
+
+
+def _expand_range(org, pos, ival, fval):
+    if pos <= 0 or pos >= 1:
+        raise ValueError("When expand=True, the position to align the origin "
+                         "should be 0 < pos < 1.")
+
+    rorg = (org - ival) / (fval - ival)
+    if rorg > pos:
+        fval = (org - ival + pos*ival) / pos
+    else:
+        ival = (org - pos*fval) / (1 - pos)
+    return ival, fval
+
+
+def shiftyaxis(ax, org, pos, expand=False):
+    bottom, top = ax.get_ylim()
+    bottom, top = _expand_range(org, pos, bottom, top)
+    ax.set_ylim(bottom, top)
+
+
+def yaxes(ax1, org1, ax2, org2):
+    # Get plotting ranges
+    try:
+        lim1 = list(ax1.get_ylim())
+        lim2 = list(ax2.get_ylim())
+    except AttributeError or TypeError:
+        raise TypeError("'ax1' and 'ax2' should be Axes objects of "
+                        "matplotlib.")
+
+    # Calculate the position
+    pos = _calc_pos(org1, org2, lim1, lim2)
+
+    # Apply the new ranges
+    shiftyaxis(ax1, org1, pos, True)
+    shiftyaxis(ax2, org2, pos, True)

pymule/plot.py

@@ -5,6 +5,7 @@ import numpy as np
 from matplotlib import rc
 from colours import *
 from errortools import *
+import mpl_axes_aligner
 
 
 rc('text', usetex=True)
@@ -75,7 +76,8 @@ def twopanel(labelx,
              upleft=[], labupleft="", colupleft=defcol,
              downleft=[], labdownleft="", coldownleft=defcol,
              upright=[], labupright="", colupright=defcol,
-             downright=[], labdownright="", coldownright=defcol):
+             downright=[], labdownright="", coldownright=defcol,
+             upalign=[], downalign=[]):
 
     if type(upleft) == np.ndarray:
         upleft = [upleft]
@@ -101,6 +103,11 @@ def twopanel(labelx,
             ax2.set_ylabel(labupright)
         for i, c in zip(upright, colupright):
             errorband(i, ax=ax2, col=c)
+        if len(upalign) == 2:
+            mpl_axes_aligner.yaxes(
+                axs[0], upalign[0],
+                ax2, upalign[1]
+            )
 
     axs[1].set_ylabel(labdownleft)
     for i, c in zip(downleft, coldownleft):
@@ -112,6 +119,11 @@ def twopanel(labelx,
             ax3.set_ylabel(labdownright)
         for i, c in zip(downright, coldownright):
             errorband(i, ax=ax3, col=c)
+        if len(downalign) == 2:
+            mpl_axes_aligner.yaxes(
+                axs[1], downalign[0],
+                ax3, downalign[1]
+            )
 
     return fig, axs
 
@@ -125,7 +137,8 @@ def kplot(sigma, labelx='x_e', labelsigma=None,
 
     kwargs = {
         "labelx":"$%s$" % labelx,
-        "labupleft": labelsigma
+        "labupleft": labelsigma,
+        "downalign": [1,1]
     }
 
     if type(sigma) != dict:
@@ -168,4 +181,8 @@ def kplot(sigma, labelx='x_e', labelsigma=None,
         orderscheme[orders[i]] for i in show
     ]
 
-    return twopanel(**kwargs)
+    fig, axs = twopanel(**kwargs)
+
+    axs[1].axhline(1, color='black', linewidth=1, zorder=1)
+
+    return fig, axs