#17 and #13: legends per default in kplot

pymule/plot.py

 import matplotlib.pyplot as plt
 import matplotlib.collections
 import matplotlib.patches
+import matplotlib.lines
 import numpy as np
 from matplotlib import rc
 from colours import *
@@ -160,8 +161,12 @@ def twopanel(labelx,
 
 def kplot(sigma, labelx='$x_e$', labelsigma=None,
           labelknlo="$K^{(1)}$", labelknnlo="$K^{(2)}$",
+          legend={'lo': r'$\rm LO$', 'nlo': r'$\rm NLO$', 'nnlo': r'$\rm NNLO$'},
+          legendopts={'what': 'l', 'loc': 'upper right'},
           show=[0,-1], showk=[1,2], nomule=False):
 
+    legend = legend.copy()
+    legendopts = legendopts.copy()
     if labelsigma is None:
         labelsigma = r"$\D\sigma$"
 
@@ -202,6 +207,14 @@ def kplot(sigma, labelx='$x_e$', labelsigma=None,
             )
         }
     else:
+        try:
+            show.remove(2)
+        except ValueError:
+            pass
+        try:
+            showk.remove(2)
+        except ValueError:
+            pass
         orders = ['lo', 'nlo']
         xsec = {
             'lo': sigma['lo'],
@@ -211,7 +224,7 @@ def kplot(sigma, labelx='$x_e$', labelsigma=None,
     if 1 in showk:
         kwargs['downleft'] = divideplots(xsec['nlo'], xsec['lo'])
         kwargs['coldownleft'] = [orderscheme['nlo']]
-    if 2 in showk and 'nnlo' in sigma:
+    if 2 in showk:
         kwargs['downright'] = divideplots(xsec['nnlo'], xsec['nlo'])
         kwargs['coldownright'] = [orderscheme['nnlo']]
 
@@ -230,6 +243,19 @@ def kplot(sigma, labelx='$x_e$', labelsigma=None,
 
     axs[1].axhline(1, color='black', linewidth=1, zorder=1)
 
+    plegend = zip(*[
+        (
+            matplotlib.lines.Line2D([0], [0], color=orderscheme[i]),
+            legend[i]
+        )
+        for i in sorted(list(set(orders[i] for i in show+showk)))
+        if i in legend
+    ])
+    what = {
+        'fig': fig, 'u': axs[0], 'l': axs[1]
+    }[legendopts.pop('what', 'u')]
+    what.legend(plegend[0], plegend[1], **legendopts)
+
     if not nomule:
         mulify(fig)
     return fig, axs