#!/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() 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207