| 8 |
|
## Define usage string for help option |
| 9 |
|
usage="""usage: %prog [options] file1.root file2.root file3.root ... |
| 10 |
|
|
| 11 |
< |
function: overlays histograms from several files with identical structure |
| 11 |
> |
function: overlays corresponding histograms from several files, dumping the |
| 12 |
> |
images into an identical directory structure in the local directory |
| 13 |
> |
and also merging all images into a single file (if output is pdf) |
| 14 |
|
|
| 15 |
|
naming: histograms whose names contain certain key terms will be handled |
| 16 |
|
specially. Use this to your advantage! |
| 19 |
|
'Logx': x-axis will be on log scale |
| 20 |
|
'Logy': y-axis will be on log scale""" |
| 21 |
|
|
| 22 |
< |
## Define colors |
| 23 |
< |
rgbcolors = [[82, 124, 219], |
| 24 |
< |
[145, 83, 207], |
| 25 |
< |
[231, 139, 77], |
| 26 |
< |
[114, 173, 117], |
| 27 |
< |
[67, 77, 83]] |
| 22 |
> |
## Define colors and styles |
| 23 |
> |
rgbvals = [[82, 124, 219], |
| 24 |
> |
[212,58,143], |
| 25 |
> |
[231, 139, 77], |
| 26 |
> |
[145, 83, 207], |
| 27 |
> |
[114, 173, 117], |
| 28 |
> |
[67, 77, 83]] |
| 29 |
> |
marker_styles = [3, 4, 5, 25, 26, 27, 28, 30] |
| 30 |
|
|
| 31 |
|
## Import python libraries |
| 32 |
|
import sys |
| 33 |
|
import optparse |
| 34 |
+ |
import shutil |
| 35 |
|
import os |
| 36 |
|
import re |
| 37 |
|
|
| 38 |
|
## Import ROOT in batch mode |
| 39 |
< |
if '-h' not in sys.argv: |
| 39 |
> |
if '-h' not in sys.argv and len(sys.argv) > 1: |
| 40 |
|
sys.argv.append('-b') |
| 41 |
|
import ROOT |
| 42 |
< |
if os.path.exists('rootlogon.C'): ROOT.gROOT.Macro('rootlogon.C') |
| 42 |
> |
if os.path.exists('rootlogon.C'): |
| 43 |
> |
ROOT.gROOT.Macro('rootlogon.C') |
| 44 |
> |
else: |
| 45 |
> |
os.system('echo -e "{\n}\n" >> rootlogon.C') |
| 46 |
> |
ROOT.gROOT.Macro('rootlogon.C') |
| 47 |
> |
os.remove('rootlogon.C') |
| 48 |
|
sys.argv.remove('-b') |
| 49 |
|
ROOT.gErrorIgnoreLevel = ROOT.kWarning |
| 50 |
< |
colors = [] |
| 51 |
< |
for rgb in rgbcolors: |
| 42 |
< |
colors.append(ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2])) |
| 43 |
< |
c1 = ROOT.TCanvas() |
| 50 |
> |
colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals] |
| 51 |
> |
canvas = ROOT.TCanvas() |
| 52 |
|
|
| 53 |
|
## Parse options |
| 54 |
|
parser = optparse.OptionParser(usage=usage) |
| 55 |
|
parser.add_option('-n', '--normalize', action="store_true", default=False, |
| 56 |
|
help="area normalize all histograms") |
| 57 |
+ |
parser.add_option('-m', '--markers', action="store_true", default=False, |
| 58 |
+ |
help="add markers to histograms") |
| 59 |
+ |
parser.add_option('--logx', action="store_true", default=False, |
| 60 |
+ |
help="Force log scale for X axis on all histograms") |
| 61 |
+ |
parser.add_option('--logy', action="store_true", default=False, |
| 62 |
+ |
help="Force log scale for Y axis on all histograms") |
| 63 |
|
parser.add_option('-e', '--ext', default="pdf", |
| 64 |
< |
help="specify the type (extension) of the output files") |
| 65 |
< |
parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME", |
| 66 |
< |
help="specify the name of the output file/directory") |
| 67 |
< |
parser.add_option('-m', '--match', default="", metavar="REGEX", |
| 68 |
< |
help="only make plots for paths containing REGEX") |
| 64 |
> |
help="choose an output extension; default is pdf") |
| 65 |
> |
parser.add_option('-o', '--opt', default="", |
| 66 |
> |
help="argument to be passed to THStack::Draw") |
| 67 |
> |
parser.add_option('--output', default="overlaidHists", metavar="NAME", |
| 68 |
> |
help="name of output directory; default is 'overlaidHists'") |
| 69 |
> |
parser.add_option('--match', default="", metavar="REGEX", |
| 70 |
> |
help="only make plots for paths containing the specified " |
| 71 |
> |
"regular expression (use '.*' for wildcard)") |
| 72 |
> |
parser.add_option('--timing', action="store_true", default=False, |
| 73 |
> |
help="output timing information") |
| 74 |
|
options, arguments = parser.parse_args() |
| 75 |
|
plot_dir = "%s/%s" % (os.path.abspath('.'), options.output) |
| 76 |
|
regex = re.compile(options.match) |
| 77 |
|
|
| 78 |
|
|
| 79 |
+ |
|
| 80 |
|
class RootFile: |
| 81 |
|
def __init__(self, file_name): |
| 82 |
< |
self.name = file_name[0:file_name.find(".root")] |
| 82 |
> |
self.name = file_name[0:-5] |
| 83 |
|
self.file = ROOT.TFile(file_name, "read") |
| 84 |
|
if self.file.IsZombie(): |
| 85 |
|
print "Error opening %s, exiting..." % file_name |
| 90 |
|
|
| 91 |
|
|
| 92 |
|
def main(): |
| 93 |
< |
files = [] |
| 94 |
< |
for filename in arguments: files.append(RootFile(filename)) |
| 93 |
> |
files = [RootFile(filename) for filename in arguments] |
| 94 |
> |
if len(files) == 0: |
| 95 |
> |
parser.print_help() |
| 96 |
> |
sys.exit(0) |
| 97 |
|
process_directory("", files) |
| 98 |
+ |
print "" |
| 99 |
|
if options.ext == "pdf": |
| 100 |
< |
os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite " |
| 78 |
< |
"-dAutoRotatePages=/All " |
| 79 |
< |
"-sOutputFile=%s.pdf " % options.output + |
| 80 |
< |
"[0-9][0-9][0-9].pdf") |
| 81 |
< |
os.system("rm [0-9]*.pdf") |
| 82 |
< |
print "Wrote %i plots to %s" % (next_counter() - 1, options.output) |
| 100 |
> |
merge_pdf() |
| 101 |
|
|
| 102 |
|
|
| 103 |
|
|
| 125 |
|
y_title = hist.GetYaxis().GetTitle() |
| 126 |
|
if "Norm" in name or options.normalize: |
| 127 |
|
y_title = "Fraction of Events in Bin" |
| 110 |
– |
hist.Draw() |
| 111 |
– |
hists = [] |
| 128 |
|
stack = ROOT.THStack("st%.3i" % int(counter), title) |
| 129 |
< |
legend = ROOT.TLegend(0.65, 0.77, 0.87, 0.89) |
| 130 |
< |
c1.SetLogx("Logx" in name) |
| 131 |
< |
c1.SetLogy("Logy" in name) |
| 129 |
> |
legend_height = 0.04 * len(files) + 0.02 |
| 130 |
> |
legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89) |
| 131 |
> |
canvas.SetLogx("Logx" in name or options.logx) |
| 132 |
> |
canvas.SetLogy("Logy" in name or options.logy) |
| 133 |
|
for i, file in enumerate(files): |
| 134 |
|
hist = file.file.GetDirectory(path).Get(name) |
| 135 |
|
if not hist: continue |
| 119 |
– |
hist.Draw() |
| 136 |
|
hist.SetTitle(file.name) |
| 137 |
|
color = colors[i % len(colors)] |
| 138 |
|
hist.SetLineColor(color) |
| 139 |
< |
hist.SetMarkerColor(color) |
| 140 |
< |
hist.SetMarkerStyle(i + 1) |
| 139 |
> |
if options.markers: |
| 140 |
> |
hist.SetMarkerColor(color) |
| 141 |
> |
hist.SetMarkerStyle(marker_styles[i]) |
| 142 |
> |
else: |
| 143 |
> |
hist.SetMarkerSize(0) |
| 144 |
> |
if "Overflow" in name: |
| 145 |
> |
nbins = hist.GetNbinsX() |
| 146 |
> |
overflow = hist.GetBinContent(nbins + 1) |
| 147 |
> |
hist.AddBinContent(nbins, overflow) |
| 148 |
> |
if "Underflow" in name: |
| 149 |
> |
underflow = hist.GetBinContent(0) |
| 150 |
> |
hist.AddBinContent(1, underflow) |
| 151 |
|
if "Norm" in name or options.normalize: |
| 152 |
|
integral = hist.Integral() |
| 153 |
< |
hist.Scale(1 / integral) |
| 153 |
> |
hist.Scale(1. / integral) |
| 154 |
|
stack.Add(hist) |
| 155 |
|
legend.AddEntry(hist) |
| 156 |
< |
stack.Draw("nostack p H") |
| 156 |
> |
stack.Draw("nostack p H" + options.opt) |
| 157 |
|
stack.SetTitle("%s;%s;%s" % (title, x_title, y_title)) |
| 158 |
|
if "Eff" in name: |
| 159 |
|
stack.Draw("nostack e p") |
| 160 |
|
stack.SetMaximum(1.) |
| 161 |
|
stack.SetMinimum(0.) |
| 162 |
+ |
if "Overflow" in name: |
| 163 |
+ |
nbins = hist.GetNbinsX() |
| 164 |
+ |
x = 0.5 * (hist.GetBinLowEdge(nbins) + |
| 165 |
+ |
hist.GetBinLowEdge(nbins + 1)) |
| 166 |
+ |
y = stack.GetMinimum("nostack") |
| 167 |
+ |
overflow_text = ROOT.TText(x, y, "Overflow") |
| 168 |
+ |
overflow_text.SetTextSize(min(1. / nbins, 0.04)) |
| 169 |
+ |
overflow_text.SetTextAlign(12) |
| 170 |
+ |
overflow_text.SetTextAngle(90) |
| 171 |
+ |
overflow_text.SetTextColor(13) |
| 172 |
+ |
overflow_text.SetTextFont(42) |
| 173 |
+ |
overflow_text.Draw() |
| 174 |
+ |
if "Underflow" in name: |
| 175 |
+ |
nbins = hist.GetNbinsX() |
| 176 |
+ |
x = hist.GetBinLowEdge(1) |
| 177 |
+ |
y = hist.GetMinimum() |
| 178 |
+ |
underflow_text = ROOT.TText(x, y, "Underflow") |
| 179 |
+ |
underflow_text.SetTextSize(1. / nbins) |
| 180 |
+ |
underflow_text.SetTextAlign(13) |
| 181 |
+ |
underflow_text.SetTextAngle(90) |
| 182 |
+ |
underflow_text.SetTextColor(13) |
| 183 |
+ |
underflow_text.Draw() |
| 184 |
|
legend.Draw() |
| 185 |
< |
if options.ext == "pdf": |
| 138 |
< |
c1.SaveAs("%.3i.pdf" % counter) |
| 139 |
< |
c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext)) |
| 185 |
> |
save_plot(stack, plot_dir, path, name, counter) |
| 186 |
|
|
| 187 |
|
|
| 188 |
|
|
| 189 |
+ |
def save_plot(stack, plot_dir, path, name, counter): |
| 190 |
+ |
output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext) |
| 191 |
+ |
canvas.SaveAs(output_file_name) |
| 192 |
+ |
if options.ext == "pdf": |
| 193 |
+ |
numbered_pdf_name = "%.3i.pdf" % counter |
| 194 |
+ |
shutil.copy(output_file_name, numbered_pdf_name) |
| 195 |
+ |
report_progress(counter, 1) |
| 196 |
+ |
|
| 197 |
+ |
|
| 198 |
+ |
|
| 199 |
+ |
def report_progress(counter, divisor): |
| 200 |
+ |
if counter % divisor == 0: |
| 201 |
+ |
print "\r%i plots written to %s" % (counter, options.output), |
| 202 |
+ |
sys.stdout.flush() |
| 203 |
+ |
|
| 204 |
+ |
|
| 205 |
+ |
|
| 206 |
+ |
def merge_pdf(): |
| 207 |
+ |
print "Writing merged pdf..." |
| 208 |
+ |
os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite " |
| 209 |
+ |
"-dAutoRotatePages=/All " |
| 210 |
+ |
"-sOutputFile=%s.pdf " % options.output + |
| 211 |
+ |
"[0-9][0-9][0-9].pdf") |
| 212 |
+ |
os.system("rm [0-9]*.pdf") |
| 213 |
+ |
|
| 214 |
+ |
|
| 215 |
|
|
| 216 |
|
def counter_generator(): |
| 217 |
|
k = 0 |
| 222 |
|
|
| 223 |
|
|
| 224 |
|
|
| 153 |
– |
|
| 225 |
|
if __name__ == "__main__": |
| 226 |
< |
sys.exit(main()) |
| 226 |
> |
if options.timing: |
| 227 |
> |
## import cProfile |
| 228 |
> |
## cProfile.run('main()', 'fooprof') |
| 229 |
> |
import profile |
| 230 |
> |
profile.run('main()', 'fooprof') |
| 231 |
> |
import pstats |
| 232 |
> |
p = pstats.Stats('fooprof') |
| 233 |
> |
p.sort_stats('cumulative').print_stats(15) |
| 234 |
> |
else: |
| 235 |
> |
sys.exit(main()) |
| 236 |
|
|
