18b: re-introduce smarter kplot

pymule/plot.py

@@ -3,6 +3,7 @@ import matplotlib.collections
 import matplotlib.patches
 import numpy as np
 from matplotlib import rc
+from errortools import *
 
 
 rc('text', usetex=True)
@@ -95,3 +96,50 @@ def twopanel(labelx,
             errorband(i, ax=ax3)
 
     return fig, axs
+
+
+def kplot(sigma, labelx='x_e', labelsigma=None,
+          labelknlo="$K^{(1)}$", labelknnlo="$K^{(2)}$",
+          show=[0,-1], showk=[1,2]):
+
+    if labelsigma is None:
+        labelsigma = r"$\D\sigma\,/\,\D " + labelx + "$"
+
+    kwargs = {
+        "labelx":"$%s$" % labelx,
+        "labupleft": labelsigma
+    }
+
+    if type(sigma) != dict:
+        raise ValueError("kplot takes dicts!")
+    if "lo" not in sigma:
+        raise KeyError("kplot: no leading order")
+    if "nlo" not in sigma:
+        raise KeyError("kplot: no next-to-leading order")
+    if "nnlo" in sigma:
+        orders = ['lo', 'nlo', 'nnlo']
+        xsec = {
+            'lo': sigma['lo'],
+            'nlo': addplots(sigma['lo'], sigma['nlo']),
+            'nnlo': addplots(
+                sigma['lo'],
+                addplots(sigma['nlo'], sigma['nnlo'])
+            )
+        }
+    else:
+        orders = ['lo', 'nlo']
+        xsec = {
+            'lo': sigma['lo'],
+            'nlo': addplots(sigma['lo'], sigma['nlo'])
+        }
+
+    if 1 in showk:
+        kwargs['downleft'] = divideplots(sigma['nlo'], sigma['lo'])
+    if 2 in showk and 'nnlo' in sigma:
+        kwargs['downright'] = divideplots(sigma['nnlo'], sigma['lo'])
+
+    kwargs['upleft'] = [
+        xsec[orders[i]] for i in show
+    ]
+
+    return twopanel(**kwargs)