| 10 |
|
|
| 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) |
| 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! |
| 17 |
– |
'Eff' : y-axis will be scaled from 0 to 1 |
| 18 |
– |
'Norm': plot will be area normalized |
| 19 |
– |
'Logx': x-axis will be on log scale |
| 20 |
– |
'Logy': y-axis will be on log scale""" |
| 15 |
|
|
| 16 |
< |
## Define colors and styles |
| 16 |
> |
#### Define colors and styles |
| 17 |
|
rgbvals = [[82, 124, 219], |
| 18 |
|
[212,58,143], |
| 19 |
|
[231, 139, 77], |
| 22 |
|
[67, 77, 83]] |
| 23 |
|
marker_styles = [3, 4, 5, 25, 26, 27, 28, 30] |
| 24 |
|
|
| 25 |
< |
## Import python libraries |
| 25 |
> |
|
| 26 |
> |
#### Import python libraries |
| 27 |
|
import sys |
| 28 |
|
import optparse |
| 29 |
|
import shutil |
| 30 |
|
import os |
| 31 |
|
import re |
| 32 |
|
|
| 33 |
< |
## Import ROOT in batch mode |
| 33 |
> |
|
| 34 |
> |
#### Import ROOT in batch mode |
| 35 |
|
if '-h' not in sys.argv and len(sys.argv) > 1: |
| 36 |
|
sys.argv.append('-b') |
| 37 |
|
import ROOT |
| 46 |
|
colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals] |
| 47 |
|
canvas = ROOT.TCanvas() |
| 48 |
|
|
| 49 |
< |
## Parse options |
| 49 |
> |
|
| 50 |
> |
#### Parse options |
| 51 |
|
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") |
| 52 |
|
parser.add_option('-e', '--ext', default="pdf", |
| 53 |
|
help="choose an output extension; default is pdf") |
| 54 |
|
parser.add_option('-o', '--opt', default="", |
| 55 |
< |
help="argument to be passed to THStack::Draw") |
| 55 |
> |
help="pass OPT to the Draw command (try 'e' for error bars)") |
| 56 |
> |
parser.add_option('-m', '--markers', action="store_true", default=False, |
| 57 |
> |
help="add markers to histograms") |
| 58 |
> |
parser.add_option('-s', '--special', action="store_true", default=False, |
| 59 |
> |
help="enable name-based special plotting options " |
| 60 |
> |
"(see below)") |
| 61 |
|
parser.add_option('--output', default="overlaidHists", metavar="NAME", |
| 62 |
|
help="name of output directory; default is 'overlaidHists'") |
| 63 |
+ |
parser.add_option('--numbering', action="store_true", default=False, |
| 64 |
+ |
help="add a page number in the upper right of each plot") |
| 65 |
+ |
parser.add_option('--timing', action="store_true", default=False, |
| 66 |
+ |
help="output timing information") |
| 67 |
|
parser.add_option('--match', default="", metavar="REGEX", |
| 68 |
|
help="only make plots for paths containing the specified " |
| 69 |
|
"regular expression (use '.*' for wildcard)") |
| 70 |
< |
parser.add_option('--timing', action="store_true", default=False, |
| 71 |
< |
help="output timing information") |
| 70 |
> |
group1 = optparse.OptionGroup( |
| 71 |
> |
parser, |
| 72 |
> |
"special plotting options", |
| 73 |
> |
"These options can be applied to all plots or to only specific plots " |
| 74 |
> |
"you choose. To affect all plots, simply include the option on the " |
| 75 |
> |
"command line. With the '-s' option, however, you may instead include " |
| 76 |
> |
"a keyword (such as 'Norm') in the ROOT name of a histogram and it will " |
| 77 |
> |
"have the option turned on regardless of the command line." |
| 78 |
> |
) |
| 79 |
> |
group1.add_option('-n', '--normalize', action="store_true", default=False, |
| 80 |
> |
help="'Norm': area normalize the histograms") |
| 81 |
> |
group1.add_option('--efficiency', action="store_true", default=False, |
| 82 |
> |
help="'Eff' : force y axis scale to run from 0 to 1") |
| 83 |
> |
group1.add_option('--logx', action="store_true", default=False, |
| 84 |
> |
help="'Logx': force log scale for x axis") |
| 85 |
> |
group1.add_option('--logy', action="store_true", default=False, |
| 86 |
> |
help="'Logy': force log scale for y axis") |
| 87 |
> |
group1.add_option('--overflow', action="store_true", default=False, |
| 88 |
> |
help="'Overflow' : display overflow content in highest bin") |
| 89 |
> |
group1.add_option('--underflow', action="store_true", default=False, |
| 90 |
> |
help="'Underflow': display underflow content in lowest bin") |
| 91 |
> |
parser.add_option_group(group1) |
| 92 |
|
options, arguments = parser.parse_args() |
| 93 |
|
plot_dir = "%s/%s" % (os.path.abspath('.'), options.output) |
| 94 |
|
regex = re.compile(options.match) |
| 95 |
|
|
| 96 |
|
|
| 97 |
< |
|
| 97 |
> |
#### Define classes and utility functions |
| 98 |
|
class RootFile: |
| 99 |
+ |
"""A wrapper for TFiles, allowing quick access to the name and Get.""" |
| 100 |
|
def __init__(self, file_name): |
| 101 |
|
self.name = file_name[0:-5] |
| 102 |
|
self.file = ROOT.TFile(file_name, "read") |
| 107 |
|
return self.file.Get(object_name) |
| 108 |
|
|
| 109 |
|
|
| 110 |
+ |
def counter_generator(): |
| 111 |
+ |
"""Incremement the counter used to number plots.""" |
| 112 |
+ |
k = 0 |
| 113 |
+ |
while True: |
| 114 |
+ |
k += 1 |
| 115 |
+ |
yield k |
| 116 |
+ |
next_counter = counter_generator().next |
| 117 |
+ |
|
| 118 |
|
|
| 119 |
+ |
#### Define the main program functions |
| 120 |
|
def main(): |
| 121 |
+ |
"""Initialize files and enter into process_directory loop""" |
| 122 |
|
files = [RootFile(filename) for filename in arguments] |
| 123 |
|
if len(files) == 0: |
| 124 |
|
parser.print_help() |
| 129 |
|
merge_pdf() |
| 130 |
|
|
| 131 |
|
|
| 103 |
– |
|
| 132 |
|
def process_directory(path, files): |
| 133 |
+ |
"""Loop through all histograms in the directory and plot them.""" |
| 134 |
|
dir_to_make = "%s/%s" % (plot_dir, path) |
| 135 |
|
if not os.path.exists(dir_to_make): |
| 136 |
|
os.mkdir(dir_to_make) |
| 142 |
|
new_path = "%s/%s" % (path, obj.GetName()) |
| 143 |
|
if obj.IsA().InheritsFrom("TDirectory"): |
| 144 |
|
process_directory(new_path, files) |
| 145 |
+ |
#### If obj is a desired histogram, process it |
| 146 |
|
if (regex.search(new_path) and |
| 147 |
|
obj.IsA().InheritsFrom("TH1") and |
| 148 |
|
not obj.IsA().InheritsFrom("TH2") and |
| 149 |
|
not obj.IsA().InheritsFrom("TH3")): |
| 150 |
< |
counter = next_counter() |
| 151 |
< |
name = obj.GetName() |
| 122 |
< |
hist = files[0].file.GetDirectory(path).Get(name) |
| 123 |
< |
title = hist.GetTitle() |
| 124 |
< |
x_title = hist.GetXaxis().GetTitle() |
| 125 |
< |
y_title = hist.GetYaxis().GetTitle() |
| 126 |
< |
if "Norm" in name or options.normalize: |
| 127 |
< |
y_title = "Fraction of Events in Bin" |
| 128 |
< |
stack = ROOT.THStack("st%.3i" % int(counter), title) |
| 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 |
| 136 |
< |
hist.SetTitle(file.name) |
| 137 |
< |
color = colors[i % len(colors)] |
| 138 |
< |
hist.SetLineColor(color) |
| 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) |
| 154 |
< |
stack.Add(hist) |
| 155 |
< |
legend.AddEntry(hist) |
| 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 |
< |
save_plot(stack, plot_dir, path, name, counter) |
| 186 |
< |
|
| 150 |
> |
process_hist(path, new_path, files, obj) |
| 151 |
> |
|
| 152 |
|
|
| 153 |
+ |
#### This is where all the plotting actually happens |
| 154 |
+ |
def process_hist(path, new_path, files, obj): |
| 155 |
+ |
"""Overlay the plots and apply options.""" |
| 156 |
+ |
counter = next_counter() |
| 157 |
+ |
name = obj.GetName() |
| 158 |
+ |
hist = files[0].file.GetDirectory(path).Get(name) |
| 159 |
+ |
title = hist.GetTitle() |
| 160 |
+ |
x_title = hist.GetXaxis().GetTitle() |
| 161 |
+ |
y_title = hist.GetYaxis().GetTitle() |
| 162 |
+ |
if options.normalize or (options.special and "Norm" in name): |
| 163 |
+ |
y_title = "Fraction of Events in Bin" |
| 164 |
+ |
stack = ROOT.THStack("st%.3i" % int(counter), title) |
| 165 |
+ |
legend_height = 0.04 * len(files) + 0.02 |
| 166 |
+ |
legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89) |
| 167 |
+ |
canvas.SetLogx(options.logx or (options.special and "Logx" in name)) |
| 168 |
+ |
canvas.SetLogy(options.logy or (options.special and "Logy" in name)) |
| 169 |
+ |
#### Fill the THStack with histograms from each file |
| 170 |
+ |
for i, file in enumerate(files): |
| 171 |
+ |
hist = file.file.GetDirectory(path).Get(name) |
| 172 |
+ |
if not hist: continue |
| 173 |
+ |
hist.SetTitle(file.name) |
| 174 |
+ |
color = colors[i % len(colors)] |
| 175 |
+ |
hist.SetLineColor(color) |
| 176 |
+ |
if options.markers: |
| 177 |
+ |
hist.SetMarkerColor(color) |
| 178 |
+ |
hist.SetMarkerStyle(marker_styles[i]) |
| 179 |
+ |
else: |
| 180 |
+ |
hist.SetMarkerSize(0) |
| 181 |
+ |
if options.overflow or (options.special and "Overflow" in name): |
| 182 |
+ |
nbins = hist.GetNbinsX() |
| 183 |
+ |
overflow = hist.GetBinContent(nbins + 1) |
| 184 |
+ |
hist.AddBinContent(nbins, overflow) |
| 185 |
+ |
if options.underflow or (options.special and "Underflow" in name): |
| 186 |
+ |
underflow = hist.GetBinContent(0) |
| 187 |
+ |
hist.AddBinContent(1, underflow) |
| 188 |
+ |
if options.normalize or (options.special and "Norm" in name): |
| 189 |
+ |
integral = hist.Integral() |
| 190 |
+ |
hist.Scale(1. / integral) |
| 191 |
+ |
stack.Add(hist) |
| 192 |
+ |
legend.AddEntry(hist) |
| 193 |
+ |
#### Draw the stack and apply text overlays |
| 194 |
+ |
stack.Draw("nostack p H" + options.opt) |
| 195 |
+ |
stack.SetTitle("%s;%s;%s" % (title, x_title, y_title)) |
| 196 |
+ |
if options.numbering: |
| 197 |
+ |
display_page_number(counter) |
| 198 |
+ |
if options.efficiency or (options.special and "Eff" in name): |
| 199 |
+ |
stack.Draw("nostack e p") |
| 200 |
+ |
stack.SetMaximum(1.) |
| 201 |
+ |
stack.SetMinimum(0.) |
| 202 |
+ |
if options.overflow or (options.special and "Overflow" in name): |
| 203 |
+ |
display_overflow(stack, hist) |
| 204 |
+ |
if options.underflow or (options.special and "Underflow" in name): |
| 205 |
+ |
display_underflow(stack, hist) |
| 206 |
+ |
legend.Draw() |
| 207 |
+ |
save_plot(stack, plot_dir, path, name, counter) |
| 208 |
|
|
| 209 |
+ |
|
| 210 |
+ |
#### Define the extra functions used by the main routines |
| 211 |
|
def save_plot(stack, plot_dir, path, name, counter): |
| 212 |
|
output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext) |
| 213 |
|
canvas.SaveAs(output_file_name) |
| 217 |
|
report_progress(counter, 1) |
| 218 |
|
|
| 219 |
|
|
| 198 |
– |
|
| 220 |
|
def report_progress(counter, divisor): |
| 221 |
|
if counter % divisor == 0: |
| 222 |
|
print "\r%i plots written to %s" % (counter, options.output), |
| 223 |
|
sys.stdout.flush() |
| 224 |
|
|
| 204 |
– |
|
| 205 |
– |
|
| 225 |
|
def merge_pdf(): |
| 226 |
|
print "Writing merged pdf..." |
| 227 |
|
os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite " |
| 231 |
|
os.system("rm [0-9]*.pdf") |
| 232 |
|
|
| 233 |
|
|
| 234 |
< |
|
| 235 |
< |
def counter_generator(): |
| 236 |
< |
k = 0 |
| 237 |
< |
while True: |
| 238 |
< |
k += 1 |
| 239 |
< |
yield k |
| 240 |
< |
next_counter = counter_generator().next |
| 241 |
< |
|
| 234 |
> |
def display_page_number(page_number): |
| 235 |
> |
page_text = ROOT.TText() |
| 236 |
> |
page_text.SetTextSize(0.03) |
| 237 |
> |
page_text.SetTextAlign(33) |
| 238 |
> |
page_text.DrawTextNDC(0.97, 0.985, "%i" % page_number) |
| 239 |
> |
|
| 240 |
> |
|
| 241 |
> |
def display_overflow(stack, hist): |
| 242 |
> |
nbins = hist.GetNbinsX() |
| 243 |
> |
x = 0.5 * (hist.GetBinLowEdge(nbins) + |
| 244 |
> |
hist.GetBinLowEdge(nbins + 1)) |
| 245 |
> |
y = stack.GetMinimum("nostack") |
| 246 |
> |
display_bin_text(x, y, nbins, "Overflow") |
| 247 |
> |
|
| 248 |
> |
|
| 249 |
> |
def display_underflow(stack, hist): |
| 250 |
> |
nbins = hist.GetNbinsX() |
| 251 |
> |
x = 0.5 * (hist.GetBinLowEdge(1) + |
| 252 |
> |
hist.GetBinLowEdge(2)) |
| 253 |
> |
y = stack.GetMinimum("nostack") |
| 254 |
> |
display_bin_text(x, y, nbins, "Underflow") |
| 255 |
> |
|
| 256 |
> |
|
| 257 |
> |
def display_bin_text(x, y, nbins, text): |
| 258 |
> |
bin_text = ROOT.TText() |
| 259 |
> |
bin_text.SetTextSize(min(1. / nbins, 0.04)) |
| 260 |
> |
bin_text.SetTextAlign(12) |
| 261 |
> |
bin_text.SetTextAngle(90) |
| 262 |
> |
bin_text.SetTextColor(13) |
| 263 |
> |
bin_text.SetTextFont(42) |
| 264 |
> |
bin_text.DrawText(x, y, text) |
| 265 |
|
|
| 266 |
|
|
| 267 |
+ |
#### Run main program |
| 268 |
|
if __name__ == "__main__": |
| 269 |
|
if options.timing: |
| 227 |
– |
## import cProfile |
| 228 |
– |
## cProfile.run('main()', 'fooprof') |
| 270 |
|
import profile |
| 271 |
|
profile.run('main()', 'fooprof') |
| 272 |
|
import pstats |
