← Back to branch summary

~christopher-hunt08/maus/maus_DOE_review_version

« back to all changes in this revision

Viewing changes to bin/user/emittance_from_virtualPlanes_OLD.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_virtualPlanes_OLD.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 xboa.Bunch
21
 
from xboa.Bunch import Bunch
22
 
 
23
 
import ROOT
24
 
import math
25
 
 
26
 
import operator
27
 
import sys
28
 
 
29
 
 
30
 
 
31
 
 
32
 
aperture_cut = 189.
33
 
number_bunches = 69
34
 
pid_cut = -13
35
 
 
36
 
number_virtual_cut = 0
37
 
number_virtual_e = 0
38
 
 
39
 
inspection_plane = 56
40
 
upstream_reference = 14
41
 
downstream_reference = 50
42
 
 
43
 
 
44
 
def main():
45
 
    global aperture_cut
46
 
    global number_bunches
47
 
    global number_virtual_e
48
 
    global number_virtual_cut
49
 
 
50
 
    if len( sys.argv ) <= 1 :
51
 
      print "No File Supplied."
52
 
      sys.exit( 0 )
53
 
 
54
 
    filename = sys.argv[1]
55
 
 
56
 
    print "Loading File"
57
 
    print filename
58
 
    print
59
 
 
60
 
    infile = ROOT.TFile( filename, "READ" )
61
 
 
62
 
    data = ROOT.MAUS.Data()
63
 
    tree = infile.Get("Spill")
64
 
    tree.SetBranchAddress("data", data)
65
 
 
66
 
    bunch_list = [ xboa.Bunch.Bunch() for i in range( number_bunches ) ]
67
 
    ana_list = []
68
 
 
69
 
    for i in range( tree.GetEntries() ) :
70
 
      tree.GetEntry( i )
71
 
      spill = data.GetSpill()
72
 
 
73
 
      if spill.GetDaqEventType() != "physics_event" :
74
 
        continue
75
 
 
76
 
      for j in range( spill.GetMCEventSize() ) :
77
 
        flag = False
78
 
        event = spill.GetAnMCEvent( j )
79
 
 
80
 
        for k in range( event.GetVirtualHitsSize() ) :
81
 
          hit = event.GetAVirtualHit( k )
82
 
 
83
 
          x = hit.GetPosition()[0]
84
 
          y = hit.GetPosition()[1]
85
 
          z = hit.GetPosition()[2]
86
 
          px = hit.GetMomentum()[0]
87
 
          py = hit.GetMomentum()[1]
88
 
          pz = hit.GetMomentum()[2]
89
 
          pid = hit.GetParticleId()
90
 
          PT = math.sqrt(px**2 + py**2)
91
 
 
92
 
# Have to cut electrons later on - so that we contain decaying muons
93
 
          if pid != pid_cut :
94
 
            continue
95
 
 
96
 
          if ( math.sqrt( x*x + y*y ) >= aperture_cut ) :
97
 
            number_virtual_cut += 1
98
 
            flag = True
99
 
            break
100
 
#          if pid != -13 : 
101
 
#            number_virtual_cut += 1
102
 
#            flag = True
103
 
#            break
104
 
#          if PT > pz :
105
 
#            number_virtual_cut += 1
106
 
#            flag = True
107
 
#            break
108
 
          if pz <= 0.0 : 
109
 
            number_virtual_cut += 1
110
 
            flag = True
111
 
            break
112
 
#          if PT > 300.0 :
113
 
#            number_virtual_cut += 1
114
 
#            flag = True
115
 
#            break
116
 
#          if math.sqrt( px*px + py*py ) > pz :
117
 
#            print "VIRTUAL - WTF", px, py, pz
118
 
#            number_virtual_cut += 1
119
 
#            flag = True
120
 
#            break
121
 
 
122
 
        if flag is True :
123
 
          continue
124
 
 
125
 
        for k in range( event.GetVirtualHitsSize() ) :
126
 
          hit = event.GetAVirtualHit( k )
127
 
 
128
 
          bunch = bunch_list[ hit.GetStationId() - 1 ]
129
 
 
130
 
          x = hit.GetPosition()[0]
131
 
          y = hit.GetPosition()[1]
132
 
          z = hit.GetPosition()[2]
133
 
          px = hit.GetMomentum()[0]
134
 
          py = hit.GetMomentum()[1]
135
 
          pz = hit.GetMomentum()[2]
136
 
          pid = hit.GetParticleId()
137
 
          bx = hit.GetBField()[0]
138
 
          by = hit.GetBField()[1]
139
 
          bz = hit.GetBField()[2]
140
 
 
141
 
          if pid != pid_cut :
142
 
            continue
143
 
 
144
 
          hit = xboa.Hit.Hit.new_from_dict( { 'x':x, 'y':y, 'z':z, 'px':px, 'py':py, 'pz':pz, 'pid':pid, 'bx':bx, 'by':by, 'bz':bz, 'mass':xboa.Common.pdg_pid_to_mass[ abs( pid ) ] }, 'energy' )
145
 
          bunch.append( hit )
146
 
 
147
 
 
148
 
    print "Read", len(bunch_list), "Bunches"
149
 
 
150
 
    for bun in bunch_list :
151
 
      if bun.bunch_weight() > 1.0e-9 and len( bun ) > 1 :
152
 
        ana_list.append( bun )
153
 
 
154
 
    print "Found", len(ana_list), "Non Empty Bunches"
155
 
    print "Cut", number_virtual_cut, "Hits and ", number_virtual_e, "Electrons"
156
 
 
157
 
    outfile = ROOT.TFile( "emittance_from_virtualPlanes.root", "RECREATE" )
158
 
 
159
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'bunch_weight', [''], 'm')
160
 
    canvas.Write()
161
 
    canvas.Print( "BunchWeight.eps", "eps")
162
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'm', 'm')
163
 
    canvas.Write()
164
 
    canvas.Print( "Emittance.eps", "eps")
165
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'beta', ['x', 'y'], 'm', 'm')
166
 
    canvas.Write()
167
 
    canvas.Print( "BetaFunc.eps", "eps")
168
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['energy'], 'm', 'MeV')
169
 
    canvas.Write()
170
 
    canvas.Print( "Energy.eps", "eps")
171
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['bz'], 'm', 'T')
172
 
    canvas.Write()
173
 
    canvas.Print( "MagneticField.eps", "eps")
174
 
    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['pid'], 'm')
175
 
    canvas.Write()
176
 
    canvas.Print( "PID.eps", "eps")
177
 
 
178
 
 
179
 
    bunch_up = ana_list[ upstream_reference ]
180
 
    bunch_down = ana_list[ downstream_reference ]
181
 
 
182
 
    up_means = bunch_up.mean( ['x','y','z','px','py','pz','pt','energy','bz'] )
183
 
    down_means = bunch_down.mean( ['x','y','z','px','py','pz','pt','energy','bz'] )
184
 
 
185
 
    print 
186
 
    print " === Upstream Results === "
187
 
    print
188
 
    print "  Collective Properties"
189
 
    print "Emittance              : ", bunch_up.get_emittance( ['x','y'] )
190
 
    print "Beta                   : ", bunch_up.get_beta( ['x','y'] )
191
 
    print "Weight                 : ", bunch_up.bunch_weight()
192
 
    print
193
 
    print "  Mean Values"
194
 
    print "X Position             : ", up_means['x']
195
 
    print "Y Position             : ", up_means['y']
196
 
    print "Z Position             : ", up_means['z']
197
 
    print "X Momentum             : ", up_means['px']
198
 
    print "Y Momentum             : ", up_means['py']
199
 
    print "Z Momentum             : ", up_means['pz']
200
 
    print "Transverese Momentum   : ", up_means['pt']
201
 
    print "Energy                 : ", up_means['energy']
202
 
    print "Z Magnetic Field       : ", up_means['bz']
203
 
 
204
 
    print 
205
 
    print " === Downstream Results === "
206
 
    print
207
 
    print "  Collective Properties"
208
 
    print "Emittance              : ", bunch_down.get_emittance( ['x','y'] )
209
 
    print "Beta                   : ", bunch_down.get_beta( ['x','y'] )
210
 
    print "Weight                 : ", bunch_down.bunch_weight()
211
 
    print
212
 
    print "  Mean Values"
213
 
    print "X Position             : ", down_means['x']
214
 
    print "Y Position             : ", down_means['y']
215
 
    print "Z Position             : ", down_means['z']
216
 
    print "X Momentum             : ", down_means['px']
217
 
    print "Y Momentum             : ", down_means['py']
218
 
    print "Z Momentum             : ", down_means['pz']
219
 
    print "Transverese Momentum   : ", down_means['pt']
220
 
    print "Energy                 : ", down_means['energy']
221
 
    print "Z Magnetic Field       : ", down_means['bz']
222
 
 
223
 
 
224
 
    if inspection_plane > 0 :
225
 
      print
226
 
      print "Saving Inspection Plane"
227
 
      print "Please Wait."
228
 
      print
229
 
      with open( 'inspection_plane.txt', 'w' ) as outfile :
230
 
        outfile.write( '# X   Y   Z   Px   Py   Pz\n' )
231
 
        for hit in ana_list[ inspection_plane ] :
232
 
          outfile.write( str( hit.get( 'x' ) ) + ' ' + 
233
 
                         str( hit.get( 'y' ) ) + ' ' +
234
 
                         str( hit.get( 'z' ) ) + ' ' +
235
 
                         str( hit.get( 'px' ) ) + ' ' +
236
 
                         str( hit.get( 'py' ) ) + ' ' +
237
 
                         str( hit.get( 'pz' ) ) + ' ' +
238
 
                         str( hit.get( 'pid' ) ) + ' ' + 
239
 
                         str( hit.get_weight() ) + '\n' )
240
 
 
241
 
 
242
 
 
243
 
 
244
 
if __name__ == "__main__":
245
 
    main()
246
 
    print "Press return to close"
247
 
    raw_input()
248
 
 
249
 
 
250
 
 
251
 
 
252
 
 
253
 
 
254
 
 
255
 
 
256
 
 
257
 
 
258
 
 
259
 
 
260
 
 
261
 
 
262
 
 
263
 
 
264
 
 
265
 
 
266
 
#def main():
267
 
#    if len( sys.argv ) <= 1 :
268
 
#      print "No File Supplied."
269
 
#      sys.exit( 0 )
270
 
#
271
 
#    filename = sys.argv[1]
272
 
#
273
 
#    print "Loading File"
274
 
#    print filename
275
 
#    print
276
 
#
277
 
#    primary_bunch = Bunch.new_list_from_read_builtin('maus_root_primary', filename)
278
 
#    bunch_list = Bunch.new_list_from_read_builtin('maus_root_virtual_hit', filename)
279
 
#    ana_list = []
280
 
#
281
 
#    print "Read", len(bunch_list), "Bunches"
282
 
#
283
 
#    for bun in bunch_list :
284
 
#      if bun.bunch_weight() <= 1.0e-9 :
285
 
#        continue
286
 
#      bun.cut( { 'amplitude x y' : cut_distance}, operator.ge )
287
 
#      if bun.bunch_weight() <= 1.0e-9 :
288
 
#        continue
289
 
#      bun.cut( {'pid':11}, operator.eq )
290
 
#      bun.cut( {'pid':-11}, operator.eq  )
291
 
#
292
 
#    for bun in bunch_list :
293
 
#      if bun.bunch_weight() > 1.0e-9 and len( bun ) > 1 :
294
 
#        ana_list.append( bun )
295
 
#
296
 
#    print "Found", len(ana_list), "Non Empty Bunches"
297
 
#
298
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'bunch_weight', [''], 'm')
299
 
#    canvas.Print( "BunchWeight.eps", "eps")
300
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'm', 'm')
301
 
#    canvas.Print( "Emittance.eps", "eps")
302
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'beta', ['x', 'y'], 'm', 'm')
303
 
#    canvas.Print( "BetaFunc.eps", "eps")
304
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['energy'], 'm', 'MeV')
305
 
#    canvas.Print( "Energy.eps", "eps")
306
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['bz'], 'm', 'T')
307
 
#    canvas.Print( "MagneticField.eps", "eps")
308
 
#    canvas, hist, graph = Bunch.root_graph(ana_list, 'mean', ['z'], 'mean', ['pid'], 'm')
309
 
#    canvas.Print( "PID.eps", "eps")
310
 
##    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['p'], 'm', 'MeV/c')
311
 
##    canvas.Print( "Momentum.eps", "eps")
312
 
#
313
 
#if __name__ == "__main__":
314
 
#    main()
315
 
#    print "Press return to close"
316
 
#    raw_input()
317
 
 
318
 
 
319
 
 
320
 
 
321
 
 
322
 
 
323