reco_momentum_tuples.py
Go to the documentation of this file.
1 import ROOT as RT
2 from math import floor, acos, cos, sqrt, tan, ceil
3 from glob import glob as ls
4 from array import array
5 from argparse import ArgumentParser as ap
6 parser = ap()
7 
8 parser.add_argument("-i", type=str, help='Input cmd', default="")
9 parser.add_argument("-o", type=str, help='Output file', default="beam_tuples.root")
10 parser.add_argument("-m", type=int, help='Max number of files', default=1000)
11 parser.add_argument("-p", type=str, help='Momentum setting', default="1")
12 parser.add_argument("-s", type=float, help='1sigma shift', default=.18)
13 args = parser.parse_args()
14 
15 RT.gROOT.SetBatch(1)
16 
17 mag_P1 = 5.82044830e-3;
18 mag_P3 = -4.68880000e-6;
19 mag_P4 = 324.573967;
20 
21 if not args.i:
22  print("Error: input cmd invalid")
23  exit()
24 files = ls(args.i)
25 t = RT.TChain("NTuples/GoodParticle")
26 
27 max_files = args.m
28 
29 
30 nFiles = 0
31 for f in files:
32 
33  if(nFiles > max_files): break
34 
35  print("Adding", f)
36  t.Add(f)
37  nFiles = nFiles + 1
38 
39 def shift_x(x):
40  return ( floor(x) + .5 )
41 
42 def to_f(x):
43  if x > 0.: return int(ceil(x))
44  else: return int(floor(x))
45 
46 def momentum_costheta(x1,x2,x3):
47  L1 = 1.98
48  L2 = 1.69472
49  L3 = 2.11666
50  fBeamBend = .12003
51 
52  a = (x2*L3 - x3*L2)*cos(fBeamBend)/(L3-L2);
53 
54 
55  numTerm = (a - x1)*( (L3 - L2)*tan(fBeamBend) + (x3 - x2)*cos(fBeamBend) ) + L1*( L3 - L2 );
56 
57  denomTerm1 = sqrt( L1*L1 + (a - x1)*(a - x1) );
58  denomTerm2 = sqrt( ( (L3 - L2)*tan(fBeamBend) + (x3 - x2)*cos(fBeamBend) )**2 + ( (L3 - L2) )**2 );
59  denom = denomTerm1 * denomTerm2;
60 
61  cosTheta = numTerm/denom;
62  return cosTheta;
63 
64 
65 fout = RT.TFile(args.o, "RECREATE")
66 h = RT.TH1D("h", "", 150,0.5,1.5)
67 hTrue = RT.TH1D("hTrue", "", 150,0.5,1.5)
68 true_vs_reco = RT.TH2D("true_vs_reco", "", 150,0.5,1.5,150,0.5,1.5)
69 
70 r = RT.TH1D("r", "", 200,-1.,1.)
71 r_vs_true = RT.TH2D("r_vs_true", "", 150,0.5,1.5, 200,-1.,1.)
72 r_vs_p = RT.TH2D("r_vs_p", "", 150,0.5,1.5, 200,-1.,1.)
73 
74 hF1 = RT.TH1D("hF1", "", 192, 0, 192)
75 hF2 = RT.TH1D("hF2", "", 192, 0, 192)
76 hF3 = RT.TH1D("hF3", "", 192, 0, 192)
77 
78 f1 = array("i", [0])
79 f2 = array("i", [0])
80 f3 = array("i", [0])
81 
82 x1 = array("d", [0])
83 x2 = array("d", [0])
84 x3 = array("d", [0])
85 
86 true_p = array("d", [0.])
87 preSpec_p = array("d", [0.])
88 TOF1_p = array("d", [0.])
89 postSpec_p = array("d", [0.])
90 NP04front_p = array("d", [0.])
91 NP04FieldCage_p = array("d", [0.])
92 reco_p = array("d", [0.])
93 reco_p_B_plus = array("d", [0.])
94 reco_p_B_minus = array("d", [0.])
95 reco_p_B_plus_2 = array("d", [0.])
96 reco_p_B_minus_2 = array("d", [0.])
97 reco_tof = array("d", [0.])
98 flip_p = array("d", [0.])
99 plus_p_1 = array("d", [0.])
100 minus_p_1 = array("d", [0.])
101 plus_p_2 = array("d", [0.])
102 minus_p_2 = array("d", [0.])
103 plus_p_3 = array("d", [0.])
104 minus_p_3 = array("d", [0.])
105 PDG = array("i", [0])
106 fcPDG = array("i", [0])
107 
108 TOF1_TrackID = array("i", [0])
109 TRIG2_TrackID = array("i", [0])
110 TOF1_ParentID = array("i", [0])
111 TRIG2_ParentID = array("i", [0])
112 TOF1_EventID = array("i", [0])
113 TRIG2_EventID = array("i", [0])
114 
115 outtree = RT.TTree("tree","")
116 outtree.Branch("f1", f1, "f1/I")
117 outtree.Branch("f2", f2, "f2/I")
118 outtree.Branch("f3", f3, "f3/I")
119 outtree.Branch("x1", x1, "x1/D")
120 outtree.Branch("x2", x2, "x2/D")
121 outtree.Branch("x3", x3, "x3/D")
122 
123 outtree.Branch("true_p", true_p, "true_p/D")
124 outtree.Branch("preSpec_p", preSpec_p, "preSpec_p/D")
125 outtree.Branch("TOF1_p", TOF1_p, "TOF1_p/D")
126 outtree.Branch("postSpec_p", postSpec_p, "postSpec_p/D")
127 outtree.Branch("NP04front_p", NP04front_p, "NP04front_p/D")
128 outtree.Branch("NP04FieldCage_p", NP04FieldCage_p, "NP04FieldCage_p/D")
129 
130 outtree.Branch("TOF1_TrackID", TOF1_TrackID, "TOF1_TrackID/I")
131 outtree.Branch("TRIG2_TrackID", TRIG2_TrackID, "TRIG2_TrackID/I")
132 outtree.Branch("TOF1_ParentID", TOF1_ParentID, "TOF1_ParentID/I")
133 outtree.Branch("TRIG2_ParentID", TRIG2_ParentID, "TRIG2_ParentID/I")
134 outtree.Branch("TOF1_EventID", TOF1_EventID, "TOF1_EventID/I")
135 outtree.Branch("TRIG2_EventID", TRIG2_EventID, "TRIG2_EventID/I")
136 
137 outtree.Branch("reco_p", reco_p, "reco_p/D")
138 outtree.Branch("reco_p_B_plus", reco_p_B_plus, "reco_p_B_plus/D")
139 outtree.Branch("reco_p_B_minus", reco_p_B_minus, "reco_p_B_minus/D")
140 outtree.Branch("reco_p_B_plus_2", reco_p_B_plus_2, "reco_p_B_plus_2/D")
141 outtree.Branch("reco_p_B_minus_2", reco_p_B_minus_2, "reco_p_B_minus_2/D")
142 reco_p_shift_n1sigma = array("d", [0.])
143 reco_p_shift_p1sigma = array("d", [0.])
144 reco_p_shift_n2sigma = array("d", [0.])
145 reco_p_shift_p2sigma = array("d", [0.])
146 outtree.Branch("reco_p_shift_n1sigma", reco_p_shift_n1sigma, "reco_p_shift_n1sigma/D")
147 outtree.Branch("reco_p_shift_p1sigma", reco_p_shift_p1sigma, "reco_p_shift_p1sigma/D")
148 outtree.Branch("reco_p_shift_n2sigma", reco_p_shift_n2sigma, "reco_p_shift_n2sigma/D")
149 outtree.Branch("reco_p_shift_p2sigma", reco_p_shift_p2sigma, "reco_p_shift_p2sigma/D")
150 reco_p_B_p1_shift_p1 = array("d", [0.])
151 reco_p_B_m1_shift_p1 = array("d", [0.])
152 reco_p_B_p1_shift_m1 = array("d", [0.])
153 reco_p_B_m1_shift_m1 = array("d", [0.])
154 outtree.Branch("reco_p_B_p1_shift_p1", reco_p_B_p1_shift_p1, "reco_p_B_p1_shift_p1/D")
155 outtree.Branch("reco_p_B_m1_shift_p1", reco_p_B_m1_shift_p1, "reco_p_B_m1_shift_p1/D")
156 outtree.Branch("reco_p_B_p1_shift_m1", reco_p_B_p1_shift_m1, "reco_p_B_p1_shift_m1/D")
157 outtree.Branch("reco_p_B_m1_shift_m1", reco_p_B_m1_shift_m1, "reco_p_B_m1_shift_m1/D")
158 
159 reco_p_B_p2_shift_p1 = array("d", [0.])
160 reco_p_B_m2_shift_p1 = array("d", [0.])
161 reco_p_B_p2_shift_m1 = array("d", [0.])
162 reco_p_B_m2_shift_m1 = array("d", [0.])
163 outtree.Branch("reco_p_B_p2_shift_p1", reco_p_B_p2_shift_p1, "reco_p_B_p2_shift_p1/D")
164 outtree.Branch("reco_p_B_m2_shift_p1", reco_p_B_m2_shift_p1, "reco_p_B_m2_shift_p1/D")
165 outtree.Branch("reco_p_B_p2_shift_m1", reco_p_B_p2_shift_m1, "reco_p_B_p2_shift_m1/D")
166 outtree.Branch("reco_p_B_m2_shift_m1", reco_p_B_m2_shift_m1, "reco_p_B_m2_shift_m1/D")
167 
168 reco_p_B_p1_shift_p2 = array("d", [0.])
169 reco_p_B_m1_shift_p2 = array("d", [0.])
170 reco_p_B_p1_shift_m2 = array("d", [0.])
171 reco_p_B_m1_shift_m2 = array("d", [0.])
172 outtree.Branch("reco_p_B_p1_shift_p2", reco_p_B_p1_shift_p2, "reco_p_B_p1_shift_p2/D")
173 outtree.Branch("reco_p_B_m1_shift_p2", reco_p_B_m1_shift_p2, "reco_p_B_m1_shift_p2/D")
174 outtree.Branch("reco_p_B_p1_shift_m2", reco_p_B_p1_shift_m2, "reco_p_B_p1_shift_m2/D")
175 outtree.Branch("reco_p_B_m1_shift_m2", reco_p_B_m1_shift_m2, "reco_p_B_m1_shift_m2/D")
176 
177 reco_p_B_p2_shift_p2 = array("d", [0.])
178 reco_p_B_m2_shift_p2 = array("d", [0.])
179 reco_p_B_p2_shift_m2 = array("d", [0.])
180 reco_p_B_m2_shift_m2 = array("d", [0.])
181 outtree.Branch("reco_p_B_p2_shift_p2", reco_p_B_p2_shift_p2, "reco_p_B_p2_shift_p2/D")
182 outtree.Branch("reco_p_B_m2_shift_p2", reco_p_B_m2_shift_p2, "reco_p_B_m2_shift_p2/D")
183 outtree.Branch("reco_p_B_p2_shift_m2", reco_p_B_p2_shift_m2, "reco_p_B_p2_shift_m2/D")
184 outtree.Branch("reco_p_B_m2_shift_m2", reco_p_B_m2_shift_m2, "reco_p_B_m2_shift_m2/D")
185 
186 outtree.Branch("reco_tof", reco_tof, "reco_tof/D")
187 outtree.Branch("flip_p", flip_p, "flip_p/D")
188 outtree.Branch("plus_p_1", plus_p_1, "plus_p_1/D")
189 outtree.Branch("minus_p_1", minus_p_1, "minus_p_1/D")
190 outtree.Branch("plus_p_2", plus_p_2, "plus_p_2/D")
191 outtree.Branch("minus_p_2", minus_p_2, "minus_p_2/D")
192 outtree.Branch("plus_p_3", plus_p_3, "plus_p_3/D")
193 outtree.Branch("minus_p_3", minus_p_3, "minus_p_3/D")
194 
195 outtree.Branch("PDG", PDG, "PDG/I")
196 outtree.Branch("fcPDG", fcPDG, "fcPDG/I")
197 
198 def momentum(x1,x2,x3, scale = 1.):
199  return 299792458*LB*scale/(1.E9 * acos(momentum_costheta(x1,x2,x3)))
200 
201 
202 count = 0
203 print("Tree has", t.GetEntries(), "entries")
204 for e in t:
205 
206  PDG[0] = int(e.NP04front_PDGid)
207  fcPDG[0] = int(e.NP04FieldCage_PDGid)
208 
209  TOF1_TrackID[0] = int(e.TOF1_TrackID)
210  TRIG2_TrackID[0] = int(e.TRIG2_TrackID)
211  TOF1_ParentID[0] = int(e.TOF1_ParentID)
212  TRIG2_ParentID[0] = int(e.TRIG2_ParentID)
213  TOF1_EventID[0] = int(e.TOF1_EventID)
214  TRIG2_EventID[0] = int(e.TRIG2_EventID)
215 
216  if not (count % 1000): print(count)
217  count = count + 1
218  px = e.AfterTarget_Px
219  py = e.AfterTarget_Py
220  pz = e.AfterTarget_Pz
221 
222  true_p[0] = sqrt( px**2 + py**2 + pz**2 )/1.e3
223  #if true_p[0] < 1.e-5: continue
224 
225  x1[0] = -1.e-3*shift_x(e.BPROF1_x)
226  x2[0] = -1.e-3*shift_x(e.BPROF2_x)
227  x3[0] = -1.e-3*shift_x(e.BPROF3_x)
228 
229  f1[0] = int( 96 - e.BPROF1_x)
230  f2[0] = int( 96 - e.BPROF2_x)
231  f3[0] = int( 96 - e.BPROF3_x)
232 
233  preSpec_p[0] = sqrt( e.BPROF1_Px**2 + e.BPROF1_Py**2 + e.BPROF1_Pz**2 )
234  TOF1_p[0] = sqrt( e.TOF1_Px**2 + e.TOF1_Py**2 + e.TOF1_Pz**2 )
235  postSpec_p[0] = sqrt( e.BPROF3_Px**2 + e.BPROF3_Py**2 + e.BPROF3_Pz**2 )
236  NP04front_p[0] = sqrt( e.NP04front_Px**2 + e.NP04front_Py**2 + e.NP04front_Pz**2 )
237  NP04FieldCage_p[0] = sqrt( e.NP04FieldCage_Px**2 + e.NP04FieldCage_Py**2 + e.NP04FieldCage_Pz**2 )
238 
239  if NP04front_p[0] < 1.e-5: continue
240 
241  hTrue.Fill(true_p[0])
242  hF1.Fill( f1[0] )
243  hF2.Fill( f2[0] )
244  hF3.Fill( f3[0] )
245 
246  reco_tof[0] = (e.TRIG2_t - e.TOF1_t)
247 
248  #print momentum_costheta(x1,x2,x3), acos(momentum_costheta(x1,x2,x3))
249 
250  #1GeV: 68.8
251  #2GeV: 137.5
252  #3GeV: 206.2
253 
254  if( args.p == "1" ): current = 68.8
255  elif( args.p == "0.5" ): current = 34.4
256  elif( args.p == "2" ): current = 137.5
257  elif( args.p == "3" ): current = 206.2
258  elif( args.p == "6" ): current = 419.7
259  elif( args.p == "7" ): current = 508.3
260 
261  LB = mag_P1 * current
262  deltaI = current - mag_P4
263  if deltaI > 0: LB = LB + mag_P3 * deltaI**2
264  #print 299792458*LB/(1.E9 * acos(momentum_costheta(x1,x2,x3)));
265  #reco_p[0] = 299792458*LB/(1.E9 * acos(momentum_costheta(x1,x2,x3)))
266  reco_p[0] = momentum(x1[0], x2[0], x3[0] )
267  reco_p_B_plus[0] = momentum(x1[0], x2[0], x3[0], 1.01 )
268  reco_p_B_minus[0] = momentum(x1[0], x2[0], x3[0], .99 )
269  reco_p_B_plus_2[0] = momentum(x1[0], x2[0], x3[0], 1.02 )
270  reco_p_B_minus_2[0] = momentum(x1[0], x2[0], x3[0], .98 )
271  reco_p_shift_n1sigma[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*1.e-3 )
272  reco_p_shift_p1sigma[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*1.e-3 )
273  reco_p_shift_n2sigma[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*2.e-3 )
274  reco_p_shift_p2sigma[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*2.e-3 )
275 
276  reco_p_B_p1_shift_p1[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*1.e-3 , 1.01)
277  reco_p_B_m1_shift_p1[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*1.e-3 , 0.99)
278  reco_p_B_p1_shift_m1[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*1.e-3 , 1.01)
279  reco_p_B_m1_shift_m1[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*1.e-3 , 0.99)
280 
281  reco_p_B_p2_shift_p1[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*1.e-3 , 1.02)
282  reco_p_B_m2_shift_p1[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*1.e-3 , 0.98)
283  reco_p_B_p2_shift_m1[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*1.e-3 , 1.02)
284  reco_p_B_m2_shift_m1[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*1.e-3 , 0.98)
285 
286  reco_p_B_p1_shift_p2[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*2.e-3 , 1.01)
287  reco_p_B_m1_shift_p2[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*2.e-3 , 0.99)
288  reco_p_B_p1_shift_m2[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*2.e-3 , 1.01)
289  reco_p_B_m1_shift_m2[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*2.e-3 , 0.99)
290 
291  reco_p_B_p2_shift_p2[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*2.e-3 , 1.02)
292  reco_p_B_m2_shift_p2[0] = momentum(x1[0], x2[0], x3[0] - (args.s)*2.e-3 , 0.98)
293  reco_p_B_p2_shift_m2[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*2.e-3 , 1.02)
294  reco_p_B_m2_shift_m2[0] = momentum(x1[0], x2[0], x3[0] + (args.s)*2.e-3 , 0.98)
295 
296  flip_p[0] = momentum(-1.*x1[0], -1.*x2[0], -1.*x3[0] )
297 
298  h.Fill(reco_p[0])
299  #r.Fill(1. - reco_p[0] / true_p[0] )
300  #r_vs_true.Fill(true_p[0], 1. - reco_p[0] / true_p[0] )
301  #r_vs_p.Fill(reco_p[0], 1. - reco_p[0] / true_p[0] )
302  true_vs_reco.Fill( reco_p[0], true_p[0] )
303 
304 
305  plus_p_1[0] = momentum(x1[0]-.5e-3, x2[0] ,x3[0])
306  minus_p_1[0] = momentum(x1[0]+.5e-3, x2[0] ,x3[0])
307 
308  plus_p_2[0] = momentum(x1[0], x2[0]-.5e-3 ,x3[0])
309  minus_p_2[0] = momentum(x1[0], x2[0]+.5e-3 ,x3[0])
310 
311  plus_p_3[0] = momentum(x1[0], x2[0], x3[0]-.5e-3)
312  minus_p_3[0] = momentum(x1[0], x2[0], x3[0]+.5e-3)
313 
314  outtree.Fill()
315 
316 fout.cd()
317 h.Write()
318 hTrue.Write()
319 r.Write()
320 r_vs_true.Write()
321 r_vs_p.Write()
322 
323 hF1.Write()
324 hF2.Write()
325 hF3.Write()
326 
327 true_vs_reco.Write()
328 
329 outtree.Write()
330 fout.Close()
reco_momentum_tuples.shift_x
def shift_x(x)
Definition: reco_momentum_tuples.py:39
reco_momentum_tuples.int
int
Definition: reco_momentum_tuples.py:10
gar::dump::array
auto array(Array const &a)
Returns a manipulator which will print the specified array.
Definition: DumpUtils.h:228
reco_momentum_tuples.momentum_costheta
def momentum_costheta(x1, x2, x3)
Definition: reco_momentum_tuples.py:46
reco_momentum_tuples.momentum
def momentum(x1, x2, x3, scale=1.)
Definition: reco_momentum_tuples.py:198
if
if(!yymsg) yymsg
reco_momentum_tuples.to_f
def to_f(x)
Definition: reco_momentum_tuples.py:42
Generated by   doxygen 1.8.11