Skip to content

Commit ed2c324

Browse files
njacazionjacazio
authored
Expand tool for drawing LUTs (#80)
- Saving pngs - Add possibility to select particles
1 parent f7f0a44 commit ed2c324

File tree

1 file changed

+297
-0
lines changed

1 file changed

+297
-0
lines changed

‎examples/smearing/draw.py‎

Lines changed: 297 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,297 @@
1+
#!/usr/bin/env python3
2+
3+
"""
4+
Script to plot the content of the LUT in terms of pointing resolution, efficiency, momentum resolution
5+
"""
6+
7+
from ROOT import gROOT, TLatex, TCanvas, TLegend, TColor, gPad
8+
from ROOT import TFile
9+
from os import path
10+
import argparse
11+
12+
13+
def main(reader_name,
14+
tags,
15+
lut_path,
16+
ptmin,
17+
ptmax,
18+
ymin=None,
19+
ymax=None,
20+
tag_name=None,
21+
logx=False,
22+
logy=False,
23+
leg_pos=[0.74, 0.2, 0.90, 0.4],
24+
particles=None,
25+
eta=0,
26+
rmin=None,
27+
add_eta_label=True,
28+
add_alice3_label=True,
29+
save=None,
30+
background=False,
31+
aod=None):
32+
gROOT.LoadMacro(reader_name)
33+
gROOT.LoadMacro("style.C")
34+
reader_name = reader_name.split(".")[-2]
35+
reader = getattr(__import__('ROOT', fromlist=[reader_name]),
36+
reader_name)
37+
style = getattr(__import__('ROOT', fromlist=["style"]),
38+
"style")
39+
40+
style()
41+
42+
p = {"el": "e", "pi": "#pi", "ka": "K", "pr": "p"}
43+
p_colors = {"el": "#e41a1c", "pi": "#377eb8",
44+
"ka": "#4daf4a", "pr": "#984ea3"}
45+
if particles is not None:
46+
to_remove = []
47+
for i in p:
48+
if i not in particles:
49+
to_remove.append(i)
50+
for i in to_remove:
51+
p.pop(i)
52+
53+
latex = TLatex()
54+
latex.SetTextAlign(33)
55+
canvas = reader_name
56+
canvas = canvas.replace("lutRead_", "")
57+
canvas = TCanvas(canvas, canvas, 800, 800)
58+
canvas.Divide(2, 2)
59+
drawn = [canvas]
60+
drawn_graphs = {}
61+
drawn_frames = {}
62+
if ymin is None:
63+
if "_dca" in reader_name:
64+
ymin = 0.1
65+
elif "_pt" in reader_name:
66+
ymin = 1.
67+
elif "_eff" in reader_name:
68+
ymin = 0.
69+
if ymax is None:
70+
if "_dca" in reader_name:
71+
ymax = 1e4
72+
elif "_pt" in reader_name:
73+
ymax = 100.
74+
elif "_eff" in reader_name:
75+
ymax = 115.
76+
77+
def adjust_pad():
78+
if logx:
79+
gPad.SetLogx()
80+
if logy:
81+
gPad.SetLogy()
82+
83+
counter = 1
84+
leg = None
85+
if tag_name is not None:
86+
leg = TLegend(*leg_pos)
87+
if add_eta_label:
88+
label = f"#eta = {int(eta)}"
89+
if rmin is not None:
90+
label += " R_{min} = " + rmin
91+
else:
92+
leg.SetHeader()
93+
leg.SetHeader(label)
94+
leg.SetLineColor(0)
95+
drawn.append(leg)
96+
97+
def draw_alice3_label(x=0.5, y=0.9):
98+
latex = TLatex()
99+
latex.SetTextAlign(13)
100+
drawn.append(latex.DrawLatexNDC(x, y, "ALICE 3 study"))
101+
for i in p:
102+
c = f"{canvas.GetName()}_{i}"
103+
c = TCanvas(c, c, 800, 800)
104+
drawn.append(c)
105+
adjust_pad()
106+
107+
frame = c.DrawFrame(ptmin, ymin,
108+
ptmax, ymax, "")
109+
frame.SetDirectory(0)
110+
drawn_frames[i] = frame
111+
g_list = []
112+
extra = {}
113+
cols = ['#e41a1c', '#377eb8', '#4daf4a',
114+
'#984ea3', '#ff7f00', '#ffff33']
115+
for k, j in enumerate(tags):
116+
lut = f"{lut_path}/lutCovm.{i}.{j}.dat"
117+
if not path.isfile(lut):
118+
print("LUT file", lut, "does not exist")
119+
return
120+
g = reader(lut, eta)
121+
if g.GetN() <= 0:
122+
print("Skipping", g.GetName())
123+
continue
124+
if len(g_list) == 0:
125+
frame.GetXaxis().SetTitle(g.GetXaxis().GetTitle())
126+
frame.GetYaxis().SetTitle(g.GetYaxis().GetTitle())
127+
col = TColor.GetColor(cols[len(g_list)])
128+
g.SetLineColor(col)
129+
g.SetLineStyle(1)
130+
g.SetLineWidth(3)
131+
g.Draw("samel")
132+
if aod is not None:
133+
if "_eff" in reader_name:
134+
f_aod = TFile(aod, "READ")
135+
extra[g.GetName()] = f_aod.Get(
136+
"qa-tracking-efficiency-kaon/pt/num")
137+
extra[g.GetName()].Divide(f_aod.Get("qa-tracking-efficiency-kaon/pt/num"),
138+
f_aod.Get("qa-tracking-efficiency-kaon/pt/den"), 1, 1, "B")
139+
extra[g.GetName()].Scale(100)
140+
extra[g.GetName()].Draw("SAME")
141+
extra[g.GetName()].SetDirectory(0)
142+
f_aod.Close()
143+
print("Drawing", g.GetName())
144+
if tag_name is not None and counter == 1:
145+
leg.AddEntry(g, tag_name[k], "l")
146+
g_list.append(g)
147+
drawn_graphs[i] = g_list
148+
if len(g_list) <= 0:
149+
print("Nothing drawn!")
150+
continue
151+
drawn.append(latex.DrawLatexNDC(0.9, 0.9, p[i]))
152+
if leg is not None:
153+
leg.Draw()
154+
draw_alice3_label(.4, .91)
155+
gPad.Update()
156+
canvas.cd(counter)
157+
clone = c.DrawClonePad()
158+
if counter != 1:
159+
l = gPad.GetListOfPrimitives()
160+
for i in l:
161+
cn = i.ClassName()
162+
if cn == "TLegend":
163+
l.Remove(i)
164+
elif cn == "TLatex":
165+
if "ALICE" in i.GetTitle():
166+
l.Remove(i)
167+
drawn.append(clone)
168+
c.SaveAs(f"/tmp/{c.GetName()}.png")
169+
gPad.Update()
170+
counter += 1
171+
if save is None:
172+
canvas.SaveAs(f"/tmp/lut_{canvas.GetName()}.root")
173+
else:
174+
fo = TFile(save, "RECREATE")
175+
fo.cd()
176+
canvas.Write()
177+
for i in drawn_graphs:
178+
for j in drawn_graphs[i]:
179+
j.Write()
180+
if len(tags) == 1:
181+
canvas_all_species = TCanvas("all_spec_"+canvas.GetName(),
182+
"all_spec_"+canvas.GetName(), 800, 800)
183+
drawn.append(canvas_all_species)
184+
canvas_all_species.cd()
185+
drawn_graphs_all_spec = {}
186+
leg_all_spec = TLegend(*leg_pos)
187+
leg_all_spec.SetNColumns(2)
188+
leg_all_spec.SetLineColor(0)
189+
drawn.append(leg_all_spec)
190+
for i in drawn_graphs:
191+
if canvas_all_species.GetListOfPrimitives().GetEntries() == 0:
192+
drawn_frames[i].Draw()
193+
g_list = []
194+
for j in drawn_graphs[i]:
195+
g_list.append(j.Clone())
196+
g_list[-1].SetLineColor(TColor.GetColor(p_colors[i]))
197+
g_list[-1].Draw("same")
198+
leg_all_spec.AddEntry(g_list[-1], p[i], "L")
199+
drawn_graphs_all_spec[i] = g_list
200+
leg_all_spec.Draw()
201+
if add_alice3_label:
202+
draw_alice3_label()
203+
latex = TLatex()
204+
latex.SetTextAlign(13)
205+
latex.SetTextSize(0.04)
206+
if tag_name is not None:
207+
drawn.append(latex.DrawLatexNDC(0.5, 0.80, tag_name[0]))
208+
drawn.append(latex.DrawLatexNDC(0.5, 0.75, f"#eta = {int(eta)}" +
209+
" R_{min} = " + rmin))
210+
211+
adjust_pad()
212+
canvas_all_species.Update()
213+
if not background:
214+
input("Done, press enter to continue")
215+
216+
217+
if __name__ == "__main__":
218+
parser = argparse.ArgumentParser(description=__doc__)
219+
parser.add_argument("reader", type=str,
220+
help="Reader macro to access the information e.g. lutRead_eff.C")
221+
parser.add_argument("--path", "-p",
222+
type=str,
223+
default="/tmp/myluts/",
224+
help="Path of the LUTs")
225+
parser.add_argument("--particles", "-P",
226+
type=str,
227+
nargs="+",
228+
default=None,
229+
help="Particles to show e.g. el pi ka mu pr")
230+
parser.add_argument("--ptmin",
231+
type=float,
232+
default=1e-2,
233+
help="Minimum pT of the plot")
234+
parser.add_argument("--ptmax",
235+
type=float,
236+
default=100.,
237+
help="Maximum pT of the plot")
238+
parser.add_argument("--tags", "-t", type=str, nargs="+",
239+
default=[".5kG.20cm.its3", ".5kG.50cm.its3",
240+
".5kG.100cm.its3", ".5kG.20cm.scenario3"],
241+
help="Tags to collect")
242+
parser.add_argument("--tags_name", "-T",
243+
type=str, nargs="+",
244+
default=None,
245+
help="Title of the tags that can be used in legend making")
246+
parser.add_argument("--eta",
247+
type=float,
248+
default=0,
249+
help="Eta position")
250+
parser.add_argument("--ymin",
251+
type=float,
252+
default=None,
253+
help="Minimum y")
254+
parser.add_argument("--ymax",
255+
type=float,
256+
default=None,
257+
help="Maximum y")
258+
parser.add_argument("--rmin",
259+
type=str,
260+
default=None,
261+
help="Label for the minimum radius")
262+
parser.add_argument("--aod",
263+
type=str,
264+
default=None,
265+
help="Results from aod to show")
266+
parser.add_argument("--save",
267+
type=str,
268+
default=None,
269+
help="Name to save the figure to")
270+
parser.add_argument("--leg_pos", "-l",
271+
type=float, nargs="+",
272+
default=[0.74, 0.2, 0.90, 0.4],
273+
help="Position of the legend in NDC coordinates")
274+
parser.add_argument("--logx", action="store_true",
275+
help="Log x")
276+
parser.add_argument("--logy", action="store_true",
277+
help="Log y")
278+
parser.add_argument("-b", action="store_true",
279+
help="Background mode")
280+
args = parser.parse_args()
281+
main(args.reader,
282+
args.tags,
283+
lut_path=args.path,
284+
ptmin=args.ptmin,
285+
ptmax=args.ptmax,
286+
logx=args.logx,
287+
logy=args.logy,
288+
tag_name=args.tags_name,
289+
particles=args.particles,
290+
leg_pos=args.leg_pos,
291+
ymin=args.ymin,
292+
ymax=args.ymax,
293+
rmin=args.rmin,
294+
eta=args.eta,
295+
save=args.save,
296+
background=args.b,
297+
aod=args.aod)

0 commit comments

Comments
 (0)