diff --git a/pymule/plot.py b/pymule/plot.py index c7aae308162bde2d98f91ae9a15f6c2f75967035..d07918e98fdfb0dc47a6ee65b12cf4bdb1dcceff 100644 --- a/pymule/plot.py +++ b/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)