make_resolution.py
Go to the documentation of this file.
1 import ROOT as RT
2 from array import array
3 from argparse import ArgumentParser as ap
4 parser = ap()
5 
6 parser.add_argument("-i", type=str, help='Input file', default="")
7 parser.add_argument("-o", type=str, help='Output file', default="beam_res.root")
8 parser.add_argument("-b", type=str, help='Binning', default="75, -.15, .15")
9 parser.add_argument("--b2", type=str, help='Binning', default="75, -.15, .15")
10 parser.add_argument("--eb", type=str, help='Binning for electrons',
11  default="75, -.15, .5")
12 parser.add_argument("--eb2", type=str, help='Binning for electrons',
13  default="75, -.15, .5")
14 args = parser.parse_args()
15 
16 if not args.i:
17  exit()
18 
19 RT.gROOT.SetBatch(1)
20 
21 inputFile = RT.TFile(args.i, "OPEN")
22 tree = inputFile.Get("tree")
23 
24 outputFile = RT.TFile(args.o, "RECREATE")
25 
26 tree.Draw("(reco_p*1.e3 - NP04front_p)/NP04front_p>>hPionsRes(" + args.b + ")", "PDG == 211 || PDG == -13")
27 tree.Draw("(reco_p*1.e3 - NP04front_p)/NP04front_p>>hProtonsRes(" + args.b + ")", "PDG == 2212")
28 tree.Draw("(reco_p*1.e3 - NP04front_p)/NP04front_p>>hElectronsRes(" + args.eb + ")", "PDG == -11")
29 tree.Draw("(reco_p*1.e3 - NP04front_p)/NP04front_p>>hKaonsRes(" + args.b + ")", "PDG == 321")
30 
31 tree.Draw("NP04front_p*1.e-3:reco_p>>hPionsRes2D(" + args.b2 + ")", "PDG == 211 || PDG == -13")
32 tree.Draw("NP04front_p*1.e-3:reco_p>>hProtonsRes2D(" + args.b2 + ")", "PDG == 2212")
33 tree.Draw("NP04front_p*1.e-3:reco_p>>hElectronsRes2D(" + args.eb2 + ")", "PDG == -11")
34 tree.Draw("NP04front_p*1.e-3:reco_p>>hKaonsRes2D(" + args.b2 + ")", "PDG == 321")
35 
36 hPionsRes = RT.gDirectory.Get("hPionsRes")
37 hProtonsRes = RT.gDirectory.Get("hProtonsRes")
38 hElectronsRes = RT.gDirectory.Get("hElectronsRes")
39 hKaonsRes = RT.gDirectory.Get("hKaonsRes")
40 
41 hPionsRes2D = RT.gDirectory.Get("hPionsRes2D")
42 hProtonsRes2D = RT.gDirectory.Get("hProtonsRes2D")
43 hElectronsRes2D = RT.gDirectory.Get("hElectronsRes2D")
44 hKaonsRes2D = RT.gDirectory.Get("hKaonsRes2D")
45 
46 outputFile.cd()
47 
48 ##New: trying to fill every bin
49 hs = [hPionsRes2D, hProtonsRes2D, hElectronsRes2D, hKaonsRes2D]
50 for h in hs:
51  for i in range(1, h.GetNbinsX()+1):
52  for j in range(1, h.GetNbinsY()+1):
53  if h.GetBinContent(i, j) < 1.: h.SetBinContent(i, j, 1)
54 
55 hPionsRes.Write()
56 hProtonsRes.Write()
57 hElectronsRes.Write()
58 hKaonsRes.Write()
59 hPionsRes2D.Write()
60 hProtonsRes2D.Write()
61 hElectronsRes2D.Write()
62 hKaonsRes2D.Write()
63 
64 ##New: trying to make a variation
65 for h in hs:
66  nbinsx = h.GetNbinsX()
67  nbinsy = h.GetNbinsY()
68 
69  hPlus = RT.TH2D(h.GetName() + "Plus", "",
70  nbinsx, h.GetXaxis().GetBinLowEdge(1), h.GetXaxis().GetBinUpEdge(nbinsx),
71  nbinsy, h.GetYaxis().GetBinLowEdge(1), h.GetYaxis().GetBinUpEdge(nbinsy))
72  hMinus = RT.TH2D(h.GetName() + "Minus", "",
73  nbinsx, h.GetXaxis().GetBinLowEdge(1), h.GetXaxis().GetBinUpEdge(nbinsx),
74  nbinsy, h.GetYaxis().GetBinLowEdge(1), h.GetYaxis().GetBinUpEdge(nbinsy))
75 
76  for i in range(1, h.GetNbinsX()+1):
77  for j in range(38, 64):
78  hPlus.SetBinContent(i, j+2, h.GetBinContent(i, j))
79  hMinus.SetBinContent(i, j-2, h.GetBinContent(i, j))
80  for j in range(1, nbinsy+1):
81  if hPlus.GetBinContent(i, j) < 1.:
82  hPlus.SetBinContent(i, j, h.GetBinContent(i,j))
83  if hMinus.GetBinContent(i, j) < 1.:
84  hMinus.SetBinContent(i, j, h.GetBinContent(i,j))
85  hPlus.Write()
86  hMinus.Write()
87 outputFile.Close()