#!/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
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()
