← 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_virtualPlanes.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.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
 
radial_cut = 189.
33
 
 
34
 
upstream_reference = 14
35
 
downstream_reference = 50
36
 
 
37
 
#inspection_plane = 56
38
 
inspection_plane = -1 # Negative means don't write a file
39
 
 
40
 
write_refernce_planes = False
41
 
print_all_planes = False
42
 
print_all_plots_all_planes = False
43
 
cut_pz_tails = False
44
 
apply_cuts = True
45
 
 
46
 
 
47
 
 
48
 
pz_tail = 0.5
49
 
 
50
 
 
51
 
def main():
52
 
    if len( sys.argv ) <= 1 :
53
 
      print "No File Supplied."
54
 
      sys.exit( 0 )
55
 
 
56
 
    filename = sys.argv[1]
57
 
 
58
 
    print "Loading File"
59
 
    print filename
60
 
    print
61
 
 
62
 
    bunch_list_tmp = Bunch.new_list_from_read_builtin('maus_root_virtual_hit', filename)
63
 
 
64
 
    print "Read", len(bunch_list_tmp), "bunches"
65
 
 
66
 
    for bun in bunch_list_tmp :
67
 
      bun.cut( {'pid':-13}, operator.ne, global_cut = True)
68
 
      if bun.bunch_weight() < 1.e-9 :
69
 
        continue
70
 
 
71
 
 
72
 
    if cut_pz_tails :
73
 
      for bun in bunch_list_tmp :
74
 
        if bun.bunch_weight() < 1.e-9 :
75
 
          continue
76
 
        mean = bun.mean( ['pz'] )['pz']
77
 
        print "PZ =", ( mean - pz_tail ), mean, ( mean + pz_tail )
78
 
        bun.cut( { 'pz' : mean + pz_tail }, operator.ge, global_cut = False )
79
 
        bun.cut( { 'pz' : mean - pz_tail }, operator.le, global_cut = False )
80
 
 
81
 
 
82
 
    if apply_cuts :
83
 
      for bun in bunch_list_tmp :
84
 
        if bun.bunch_weight() < 1.e-9 :
85
 
          continue
86
 
        bun.cut( {'pid':-13}, operator.ne, global_cut = True)
87
 
        if bun.bunch_weight() < 1.e-9 :
88
 
          continue
89
 
        bun.cut( {'r' : radial_cut}, operator.ge, global_cut = True)
90
 
        if bun.bunch_weight() < 1.e-9 :
91
 
          continue
92
 
        bun.transmission_cut( bunch_list_tmp[ downstream_reference ], True )
93
 
 
94
 
 
95
 
    bunch_list = [ bun for bun in bunch_list_tmp if len( bun ) > 1 and bun.bunch_weight()>1.0 ]
96
 
 
97
 
    outfile = ROOT.TFile( "emittance_from_virtualPlanes.root", "RECREATE" )
98
 
 
99
 
#    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'standard_deviation', ['pz'], 'MeV/c')
100
 
#    canvas.Write()
101
 
#    canvas.Print( "RMS_pz.eps", "eps")
102
 
#    canvas.Close()
103
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'bunch_weight', [''], 'm')
104
 
    canvas.Write()
105
 
    canvas.Print( "BunchWeight.eps", "eps")
106
 
    canvas.Close()
107
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'm', 'm')
108
 
    canvas.Write()
109
 
    canvas.Print( "Emittance.eps", "eps")
110
 
    canvas.Close()
111
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['x', 'y'], 'm', 'm')
112
 
    canvas.Write()
113
 
    canvas.Print( "BetaFunc.eps", "eps")
114
 
    canvas.Close()
115
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['energy'], 'm', 'MeV')
116
 
    canvas.Write()
117
 
    canvas.Print( "Energy.eps", "eps")
118
 
    canvas.Close()
119
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['bz'], 'm', 'T')
120
 
    canvas.Write()
121
 
    canvas.Print( "MagneticField.eps", "eps")
122
 
    canvas.Close()
123
 
    canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['pid'], 'm')
124
 
    canvas.Write()
125
 
    canvas.Print( "PID.eps", "eps")
126
 
    canvas.Close()
127
 
 
128
 
    bunch_up = bunch_list[ upstream_reference ]
129
 
    bunch_down = bunch_list[ downstream_reference ]
130
 
 
131
 
    up_means = bunch_up.mean( ['x','y','z','px','py','pz','pt','energy','bz'] )
132
 
    down_means = bunch_down.mean( ['x','y','z','px','py','pz','pt','energy','bz'] )
133
 
 
134
 
    if write_refernce_planes is True :
135
 
      print
136
 
      print "Writing Raw Output Files"
137
 
      print
138
 
 
139
 
      with open( "upstream_virtual_hits.txt", 'w' ) as outfile :
140
 
        outfile.write( '#  Position (x,y,z)  |  Momentum (px,py,pz)  |  Magnetic File (bz) \n\n' )
141
 
 
142
 
        for hit in bunch_up :
143
 
          outfile.write( str(hit.get('x')) + ' '
144
 
                         + str(hit.get( 'y' )) + ' '
145
 
                         + str(hit.get( 'z' )) + ' '
146
 
                         + str(hit.get( 'px' )) + ' '
147
 
                         + str(hit.get( 'py' )) + ' '
148
 
                         + str(hit.get( 'pz' )) + ' '
149
 
                         + str(hit.get( 'bz' )) + '\n' )
150
 
 
151
 
      with open( "downstream_virtual_hits.txt", 'w' ) as outfile :
152
 
        outfile.write( '#  Position (x,y,z)  |  Momentum (px,py,pz)  |  Magnetic File (bz) \n\n' )
153
 
 
154
 
        for hit in bunch_down :
155
 
          outfile.write( str(hit.get('x')) + ' '
156
 
                         + str(hit.get( 'y' )) + ' '
157
 
                         + str(hit.get( 'z' )) + ' '
158
 
                         + str(hit.get( 'px' )) + ' '
159
 
                         + str(hit.get( 'py' )) + ' '
160
 
                         + str(hit.get( 'pz' )) + ' '
161
 
                         + str(hit.get( 'bz' )) + '\n' )
162
 
 
163
 
 
164
 
    print 
165
 
    print " === Upstream Results === "
166
 
    print
167
 
    print "  Collective Properties"
168
 
    print "Emittance              : ", bunch_up.get_emittance( ['x','y'] )
169
 
    print "Beta                   : ", bunch_up.get_beta( ['x','y'] )
170
 
    print "Weight                 : ", bunch_up.bunch_weight()
171
 
    print
172
 
    print "  Mean Values"
173
 
    print "X Position             : ", up_means['x']
174
 
    print "Y Position             : ", up_means['y']
175
 
    print "Z Position             : ", up_means['z']
176
 
    print "X Momentum             : ", up_means['px']
177
 
    print "Y Momentum             : ", up_means['py']
178
 
    print "Z Momentum             : ", up_means['pz']
179
 
    print "Transverese Momentum   : ", up_means['pt']
180
 
    print "Energy                 : ", up_means['energy']
181
 
    print "Z Magnetic Field       : ", up_means['bz']
182
 
 
183
 
    print 
184
 
    print " === Downstream Results === "
185
 
    print
186
 
    print "  Collective Properties"
187
 
    print "Emittance              : ", bunch_down.get_emittance( ['x','y'] )
188
 
    print "Beta                   : ", bunch_down.get_beta( ['x','y'] )
189
 
    print "Weight                 : ", bunch_down.bunch_weight()
190
 
    print
191
 
    print "  Mean Values"
192
 
    print "X Position             : ", down_means['x']
193
 
    print "Y Position             : ", down_means['y']
194
 
    print "Z Position             : ", down_means['z']
195
 
    print "X Momentum             : ", down_means['px']
196
 
    print "Y Momentum             : ", down_means['py']
197
 
    print "Z Momentum             : ", down_means['pz']
198
 
    print "Transverese Momentum   : ", down_means['pt']
199
 
    print "Energy                 : ", down_means['energy']
200
 
    print "Z Magnetic Field       : ", down_means['bz']
201
 
 
202
 
    if inspection_plane > 0 :
203
 
      print 
204
 
      print "Saving Inspection Plane"
205
 
      print "Please Wait."
206
 
      print
207
 
 
208
 
      with open( 'inspection_plane.txt', 'w' ) as outfile :
209
 
        outfile.write( '# X   Y   Z   Px   Py   Pz\n' )
210
 
        for hit in bunch_list[ inspection_plane ] :
211
 
          outfile.write( str( hit.get( 'x' ) ) + ' ' + 
212
 
                         str( hit.get( 'y' ) ) + ' ' +
213
 
                         str( hit.get( 'z' ) ) + ' ' +
214
 
                         str( hit.get( 'px' ) ) + ' ' +
215
 
                         str( hit.get( 'py' ) ) + ' ' +
216
 
                         str( hit.get( 'pz' ) ) + ' ' +
217
 
                         str( hit.get( 'pid' ) ) + ' ' +
218
 
                         str( hit.get_weight() ) + '\n' )
219
 
      print 'inspection_plane.txt Created'
220
 
 
221
 
    if print_all_planes is True :
222
 
      print
223
 
      print "Saving Plane Data"
224
 
      print "Please Wait."
225
 
      print
226
 
      with open( 'plane_data.txt', 'w' ) as outfile :
227
 
        outfile.write( '# Mean X    Mean Y   Mean Z    Mean Px   Mean Py   Mean Pz   Mean Pt   Mean Bz  Energy   Beta    Emittance   Weight \n' )
228
 
        bunch_count = 0
229
 
        for bun in bunch_list :
230
 
          means = bun.mean( ['x','y','z','px','py','pz','pt','energy','bz'] )
231
 
          outfile.write( str( bunch_count ) + ' ' +
232
 
                         str( means[ 'x' ] ) + ' ' +
233
 
                         str( means[ 'y' ] ) + ' ' +
234
 
                         str( means[ 'z' ] ) + ' ' +
235
 
                         str( means[ 'px' ] ) + ' ' +
236
 
                         str( means[ 'py' ] ) + ' ' +
237
 
                         str( means[ 'pz' ] ) + ' ' +
238
 
                         str( means[ 'pt' ] ) + ' ' +
239
 
                         str( means[ 'bz' ] ) + ' ' +
240
 
                         str( means[ 'energy' ] ) + ' ' +
241
 
                         str( bun.get_emittance( [ 'x','y' ] ) ) + ' ' +
242
 
                         str( bun.get_beta( ['x','y'] ) ) + ' ' +
243
 
                         str( bun.bunch_weight() ) + "\n" )
244
 
          bunch_count += 1
245
 
      print 'plane_data.txt Created'
246
 
 
247
 
    if print_all_plots_all_planes is True :
248
 
      print 
249
 
      print "Saving plots for every plane. Theres going to be a lot of them!"
250
 
      print "Please Wait."
251
 
      print
252
 
 
253
 
      for num in range( len( bunch_list ) ) :
254
 
        bunch = bunch_list[num]
255
 
        prefix = 'Virtual_plane_' + str( num ) + '_'
256
 
 
257
 
        canvas, hist = bunch.root_histogram( 'x', 'mm', 'y', 'mm' )
258
 
        canvas.Print( prefix+'X-Y.eps', 'eps' )
259
 
        canvas.Close()
260
 
        canvas, hist = bunch.root_histogram( 'px', 'MeV/c', 'py', 'MeV/c' )
261
 
        canvas.Print( prefix+'PX-PY.eps', 'eps' )
262
 
        canvas.Close()
263
 
        canvas, hist = bunch.root_histogram( 'x', 'mm', 'px', 'MeV/c' )
264
 
        canvas.Print( prefix+'X-PX.eps', 'eps' )
265
 
        canvas.Close()
266
 
        canvas, hist = bunch.root_histogram( 'y', 'mm', 'py', 'MeV/c' )
267
 
        canvas.Print( prefix+'Y-PY.eps', 'eps' )
268
 
        canvas.Close()
269
 
        canvas, hist = bunch.root_histogram( 'pz', 'MeV/c' )
270
 
        canvas.Print( prefix+'PZ.eps', 'eps' )
271
 
        canvas.Close()
272
 
        canvas, hist = bunch.root_histogram( 'pt', 'MeV/c' )
273
 
        canvas.Print( prefix+'PT.eps', 'eps' )
274
 
        canvas.Close()
275
 
 
276
 
 
277
 
if __name__ == "__main__":
278
 
    main()
279