diff options
-rwxr-xr-x | utils/plot-root-results | 308 |
1 files changed, 308 insertions, 0 deletions
diff --git a/utils/plot-root-results b/utils/plot-root-results new file mode 100755 index 0000000..e8dcf19 --- /dev/null +++ b/utils/plot-root-results @@ -0,0 +1,308 @@ +#!/usr/bin/env python +# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago> +# +# This program is free software: you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the Free +# Software Foundation, either version 3 of the License, or (at your option) +# any later version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for +# more details. +# +# You should have received a copy of the GNU General Public License along with +# this program. If not, see <https://www.gnu.org/licenses/>. + +from __future__ import print_function, division +import numpy as np + +# Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py +def despine(fig=None, ax=None, top=True, right=True, left=False, + bottom=False, offset=None, trim=False): + """Remove the top and right spines from plot(s). + + fig : matplotlib figure, optional + Figure to despine all axes of, default uses current figure. + ax : matplotlib axes, optional + Specific axes object to despine. + top, right, left, bottom : boolean, optional + If True, remove that spine. + offset : int or dict, optional + Absolute distance, in points, spines should be moved away + from the axes (negative values move spines inward). A single value + applies to all spines; a dict can be used to set offset values per + side. + trim : bool, optional + If True, limit spines to the smallest and largest major tick + on each non-despined axis. + + Returns + ------- + None + + """ + # Get references to the axes we want + if fig is None and ax is None: + axes = plt.gcf().axes + elif fig is not None: + axes = fig.axes + elif ax is not None: + axes = [ax] + + for ax_i in axes: + for side in ["top", "right", "left", "bottom"]: + # Toggle the spine objects + is_visible = not locals()[side] + ax_i.spines[side].set_visible(is_visible) + if offset is not None and is_visible: + try: + val = offset.get(side, 0) + except AttributeError: + val = offset + _set_spine_position(ax_i.spines[side], ('outward', val)) + + # Potentially move the ticks + if left and not right: + maj_on = any( + t.tick1line.get_visible() + for t in ax_i.yaxis.majorTicks + ) + min_on = any( + t.tick1line.get_visible() + for t in ax_i.yaxis.minorTicks + ) + ax_i.yaxis.set_ticks_position("right") + for t in ax_i.yaxis.majorTicks: + t.tick2line.set_visible(maj_on) + for t in ax_i.yaxis.minorTicks: + t.tick2line.set_visible(min_on) + + if bottom and not top: + maj_on = any( + t.tick1line.get_visible() + for t in ax_i.xaxis.majorTicks + ) + min_on = any( + t.tick1line.get_visible() + for t in ax_i.xaxis.minorTicks + ) + ax_i.xaxis.set_ticks_position("top") + for t in ax_i.xaxis.majorTicks: + t.tick2line.set_visible(maj_on) + for t in ax_i.xaxis.minorTicks: + t.tick2line.set_visible(min_on) + + if trim: + # clip off the parts of the spines that extend past major ticks + xticks = ax_i.get_xticks() + if xticks.size: + firsttick = np.compress(xticks >= min(ax_i.get_xlim()), + xticks)[0] + lasttick = np.compress(xticks <= max(ax_i.get_xlim()), + xticks)[-1] + ax_i.spines['bottom'].set_bounds(firsttick, lasttick) + ax_i.spines['top'].set_bounds(firsttick, lasttick) + newticks = xticks.compress(xticks <= lasttick) + newticks = newticks.compress(newticks >= firsttick) + ax_i.set_xticks(newticks) + + yticks = ax_i.get_yticks() + if yticks.size: + firsttick = np.compress(yticks >= min(ax_i.get_ylim()), + yticks)[0] + lasttick = np.compress(yticks <= max(ax_i.get_ylim()), + yticks)[-1] + ax_i.spines['left'].set_bounds(firsttick, lasttick) + ax_i.spines['right'].set_bounds(firsttick, lasttick) + newticks = yticks.compress(yticks <= lasttick) + newticks = newticks.compress(newticks >= firsttick) + ax_i.set_yticks(newticks) + +if __name__ == '__main__': + import ROOT + import argparse + from os.path import split + from matplotlib.ticker import FuncFormatter + + parser = argparse.ArgumentParser("plot ROOT fit results") + parser.add_argument("filename", help="input file") + parser.add_argument("--save", action="store_true", default=False, help="save plots") + args = parser.parse_args() + + if args.save: + # default \textwidth for a fullpage article in Latex is 16.50764 cm. + # You can figure this out by compiling the following TeX document: + # + # \documentclass{article} + # \usepackage{fullpage} + # \usepackage{layouts} + # \begin{document} + # textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth} + # \end{document} + + width = 16.50764 + width /= 2.54 # cm -> inches + # According to this page: + # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf, + # Tufte suggests an aspect ratio of 1.5 - 1.6. + height = width/1.5 + FIGSIZE = (width,height) + + import matplotlib.pyplot as plt + + font = {'family':'serif', 'serif': ['computer modern roman']} + plt.rc('font',**font) + + plt.rc('text', usetex=True) + else: + # on retina screens, the default plots are way too small + # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1 + # Qt5 will scale everything using the dpi in ~/.Xresources + import matplotlib + matplotlib.use("Qt5Agg") + + import matplotlib.pyplot as plt + + # Default figure size. Currently set to my monitor width and height so that + # things are properly formatted + FIGSIZE = (13.78,7.48) + + # Make the defalt font bigger + plt.rc('font', size=22) + + root_file = ROOT.TFile(args.filename) + + head, tail = split(args.filename) + + if tail.startswith("e_") or tail.startswith("electron"): + prefix = "electron" + elif tail.startswith("mu_") or tail.startswith("muon"): + prefix = "muon" + else: + prefix = "" + + try: + if root_file.Get("h1"): + for hist_number, tf1_number in zip([1,2,4,5],[1,2,3,None]): + h = root_file.Get("h%i" % hist_number) + if tf1_number: + f = root_file.Get("f%i" % tf1_number) + + bins = [h.GetXaxis().GetBinLowEdge(i) for i in range(1,h.GetNbinsX()+1)] + [h.GetXaxis().GetBinUpEdge(h.GetNbinsX())] + hist = [h.GetBinContent(i) for i in range(1,h.GetNbinsX()+1)] + + bins = np.array(bins) + hist = np.array(hist) + + bincenters = (bins[1:] + bins[:-1])/2 + + norm = np.trapz(hist,bincenters) + + hist /= norm + + fig = plt.figure(figsize=FIGSIZE) + plt.hist(bincenters,weights=hist,bins=bins,histtype='step') + x = np.linspace(bins[0],bins[-1],10000) + if tf1_number: + plt.plot(x,[f(xval)/norm for xval in x],color='red') + despine(fig,trim=True) + if hist_number == 1: + plt.gca().set_xlim(-1,1) + plt.ylabel("Arbitrary Units") + plt.xlabel(r"$\cos\theta$") + if args.save: + plt.savefig("%s_shower_angular_distribution.pdf" % prefix) + plt.savefig("%s_shower_angular_distribution.eps" % prefix) + else: + plt.title("%s Shower Angular Distribution" % prefix.capitalize()) + elif hist_number == 2: + plt.ylabel("Arbitrary Units") + plt.xlabel(r"Distance along Track (cm)") + if args.save: + plt.savefig("%s_shower_position_distribution.pdf" % prefix) + plt.savefig("%s_shower_position_distribution.eps" % prefix) + else: + plt.title("%s Shower Position Distribution" % prefix.capitalize()) + elif hist_number == 4: + plt.ylabel("Arbitrary Units") + plt.xlabel(r"$\cos\theta$") + if args.save: + plt.savefig("%s_delta_ray_angular_distribution.pdf" % prefix) + plt.savefig("%s_delta_ray_angular_distribution.eps" % prefix) + else: + plt.title("%s Delta Ray Angular Distribution" % prefix.capitalize()) + elif hist_number == 5: + plt.ylabel("Arbitrary Units") + plt.xlabel(r"Distance along Track (cm)") + if args.save: + plt.savefig("%s_delta_ray_position_distribution.pdf" % prefix) + plt.savefig("%s_delta_ray_position_distribution.eps" % prefix) + else: + plt.title("%s Delta Ray Position Distribution" % prefix.capitalize()) + else: + for graph_name, tf1_number, ylabel in zip(["g_dir_alpha","g_dir_beta","g_pos_alpha","g_pos_beta","g_dir_alpha_delta","g_dir_beta_delta"], + [1,2,None,None,3,4], + [r"$\alpha$",r"$\beta$",r"$k$",r"$\theta$",r"$\alpha$",r"$\beta$"]): + g = root_file.Get(graph_name) + if tf1_number: + f = g.GetFunction("f%i" % tf1_number) + + x = [g.GetX()[i] for i in range(g.GetN())] + y = [g.GetY()[i] for i in range(g.GetN())] + yerr = [g.GetEY()[i] for i in range(g.GetN())] + + x = np.array(x) + y = np.array(y) + yerr = np.array(yerr) + + fig = plt.figure(figsize=FIGSIZE) + plt.errorbar(x,y,yerr=yerr,fmt='o') + x = np.linspace(x[0],x[-1],10000) + if tf1_number: + plt.plot(x,[f(xval) for xval in x],color='red') + despine(fig,trim=True) + plt.xlabel("Kinetic Energy (MeV)") + plt.ylabel(ylabel) + + if graph_name == "g_dir_alpha": + if args.save: + plt.savefig("%s_shower_angular_distribution_alpha.pdf" % prefix) + plt.savefig("%s_shower_angular_distribution_alpha.eps" % prefix) + else: + plt.title("%s Shower Angular Distribution" % prefix.capitalize()) + elif graph_name == "g_dir_beta": + if args.save: + plt.savefig("%s_shower_angular_distribution_beta.pdf" % prefix) + plt.savefig("%s_shower_angular_distribution_beta.eps" % prefix) + else: + plt.title("%s Shower Position Distribution" % prefix.capitalize()) + elif graph_name == "g_pos_alpha": + if args.save: + plt.savefig("%s_shower_position_distribution_alpha.pdf" % prefix) + plt.savefig("%s_shower_position_distribution_alpha.eps" % prefix) + else: + plt.title("%s Shower Position Distribution" % prefix.capitalize()) + elif graph_name == "g_pos_beta": + if args.save: + plt.savefig("%s_shower_position_distribution_beta.pdf" % prefix) + plt.savefig("%s_shower_position_distribution_beta.eps" % prefix) + else: + plt.title("%s Shower Position Distribution" % prefix.capitalize()) + elif graph_name == "g_dir_alpha_delta": + if args.save: + plt.savefig("%s_delta_ray_angular_distribution_alpha.pdf" % prefix) + plt.savefig("%s_delta_ray_angular_distribution_alpha.eps" % prefix) + else: + plt.title("%s Delta Ray Angular Distribution" % prefix.capitalize()) + elif graph_name == "g_dir_beta_delta": + if args.save: + plt.savefig("%s_delta_ray_angular_distribution_beta.pdf" % prefix) + plt.savefig("%s_delta_ray_angular_distribution_beta.eps" % prefix) + else: + plt.title("%s Delta Ray Position Distribution" % prefix.capitalize()) + + plt.show() + + finally: + root_file.Close() |