← Back to branch summary

~chris-rogers/maus/emr_mc_digitization

« back to all changes in this revision

Viewing changes to src/map/MapPyBeamlineSimulation/CallG4bl.py

  • Committer: Chris Rogers
  • Date: 2014-04-16 11:48:45 UTC
  • mfrom: (707 merge)
  • mto: This revision was merged to the branch mainline in revision 711.
  • Revision ID: chris.rogers@stfc.ac.uk-20140416114845-h3u3q7pdcxkxvovs
Update to trunk

Show diffs side-by-side

added added

removed removed

Lines of Context:
src/map/MapPyBeamlineSimulation/CallG4bl.py
 
1
"""
 
2
Class to call G4BL and generate primaries for MAUS
 
3
"""
 
4
# This file is part of MAUS: http://micewww.pp.rl.ac.uk/projects/maus
 
5
#
 
6
# MAUS is free software: you can redistribute it and/or modify
 
7
# it under the terms of the GNU General Public License as published by
 
8
# the Free Software Foundation, either version 3 of the License, or
 
9
# (at your option) any later version.
 
10
#
 
11
# MAUS is distributed in the hope that it will be useful,
 
12
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
13
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
14
# GNU General Public License for more details.
 
15
#
 
16
# You should have received a copy of the GNU General Public License
 
17
# along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 
18
 
 
19
#pylint: disable=W0611
 
20
 
 
21
import cdb 
 
22
from cdb import Beamline
 
23
import random
 
24
import sys
 
25
from subprocess import call
 
26
import json
 
27
import os
 
28
import fileinput   
 
29
import math
 
30
import numpy
 
31
 
 
32
class CallG4bl: #pylint: disable=R0903
 
33
    """
 
34
    Calls G4BL and writes output to python dictionary
 
35
 
 
36
    Creates primaries for MAUS depending on the parameters
 
37
    set in the dictionary 'g4bl' in the json configuration
 
38
    document
 
39
    - magnet currents and proton absorber thickness can be set
 
40
    along with the number of particles to generate
 
41
    - to retrieve magnet currents and proton absorber
 
42
    thickness from CDB set get_magnet_currents_pa_cdb to
 
43
    True and run_number to the run of interest in json
 
44
    configuration document
 
45
    """
 
46
    
 
47
    def __init__(#pylint: disable=C0103, R0903, R0912, R0913, R0914, R0915 
 
48
                 self,newline='', file_path='', theta=0, deltaz=0, 
 
49
                 path_g4bl='', output_path='', run_number=0, \
 
50
                 get_magnet_currents_pa_cdb='', random_seed=0):
 
51
        """
 
52
        - particles will be a dictionary filled by G4BL
 
53
        - pdgid masses are input to calculate total energy,
 
54
        masses are in MeV/C^2
 
55
        """
 
56
 
 
57
        self.particles = 0
 
58
        self.pdgid_mass = {}
 
59
 
 
60
        self.pdgid_mass = {'211':139.5700, '-211':139.5700, '2212':938.27231, \
 
61
        '-13':105.6584, '13':105.6584, '11':0.5109906, '-11':0.5109906, \
 
62
        '22':0, '2112':939.5653, '14':0, '-14':0} #pylint: disable=C0103
 
63
 
 
64
        try:
 
65
            off = open(file_path+'/MAY09-B1B2-positives.in',"r")
 
66
        except IOError as err:
 
67
            raise IOError(err)
 
68
        text = off.read()
 
69
 
 
70
        if str(get_magnet_currents_pa_cdb) in ['True']:
 
71
            try:
 
72
                beamline = Beamline()
 
73
                sf3 = open(file_path+'/key.txt',"w")
 
74
                sf3.close()
 
75
                for k, v in beamline.get_beamline_for_run(run_number).iteritems(): #pylint: disable = C0301
 
76
                    for k, v in v.iteritems():
 
77
                        if k == 'magnets':
 
78
                            for k, v in v.iteritems():
 
79
                                newlines = str(k)+'\n'
 
80
                                sff = open(file_path+'/key.txt',"r")
 
81
                                text5 = sff.read()
 
82
                                sff.close()
 
83
                                sf2 = open(file_path+'/key.txt',"w")
 
84
                                sf2.write(text5)
 
85
                                sf2.write(newlines)
 
86
                                sf2.close()
 
87
                        if k == 'proton_absorber_thickness':
 
88
                            newlines = str(k)+'\n'
 
89
                            sff = open(file_path+'/key.txt',"r")
 
90
                            text5 = sff.read()
 
91
                            sff.close()
 
92
                            sf2 = open(file_path+'/key.txt',"w")
 
93
                            sf2.write(text5)
 
94
                            sf2.write(newlines)
 
95
                            sf2.close()
 
96
                line = []
 
97
                for row in fileinput.input([file_path+'/key.txt']):
 
98
                    line.append(row[:-1])
 
99
                sf2.close()
 
100
                sf4 = open(file_path+'/value.txt',"w")
 
101
                sf4.close()
 
102
                for k, v in beamline.get_beamline_for_run(run_number).iteritems(): #pylint: disable = C0301
 
103
                    for k, v in v.iteritems():
 
104
                        if k == 'magnets':
 
105
                            for k, v in v.iteritems():
 
106
                                if k in ['q1', 'q3']:
 
107
                                    for k, v in v.iteritems():
 
108
                                        if k == 'set_current':
 
109
                                            v = v / 96
 
110
                                            newlines = str(v)+'\n'
 
111
                                            sf5 = open(file_path+'/value.txt',"r") #pylint: disable = C0301
 
112
                                            text6 = sf5.read()
 
113
                                            sf5.close()
 
114
                                            sf6 = open(file_path+'/value.txt',"w") #pylint: disable = C0301
 
115
                                            sf6.write(text6)
 
116
                                            sf6.write(newlines)
 
117
                                            sf6.close()
 
118
                                if k in ['q2']:
 
119
                                    for k, v in v.iteritems():
 
120
                                        if k == 'set_current':
 
121
                                            v = - v / 96
 
122
                                            newlines = str(v)+'\n'
 
123
                                            sf5 = open(file_path+'/value.txt',"r") #pylint: disable = C0301
 
124
                                            text6 = sf5.read()
 
125
                                            sf5.close() 
 
126
                                            sf6 = open(file_path+'/value.txt',"w") #pylint: disable = C0301
 
127
                                            sf6.write(text6)
 
128
                                            sf6.write(newlines)
 
129
                                            sf6.close()
 
130
                                if k in ['d2']: #pylint: disable = W0631
 
131
                                    for k, v in v.iteritems(): #pylint: disable = W0631, C0301
 
132
                                        if k == 'set_current':
 
133
                                            coeff = [39.59, -55.998, 256.914, -v] #pylint: disable = C0301
 
134
                                            roots = numpy.roots(coeff)
 
135
                                            root =  roots[2]
 
136
                                            newlines = str(root)[1:-4]+'\n'
 
137
                                            sf5 = open(file_path+'/value.txt',"r") #pylint: disable = C0301
 
138
                                            text6 = sf5.read()
 
139
                                            sf5.close()
 
140
                                            sf6 = open(file_path+'/value.txt',"w") #pylint: disable = C0301
 
141
                                            sf6.write(text6)
 
142
                                            sf6.write(newlines)
 
143
                                            sf6.close()
 
144
                                if k in ['ds']:
 
145
                                    for k, v in v.iteritems():
 
146
                                        if k == 'set_current':
 
147
                                            v = v / 174
 
148
                                            newlines = str(v)+'\n'
 
149
                                            sf5 = open(file_path+'/value.txt',"r") #pylint: disable = C0301
 
150
                                            text6 = sf5.read()
 
151
                                            sf5.close()
 
152
                                            sf6 = open(file_path+'/value.txt',"w") #pylint: disable = C0301
 
153
                                            sf6.write(text6)
 
154
                                            sf6.write(newlines)
 
155
                                            sf6.close()
 
156
                                if k in ['d1']:
 
157
                                    for k, v in v.iteritems():
 
158
                                        if k == 'set_current':
 
159
                                            coeff = [39.59, -55.998, 256.914, -v] #pylint: disable = C0301
 
160
                                            roots = numpy.roots(coeff)
 
161
                                            root =  roots[2]
 
162
                                            newlines = str(root)[1:-4]+'\n'
 
163
                                            sf5 = open(file_path+'/value.txt',"r") #pylint: disable = C0301
 
164
                                            text6 = sf5.read()
 
165
                                            sf5.close()
 
166
                                            sf6 = open(file_path+'/value.txt',"w") #pylint: disable = C0301
 
167
                                            sf6.write(text6)
 
168
                                            sf6.write(newlines)
 
169
                                            sf6.close()
 
170
                        if k == 'proton_absorber_thickness':
 
171
                            newlines = str(v)+'\n'
 
172
                            sf5 = open(file_path+'/value.txt',"r")
 
173
                            text6 = sf5.read()
 
174
                            sf5.close()
 
175
                            sf6 = open(file_path+'/value.txt',"w")
 
176
                            sf6.write(text6)
 
177
                            sf6.write(newlines)
 
178
                            sf6.close()
 
179
                     
 
180
                line2 = []
 
181
                for row in fileinput.input([file_path+'/value.txt']):
 
182
                    line2.append(row[:-1])
 
183
                sf6.close()
 
184
 
 
185
                sf7 = open(file_path+'/magnet_currents.txt',"w")
 
186
                sf7.close()
 
187
 
 
188
                for c, d in zip(line, line2):
 
189
                    text2 = 'param {0}={1}'.format(c, d)+'\n'              
 
190
                    sf8 = open(file_path+'/magnet_currents.txt',"r")
 
191
                    text10 = sf8.read()
 
192
                    sf8.close()
 
193
                    sf9 = open(file_path+'/magnet_currents.txt',"w")
 
194
                    sf9.write(text10)
 
195
                    sf9.write(text2)
 
196
                    sf9.close()
 
197
 
 
198
                sf8 = open(file_path+'/magnet_currents.txt',"r")
 
199
                text3 = sf8.read()
 
200
         
 
201
                os.remove(file_path+'/key.txt')
 
202
                os.remove(file_path+'/value.txt')
 
203
                                
 
204
            except cdb.CdbTemporaryError:
 
205
                sys.excepthook(*sys.exc_info())
 
206
                print """
 
207
                Failed to run - check network connection!
 
208
                """
 
209
 
 
210
            try:  
 
211
                sf = open(file_path+'/MAY09-B1B2-positives-param.in', "w")   
 
212
            except IOError as err:
 
213
                raise IOError(err)
 
214
            sf.write(newline)
 
215
            sf.write('param random_seed='+str(random_seed)+'\n')
 
216
            sf.write(text3)
 
217
            sf.write(text)
 
218
            sf.close()
 
219
            off.close()
 
220
 
 
221
        else:
 
222
 
 
223
            try:
 
224
                sf = open(file_path+'/MAY09-B1B2-positives-param.in', "w")
 
225
            except IOError as err:
 
226
                raise IOError(err)
 
227
            sf.write(newline)
 
228
            sf.write('param random_seed='+str(random_seed))
 
229
            sf.write(text)
 
230
            sf.close()
 
231
            off.close()
 
232
        
 
233
        g4bl_command_line_input = path_g4bl + '/g4bl ' + file_path + '/MAY09-B1B2-positives-param.in' #pylint: disable = C0301
 
234
        try:
 
235
            call(g4bl_command_line_input, shell=True)
 
236
        except Exception: #pylint: disable = W0703
 
237
            print('Error: Cannot execute G4Beamline.')
 
238
 
 
239
        try:
 
240
            of = open(output_path+'/G4BLoutput.txt',"r")
 
241
        except IOError as err:
 
242
            raise IOError(err)   
 
243
        line = []
 
244
        for row in fileinput.input([output_path+'/G4BLoutput.txt']):
 
245
            line.append(row)
 
246
        of.close()
 
247
 
 
248
        for i in range(3):
 
249
            line.pop(0)
 
250
        random.shuffle(line)
 
251
 
 
252
        self.particles = {} 
 
253
        theta = math.radians(theta)
 
254
        for i in range(0, len(line)):
 
255
            element = []
 
256
            for word in line[i].split(' '):
 
257
                element.append(word) 
 
258
            px = float(element[3])
 
259
            py = float(element[4])
 
260
            pz = float(element[5])
 
261
         
 
262
            rest_mass = self.pdgid_mass[str(element[7])]
 
263
            part_energy = math.sqrt(px**2+py**2+pz**2+rest_mass**2)
 
264
         
 
265
            try:     
 
266
                key = 'entry'+str(i)
 
267
 
 
268
                self.particles[key] = dict(position = dict(\
 
269
                x = math.cos(theta) * float(element[0]) + 737.63,
 
270
                y = float(element[1]), 
 
271
                z = math.sin(theta) * float(element[0]) - deltaz))
 
272
 
 
273
                self.particles[key].update(momentum = dict(\
 
274
                x = math.cos(theta) * float(element[3]) - math.sin(theta) * float(element[5]), #pylint: disable = C0301
 
275
                y = float(element[4]),
 
276
                z = math.sin(theta) * float(element[3]) + math.cos(theta) * float(element[5]))) #pylint: disable = C0301
 
277
                self.particles[key].update(time = float(element[6]), \
 
278
                particle_id = int(element[7]))
 
279
                self.particles[key].update(random_seed = random_seed, \
 
280
                energy = part_energy)    
 
281
            except IOError as err:
 
282
                print('Bad G4beamline file, not enough elements or empty file')
 
283
                raise IOError(err)
 
284
#        with open('G4BLoutput.txt', 'w') as outfile:
 
285
#              json.dump(self.particles, outfile)
 
286
 
 
287
#        for i in range(0, len(line)):
 
288
#            for j in range(i + 1, len(line)):
 
289
###                if self.particles['entry' + str(i)]['random_seed']\
 
290
#                   == self.particles['entry' + str(j)]['random_seed']:
 
291
#                    self.particles['entry' + str(j)]['random_seed']\
 
292
#                   =self.particles['entry' + str(j)]['random_seed']\
 
293
#                   +1            
 
294
 
 
295
        if(os.path.exists(file_path + '/magnet_currents.txt') == True):
 
296
            os.remove(file_path + '/magnet_currents.txt')