-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathPlotter.py
143 lines (123 loc) · 6.01 KB
/
Plotter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import copy
import importlib
import os
import ROOT
import sys
import yaml
def mk_smart_hist(hist, hist_desc):
use_log_y = hist_desc.get('use_log_y', False)
use_log_x = hist_desc.get('use_log_x', False)
max_y_sf = hist_desc.get('max_y_sf', 1.5)
divide_by_bin_width = hist_desc.get('divide_by_bin_width', False)
smart_hist = ROOT.root_ext.SmartHistogram('TH1D')(hist, use_log_y,use_log_x, max_y_sf, divide_by_bin_width)
if 'x_title' in hist_desc:
smart_hist.GetXaxis().SetTitle(hist_desc['x_title'])
if 'y_title' in hist_desc:
smart_hist.GetYaxis().SetTitle(hist_desc['y_title'])
return smart_hist
def to_str(value):
if isinstance(value, bool):
return 'true' if value else 'false'
if isinstance(value, list):
return ','.join(map(to_str, value))
return str(value)
def ToItem(entry):
map = ROOT.std.map('std::string', 'std::string')()
for key, value in entry.items():
map[key] = to_str(value)
item = ROOT.root_ext.draw_options.Item()
item.properties = ROOT.analysis.PropertyList(map)
return item
def LoadHistOptions(entry):
item = ToItem(entry)
return ROOT.root_ext.draw_options.Histogram(item)
def LoadPageOptions(entry):
item = ToItem(entry)
return ROOT.root_ext.draw_options.Page(item)
class Plotter(object):
initialized = False
def read_config(self, config):
if type(config)==str:
with open(config, 'r') as f:
return yaml.safe_load(f)
elif type(config) == dict:
return config
elif type(config) == list:
return config
else:
raise RuntimeError(f'Unknown config type: {type(config)}')
def __init__(self, page_cfg, page_cfg_custom, hist_cfg, inputs_cfg):
if not Plotter.initialized:
ROOT.gROOT.SetBatch(True)
python_dir = os.path.dirname(os.path.abspath(__file__))
header_dir = os.path.join(python_dir, "include")
ROOT.gInterpreter.Declare(f'#include "{header_dir}/PdfPrinter.h"')
ROOT.gInterpreter.Declare(f'#include "{header_dir}/StackedPlotDescriptor.h"')
Plotter.initialized = True
self.page_cfg = self.read_config(page_cfg)
if page_cfg_custom:
self.page_cfg.update(self.read_config(page_cfg_custom))
self.hist_cfg = self.read_config(hist_cfg)
self.inputs_cfg = self.read_config(inputs_cfg)
self.sgn_hist_opt = LoadHistOptions(self.page_cfg['sgn_hist'])
self.bkg_hist_opt = LoadHistOptions(self.page_cfg['bkg_hist'])
self.data_hist_opt = LoadHistOptions(self.page_cfg['data_hist'])
self.bkg_unc_hist_opt = LoadHistOptions(self.page_cfg['bkg_unc_hist'])
self.page = LoadPageOptions(self.page_cfg['page_setup'])
def plot(self, hist_name, histograms, output_file, want_data=True, custom=None):
page_cfg = copy.deepcopy(self.page_cfg)
if custom:
for key, value in custom.items():
page_cfg[key]['text'] = value
items = ROOT.root_ext.draw_options.ItemCollection()
for element in page_cfg['page_setup']['text_boxes'] + [ page_cfg['page_setup']['legend'] ]:
items[element] = ToItem(page_cfg[element])
printer = ROOT.root_ext.PdfPrinter(output_file, items, self.page)
desc = ROOT.root_ext.StackedPlotDescriptor(self.page, self.sgn_hist_opt, self.bkg_hist_opt, self.data_hist_opt,
self.bkg_unc_hist_opt)
smart_hists = {}
for input in self.inputs_cfg:
name = input['name']
smart_hists[name] = mk_smart_hist(histograms[name], self.hist_cfg[hist_name])
hist_type = input.get('type', 'background')
if hist_type == 'signal':
desc.AddSignalHistogram(smart_hists[name], input['title'], ROOT.root_ext.Color.Parse(input['color']),
input.get('scale', 1.))
elif hist_type == 'background':
desc.AddBackgroundHistogram(smart_hists[name], input['title'], ROOT.root_ext.Color.Parse(input['color']))
elif hist_type == 'data':
if want_data:
desc.AddDataHistogram(smart_hists[name], input['title'])
else:
raise RuntimeError(f'Unknown histogram type: {hist_type}')
printer.Print(hist_name, desc, True)
def load(module_file):
if not os.path.exists(module_file):
raise RuntimeError(f"Cannot find path to {module_file}.")
module_name, module_ext = os.path.splitext(module_file)
spec = importlib.util.spec_from_file_location(module_name, module_file)
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
return module
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description='Plotting tool.')
parser.add_argument('--page-cfg', required=True, type=str, help="Page style config file")
parser.add_argument('--page-cfg-custom', required=False, default=None, type=str,
help="Additional page style customisations")
parser.add_argument('--hist-cfg', required=True, type=str, help="Histogram style config file")
parser.add_argument('--inputs-cfg', required=True, type=str, help="Inputs style config file")
parser.add_argument('--hist-name', required=True, type=str, help="Histogram name")
parser.add_argument('--custom', required=False, default=None, type=str, help="Customizations in format key1=value1,key2=value2,...")
parser.add_argument('--output', required=True, type=str, help="Output pdf file.")
parser.add_argument('--hist-maker', required=True, type=str, help="Path to histogram maker module.")
parser.add_argument('--verbose', required=False, type=int, default=0, help="verbosity level")
parser.add_argument('hist_maker_args', type=str, nargs='*', help="hist maker arguments")
args = parser.parse_args()
plotter = Plotter(page_cfg=args.page_cfg, page_cfg_custom=args.page_cfg_custom, hist_cfg=args.hist_cfg,
inputs_cfg=args.inputs_cfg)
hist_maker = load(args.hist_maker)
hists = hist_maker.make_histograms(*args.hist_maker_args, hist_name=args.hist_name, hist_cfg=plotter.hist_cfg)
custom = None if args.custom is None else dict(item.split('=') for item in args.custom.split(','))
plotter.plot(args.hist_name, hists, args.output, custom=custom)