← Back to branch summary

~maus-step4-study/maus/step4_reducedCurrent_study

« back to all changes in this revision

Viewing changes to bin/user/emittance_from_trackers.py

  • Committer: Christopher Hunt
  • Date: 2014-07-30 14:09:43 UTC
  • Revision ID: christopher.hunt08@imperial.ac.uk-20140730140943-gad3vxtlziw2dwx7
Cleaning up some crap

Show diffs side-by-side

added added

removed removed

Lines of Context:
bin/user/emittance_from_trackers.py
1
 
#!/usr/bin/env python
2
 
 
3
 
# This file is part of MAUS: http://micewww.pp.rl.ac.uk/projects/maus
4
 
#
5
 
# MAUS is free software: you can redistribute it and/or modify
6
 
# it under the terms of the GNU General Public License as published by
7
 
# the Free Software Foundation, either version 3 of the License, or
8
 
# (at your option) any later version.
9
 
#
10
 
# MAUS is distributed in the hope that it will be useful,
11
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
12
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13
 
# GNU General Public License for more details.
14
 
#
15
 
# You should have received a copy of the GNU General Public License
16
 
# along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
17
 
#
18
 
 
19
 
 
20
 
import io   #  generic python library for I/O
21
 
import MAUS
22
 
import ROOT
23
 
import xboa
24
 
 
25
 
# definitions of MAUS data structure for PyROOT
26
 
import libMausCpp #pylint: disable = W0611
27
 
 
28
 
import math
29
 
import sys
30
 
 
31
 
 
32
 
 
33
 
p_value_cut = 0.05
34
 
pt_cut = 150.0
35
 
pz_cut = 300.0
36
 
aperture_cut = 189.0
37
 
pid_cut = -13
38
 
 
39
 
upstream_reference = 9001
40
 
downstream_reference = 9002
41
 
 
42
 
print_everything = False
43
 
print_data = False
44
 
 
45
 
 
46
 
 
47
 
number_recon_cut = 0
48
 
number_reference_cut = 0
49
 
number_reference_e = 0
50
 
 
51
 
reference_hit_count = 0
52
 
 
53
 
 
54
 
 
55
 
def run() :
56
 
  if len( sys.argv ) <= 1 :
57
 
    print "No File Supplied."
58
 
    sys.exit( 0 )
59
 
 
60
 
  filename = sys.argv[1]
61
 
 
62
 
  print "Loading File"
63
 
  print filename
64
 
  print
65
 
 
66
 
  infile = ROOT.TFile( filename, "READ" )
67
 
 
68
 
  data = ROOT.MAUS.Data()
69
 
  tree = infile.Get("Spill")
70
 
  tree.SetBranchAddress("data", data)
71
 
 
72
 
  kalman_up_bunches = [ xboa.Bunch.Bunch() for i in range( 5 ) ]
73
 
  kalman_down_bunches = [ xboa.Bunch.Bunch() for i in range( 5 ) ]
74
 
  reference1_bunch = xboa.Bunch.Bunch()
75
 
  reference2_bunch = xboa.Bunch.Bunch()
76
 
 
77
 
  print "Loading Data"
78
 
  print "Please Wait..."
79
 
 
80
 
 
81
 
  global number_recon_cut
82
 
  global number_reference_cut
83
 
  global number_reference_e
84
 
 
85
 
  global upstream_reference
86
 
  global downstream_reference
87
 
 
88
 
  global print_data
89
 
 
90
 
  global reference_hit_count
91
 
 
92
 
  
93
 
  for i in range( tree.GetEntries() ) :
94
 
    tree.GetEntry( i )
95
 
    spill = data.GetSpill()
96
 
 
97
 
    if spill.GetDaqEventType() != "physics_event" :
98
 
      continue
99
 
 
100
 
##  KALMAN TRACKS  ##
101
 
    for j in range( spill.GetReconEvents().size() ) :
102
 
      for k in range( spill.GetReconEvents()[j].GetSciFiEvent().scifitracks().size() ) :
103
 
        track = spill.GetReconEvents()[j].GetSciFiEvent().scifitracks()[k]
104
 
 
105
 
        if track.tracker() == 0 :
106
 
          appendHits( kalman_up_bunches, track )
107
 
          appendMC( reference1_bunch, track )
108
 
        elif track.tracker() == 1 :
109
 
          appendHits( kalman_down_bunches, track )
110
 
          appendMC( reference2_bunch, track )
111
 
        else :
112
 
          print "Unknown Tracker,", track.tracker()
113
 
 
114
 
 
115
 
 
116
 
  number_kalman = 0
117
 
  for i in range( 5 ) :
118
 
    number_kalman += len( kalman_up_bunches[i] )
119
 
    number_kalman += len( kalman_down_bunches[i] )
120
 
 
121
 
  print
122
 
  print "Found ", number_kalman, " Kalman Points"
123
 
  print "Found ", len( reference1_bunch ) + len( reference2_bunch ), " Kalman MC Hits"
124
 
  print
125
 
  print "Cut", number_recon_cut, "Particles from the Reconstruction."
126
 
#  print "Cut", number_reference_cut, "Particles from the Reference Planes."
127
 
#  print "Removed", number_reference_e, "Electrons from Reference Planes"
128
 
  print
129
 
  print "Calculating Emittances"
130
 
  print
131
 
  print reference_hit_count
132
 
  print
133
 
 
134
 
 
135
 
  reference1_means = reference1_bunch.mean( ['x', 'y', 'z', 'pz', 'p', 'energy', 'bz' ] )
136
 
  reference1_emittance = reference1_bunch.get_emittance( ['x', 'y'] )
137
 
  reference1_weight = reference1_bunch.bunch_weight()
138
 
  reference1_beta = reference1_bunch.get_beta( ['x', 'y'] )
139
 
 
140
 
  reference2_means = reference2_bunch.mean( ['x', 'y', 'z', 'pz', 'p', 'energy', 'bz' ] )
141
 
  reference2_emittance = reference2_bunch.get_emittance( ['x', 'y'] )
142
 
  reference2_weight = reference2_bunch.bunch_weight()
143
 
  reference2_beta = reference2_bunch.get_beta( ['x', 'y'] )
144
 
 
145
 
 
146
 
  print
147
 
  print "=========================================="
148
 
  print " ---    USING KALMAN MC POINTS    --- "
149
 
  print "=========================================="
150
 
  print
151
 
 
152
 
  print
153
 
  print " ---  Tracker 1  --- "
154
 
  print 
155
 
  print "Weight       =  " + str( reference1_weight )
156
 
  print "Beta         =  " + str( reference1_beta )
157
 
  print "Mean x       =  " + str( reference1_means[ 'x' ] )
158
 
  print "Mean y       =  " + str( reference1_means[ 'y' ] )
159
 
  print "Mean z       =  " + str( reference1_means[ 'z' ] )
160
 
  print "Mean pz      =  " + str( reference1_means[ 'pz' ] )
161
 
  print "Mean p       =  " + str( reference1_means[ 'p' ] )
162
 
  print "Mean Energy  =  " + str( reference1_means[ 'energy' ] )
163
 
  print "Mean Bz      =  " + str( reference1_means[ 'bz' ] )
164
 
  print "Emittance    =  " + str( reference1_emittance )
165
 
  print
166
 
  print
167
 
 
168
 
  print " ---  Tracker 2  --- "
169
 
  print 
170
 
  print "Weight       =  " + str( reference2_weight )
171
 
  print "Beta         =  " + str( reference2_beta )
172
 
  print "Mean x       =  " + str( reference2_means[ 'x' ] )
173
 
  print "Mean y       =  " + str( reference2_means[ 'y' ] )
174
 
  print "Mean z       =  " + str( reference2_means[ 'z' ] )
175
 
  print "Mean pz      =  " + str( reference2_means[ 'pz' ] )
176
 
  print "Mean p       =  " + str( reference2_means[ 'p' ] )
177
 
  print "Mean Energy  =  " + str( reference2_means[ 'energy' ] )
178
 
  print "Mean Bz      =  " + str( reference2_means[ 'bz' ] )
179
 
  print "Emittance    =  " + str( reference2_emittance )
180
 
  print
181
 
 
182
 
 
183
 
  kalman1_means = kalman_up_bunches[0].mean( ['x', 'y', 'z', 'pz', 'p', 'energy', 'bz' ] )
184
 
  kalman1_emittance = kalman_up_bunches[0].get_emittance( ['x', 'y'] )
185
 
  kalman1_weight = kalman_up_bunches[0].bunch_weight()
186
 
  kalman1_beta = kalman_up_bunches[0].get_beta( ['x', 'y'] )
187
 
 
188
 
  kalman2_means = kalman_down_bunches[0].mean( ['x', 'y', 'z', 'pz', 'p', 'energy', 'bz' ] )
189
 
  kalman2_emittance = kalman_down_bunches[0].get_emittance( ['x', 'y'] )
190
 
  kalman2_weight = kalman_down_bunches[0].bunch_weight()
191
 
  kalman2_beta = kalman_down_bunches[0].get_beta( ['x', 'y'] )
192
 
 
193
 
  print
194
 
  print "=========================================="
195
 
  print " --- USING KALMAN FITTED TRACK POINTS --- "
196
 
  print "=========================================="
197
 
  print
198
 
 
199
 
  print
200
 
  print " ---  Tracker 1  --- "
201
 
  print 
202
 
  print "Weight       =  " + str( kalman1_weight )
203
 
  print "Beta         =  " + str( kalman1_beta )
204
 
  print "Mean x       =  " + str( kalman1_means[ 'x' ] )
205
 
  print "Mean y       =  " + str( kalman1_means[ 'y' ] )
206
 
  print "Mean z       =  " + str( kalman1_means[ 'z' ] )
207
 
  print "Mean pz      =  " + str( kalman1_means[ 'pz' ] )
208
 
  print "Mean p       =  " + str( kalman1_means[ 'p' ] )
209
 
  print "Mean Energy  =  " + str( kalman1_means[ 'energy' ] )
210
 
  print "Mean Bz      =  " + str( kalman1_means[ 'bz' ] )
211
 
  print "Emittance    =  " + str( kalman1_emittance )
212
 
  print
213
 
  print
214
 
 
215
 
  print " ---  Tracker 2  --- "
216
 
  print 
217
 
  print "Weight       =  " + str( kalman2_weight )
218
 
  print "Beta         =  " + str( kalman2_beta )
219
 
  print "Mean x       =  " + str( kalman2_means[ 'x' ] )
220
 
  print "Mean y       =  " + str( kalman2_means[ 'y' ] )
221
 
  print "Mean z       =  " + str( kalman2_means[ 'z' ] )
222
 
  print "Mean pz      =  " + str( kalman2_means[ 'pz' ] )
223
 
  print "Mean p       =  " + str( kalman2_means[ 'p' ] )
224
 
  print "Mean Energy  =  " + str( kalman2_means[ 'energy' ] )
225
 
  print "Mean Bz      =  " + str( kalman2_means[ 'bz' ] )
226
 
  print "Emittance    =  " + str( kalman2_emittance )
227
 
  print
228
 
 
229
 
 
230
 
 
231
 
  if print_everything :
232
 
    printPlots( reference1_bunch, "Ref1_" )
233
 
    printPlots( reference2_bunch, "Ref2_" )
234
 
    printPlots( kalman_up_bunches[0], "Kal_up_1_" )
235
 
    printPlots( kalman_up_bunches[1], "Kal_up_2_" )
236
 
    printPlots( kalman_up_bunches[2], "Kal_up_3_" )
237
 
    printPlots( kalman_up_bunches[3], "Kal_up_4_" )
238
 
    printPlots( kalman_up_bunches[4], "Kal_up_5_" )
239
 
    printPlots( kalman_down_bunches[0], "Kal_down_1_" )
240
 
    printPlots( kalman_down_bunches[1], "Kal_down_2_" )
241
 
    printPlots( kalman_down_bunches[2], "Kal_down_3_" )
242
 
    printPlots( kalman_down_bunches[3], "Kal_down_4_" )
243
 
    printPlots( kalman_down_bunches[4], "Kal_down_5_" )
244
 
 
245
 
 
246
 
  if print_data is True :
247
 
    print
248
 
    print "Saving Station Data"
249
 
    print "Please Wait."
250
 
    print
251
 
    with open( 'station_data.txt', 'w' ) as outfile :
252
 
      outfile.write( '# StationNum    Mean X    Mean Y   Mean Px   Mean Py   Mean Pz   Mean Pt   Mean Bz  Energy   Beta    Emittance   Weight \n' )
253
 
      bunch_count = 0
254
 
      for i in range( 5 ) :
255
 
        bun_up = kalman_up_bunches[i]
256
 
        means_up = bun_up.mean( ['x','y','px','py','pz','pt','energy','bz'] )
257
 
 
258
 
        outfile.write( "UPSTREAM " + str( i ) + ' ' +
259
 
                       str( means_up[ 'x' ] ) + ' ' +
260
 
                       str( means_up[ 'y' ] ) + ' ' +
261
 
                       str( means_up[ 'px' ] ) + ' ' +
262
 
                       str( means_up[ 'py' ] ) + ' ' +
263
 
                       str( means_up[ 'pz' ] ) + ' ' +
264
 
                       str( means_up[ 'pt' ] ) + ' ' +
265
 
                       str( means_up[ 'bz' ] ) + ' ' +
266
 
                       str( means_up[ 'energy' ] ) + ' ' +
267
 
                       str( bun_up.get_emittance( [ 'x','y' ] ) ) + ' ' +
268
 
                       str( bun_up.get_beta( ['x','y'] ) ) + ' ' +
269
 
                       str( bun_up.bunch_weight() ) + "\n" )
270
 
 
271
 
      outfile.write( '\n' )
272
 
      for i in range( 5 ) :
273
 
        bun_down = kalman_down_bunches[i]
274
 
        means_down = bun_down.mean( ['x','y','px','py','pz','pt','energy','bz'] )
275
 
 
276
 
        outfile.write( "DOWNSTREAM " + str( i ) + ' ' +
277
 
                       str( means_down[ 'x' ] ) + ' ' +
278
 
                       str( means_down[ 'y' ] ) + ' ' +
279
 
                       str( means_down[ 'px' ] ) + ' ' +
280
 
                       str( means_down[ 'py' ] ) + ' ' +
281
 
                       str( means_down[ 'pz' ] ) + ' ' +
282
 
                       str( means_down[ 'pt' ] ) + ' ' +
283
 
                       str( means_down[ 'bz' ] ) + ' ' +
284
 
                       str( means_down[ 'energy' ] ) + ' ' +
285
 
                       str( bun_down.get_emittance( [ 'x','y' ] ) ) + ' ' +
286
 
                       str( bun_down.get_beta( ['x','y'] ) ) + ' ' +
287
 
                       str( bun_down.bunch_weight() ) + "\n" )
288
 
 
289
 
    print 'station_data.txt Created'
290
 
 
291
 
 
292
 
 
293
 
def appendHits( bunch_list, track ) :
294
 
  global number_recon_cut
295
 
  if track.P_value() < p_value_cut :
296
 
    number_recon_cut += 1
297
 
    return
298
 
 
299
 
  trackPoints = track.scifitrackpoints()
300
 
 
301
 
  for point in trackPoints :
302
 
    if point.plane() != 0 : 
303
 
      continue
304
 
 
305
 
    x = point.x()
306
 
    y = point.y()
307
 
    z = 0.0
308
 
 
309
 
    px = point.px()
310
 
    py = point.py()
311
 
    pz = point.pz()
312
 
 
313
 
    if math.sqrt(px**2 + py**2) > pt_cut or pz > pz_cut : 
314
 
      print "KAL - OOPS :", px, py, pz, track.P_value()
315
 
      number_recon_cut += 1
316
 
      continue
317
 
 
318
 
    pid = 13
319
 
 
320
 
    bz = 4.0
321
 
 
322
 
    hit = xboa.Hit.Hit.new_from_dict( { 'x':x, 'y':y, 'z':z, 'px':px, 'py':py, 'pz':pz, 'pid':pid, 'bz':bz, 'mass':xboa.Common.pdg_pid_to_mass[pid] }, 'energy' )
323
 
    bunch_list[ point.station() ].append( hit )
324
 
 
325
 
 
326
 
 
327
 
def appendMC( bunch, track ) :
328
 
  global reference_hit_count
329
 
 
330
 
  trackPoints = track.scifitrackpoints()
331
 
 
332
 
  for point in trackPoints :
333
 
    if point.station() != 0 or point.plane() != 0 : 
334
 
      continue
335
 
 
336
 
    x = point.mc_x()
337
 
    y = point.mc_y()
338
 
    z = 0.0
339
 
 
340
 
    px = point.mc_px()
341
 
    py = point.mc_py()
342
 
    pz = point.mc_pz()
343
 
 
344
 
    pid = 13
345
 
 
346
 
    bz = 4.0
347
 
 
348
 
    hit = xboa.Hit.Hit.new_from_dict( { 'x':x, 'y':y, 'z':z, 'px':px, 'py':py, 'pz':pz, 'pid':pid, 'bz':bz, 'mass':xboa.Common.pdg_pid_to_mass[pid] }, 'energy' )
349
 
    bunch.append( hit )
350
 
    reference_hit_count += 1
351
 
 
352
 
 
353
 
def appendReferenceHit( bunch, spechit ) :
354
 
 
355
 
  x = spechit.GetPosition()[0]
356
 
  y = spechit.GetPosition()[1]
357
 
  z = spechit.GetPosition()[2]
358
 
  px = spechit.GetMomentum()[0]
359
 
  py = spechit.GetMomentum()[1]
360
 
  pz = spechit.GetMomentum()[2]
361
 
  pid = abs( spechit.GetParticleId() )
362
 
 
363
 
  bz = 4.0
364
 
 
365
 
  hit = xboa.Hit.Hit.new_from_dict( { 'x':x, 'y':y, 'z':z, 'px':px, 'py':py, 'pz':pz, 'pid':pid, 'bz':bz, 'mass':xboa.Common.pdg_pid_to_mass[pid] }, 'energy' )
366
 
  bunch.append( hit )
367
 
 
368
 
 
369
 
#def check_hit( hit ) :
370
 
#  if hit.GetParticleId() != pid_cut : 
371
 
#    return False
372
 
#
373
 
#  x = hit.GetPosition()[0]
374
 
#  y = hit.GetPosition()[1]
375
 
#  z = hit.GetPosition()[2]
376
 
#
377
 
#  if ( math.sqrt( x*x + y*y ) > aperture_cut ) :
378
 
#    return False
379
 
#
380
 
#  px = hit.GetMomentum()[0]
381
 
#  py = hit.GetMomentum()[1]
382
 
#  pz = hit.GetMomentum()[2]
383
 
#  PT = math.sqrt(px**2 + py**2)
384
 
#
385
 
#  if  PT > pt_cut : 
386
 
#    return False
387
 
#  if pz > pz_cut or pz < 0.0 :
388
 
#    return False
389
 
#
390
 
#  return True
391
 
 
392
 
 
393
 
 
394
 
def printPlots( bunch, prefix ) :
395
 
    canvas, hist = bunch.root_histogram( 'x', 'm', 'px', 'MeV/c' )
396
 
    canvas.Print( prefix+"EmittanceX.eps", "eps")
397
 
    canvas.Close()
398
 
    canvas, hist = bunch.root_histogram( 'y', 'm', 'py', 'MeV/c' )
399
 
    canvas.Print( prefix+"EmittanceY.eps", "eps")
400
 
    canvas.Close()
401
 
    canvas, hist = bunch.root_histogram( 'pt', 'MeV/c' )
402
 
    canvas.Print( prefix+"PT.eps", "eps")
403
 
    canvas.Close()
404
 
    canvas, hist = bunch.root_histogram( 'pz', 'MeV/c' )
405
 
    canvas.Print( prefix+"PZ.eps", "eps")
406
 
    canvas.Close()
407
 
    canvas, hist = bunch.root_histogram( 'x', 'mm', 'y', 'MeV/c' )
408
 
    canvas.Print( prefix+"X-Y.eps", "eps")
409
 
    canvas.Close()
410
 
    canvas, hist = bunch.root_histogram( 'px', 'mm', 'py', 'MeV/c' )
411
 
    canvas.Print( prefix+"PX-PY.eps", "eps")
412
 
    canvas.Close()
413
 
 
414
 
    filename = prefix+"hitList.txt"
415
 
 
416
 
    with open( filename, 'w' ) as outfile :
417
 
      outfile.write( '#  Position (x,y,z)  |  Momentum (px,py,pz)  |  Magnetic Field (bz) \n\n' )
418
 
 
419
 
      for hit in bunch :
420
 
        outfile.write( str(hit.get('x')) + ' '
421
 
                       + str(hit.get( 'y' )) + ' '
422
 
                       + str(hit.get( 'z' )) + ' '
423
 
                       + str(hit.get( 'px' )) + ' '
424
 
                       + str(hit.get( 'py' )) + ' '
425
 
                       + str(hit.get( 'pz' )) + ' '
426
 
                       + str(hit.get( 'bz' )) + '\n' )
427
 
 
428
 
 
429
 
if __name__ == "__main__":
430
 
  run()
431
 
 
432