#!/usr/bin/env python # Copyright (c) 2019, Anthony Latorre # # 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 . from __future__ import print_function, division import numpy as np if __name__ == '__main__': import ROOT import argparse from os.path import split from sddm.plot import despine from sddm import setup_matplotlib 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() setup_matplotlib(args.save) import matplotlib.pyplot as plt 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() 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() 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()