16: Added kfactor plotter in re #17 and #13

pymule/__init__.py

@@ -16,6 +16,6 @@ from loader import importreg, pattern, setup, sigma,                          \
                    mergeset, mergeseeds, mergefks,                            \
                    addsets, scaleset
 
-from plot import errorband
+from plot import errorband, kplot
 
 from constants import pi, alpha, GF, conv, Mmu, Mel, Mtau

pymule/plot.py

@@ -2,6 +2,14 @@ import matplotlib.pyplot as plt
 import matplotlib.collections
 import matplotlib.patches
 import numpy as np
+from matplotlib import rc
+
+
+rc('text', usetex=True)
+rc('text.latex', preamble="\n".join([
+    '\\usepackage{amsmath}',
+    '\\newcommand{\\D}{\mathrm{d}}'
+]))
 
 
 def errorband(p, ax=None, underflow=False, overflow=False):
@@ -42,3 +50,39 @@ def errorband(p, ax=None, underflow=False, overflow=False):
 
     ax.add_collection(pc)
     return artist
+
+
+def kplot(sigma, k, sigmar=[], kr=[],
+          labelx='x_e',
+          labelk="$K$", labelkr=None,
+          labelsigma=None, labelsigmar=None):
+
+    if labelsigma is None:
+        labelsigma = r"$\D\sigma\,/\,\D " + labelx + "$"
+
+    fig, axs = plt.subplots(2, sharex=True, sharey=True, gridspec_kw={'hspace': 0})
+    axs[1].set_xlabel("$"+labelx+"$")
+
+    axs[0].set_ylabel(labelsigma)
+    for i in sigma:
+        errorband(i, ax=axs[0])
+
+    if len(sigmar) > 0:
+        ax2 = axs[0].twinx()
+        if labelsigmar is not None:
+            ax2.set_ylabel(labelsigmar)
+        for i in sigmar:
+            errorband(i, ax=ax2)
+
+    axs[1].set_ylabel(labelk)
+    for i in k:
+        errorband(i, ax=axs[1])
+
+    if len(kr) > 0:
+        ax3 = axs[1].twinx()
+        if labelkr is not None:
+            ax3.set_ylabel(labelkr)
+        for i in kr:
+            errorband(i, ax=ax3)
+
+    return fig, axs