1 | #!/bin/env python |
---|
2 | """ |
---|
3 | Fix frame wrapping events in powgen |
---|
4 | |
---|
5 | """ |
---|
6 | import os, sys, struct |
---|
7 | import numpy, math, fnmatch |
---|
8 | from shutil import move, copy2 |
---|
9 | |
---|
10 | |
---|
11 | __version__ = "0.3" |
---|
12 | |
---|
13 | INSTRUMENT = 'PG3' |
---|
14 | TOP_DIR = '.' |
---|
15 | DSPACEMAP = 'powgen_dspacemap_d1370_2010_09_12.dat' |
---|
16 | |
---|
17 | PIX_MASK = 0x000FFFFF # PG3 "good" pixels |
---|
18 | FIX_MASK = 0x00F00000 # PG3 mask for fixed pixels |
---|
19 | |
---|
20 | TOF_FLAG = 0x00100000 # |
---|
21 | RES_FLAG = 0x00200000 # |
---|
22 | |
---|
23 | MAX_PIXEL = 300000 # |
---|
24 | |
---|
25 | SNS_FREQ = 60.0 # 60 Hz |
---|
26 | |
---|
27 | MIN_EVENTS = 4000 # need at least this many neutrons |
---|
28 | |
---|
29 | NeutronEvent = numpy.dtype([('tof','uint32'),('pix','uint32')]) # event data format |
---|
30 | DataChunk = NeutronEvent.itemsize*1024L*1024L*24L # read 24 Mevents at once, otherwise a small DFS might choke |
---|
31 | |
---|
32 | |
---|
33 | fix_log = sys.stdout |
---|
34 | |
---|
35 | def log_msg(*args): |
---|
36 | "Simple log - to the screen and to the file" |
---|
37 | for arg in args: |
---|
38 | fix_log.write(arg) |
---|
39 | if fix_log != sys.stdout: |
---|
40 | print arg, |
---|
41 | else: |
---|
42 | fix_log.write('\n') |
---|
43 | if fix_log != sys.stdout: |
---|
44 | print |
---|
45 | |
---|
46 | |
---|
47 | |
---|
48 | def find_prenexus(file_type='*neutron_event.dat', top_dir=TOP_DIR): |
---|
49 | "Find DAS files" |
---|
50 | for root, dirs, files in os.walk(top_dir): |
---|
51 | for file_name in files: |
---|
52 | if fnmatch.fnmatch(file_name, file_type): |
---|
53 | yield os.path.join(root, file_name) |
---|
54 | |
---|
55 | |
---|
56 | def plot_dspace(data, file_name, log_scale=True): |
---|
57 | """Here is some plotting just for purpose of seeing |
---|
58 | that selected neutrons not needed for fix""" |
---|
59 | base_file = os.path.basename(file_name) |
---|
60 | d_max = 20.0 |
---|
61 | n_bin = 2000 |
---|
62 | pylab.figure() |
---|
63 | pylab.title('%s All Modules' % file_name) |
---|
64 | for label, d in data: |
---|
65 | dhist,dbins = numpy.histogram(d, bins=n_bin, range=(0.0,d_max), new=True) |
---|
66 | if log_scale: |
---|
67 | pylab.semilogy(dbins[:-1], dhist, marker='.', linestyle='steps', label=label) |
---|
68 | else: |
---|
69 | pylab.plot (dbins[:-1], dhist, marker='.', linestyle='steps', label=label) |
---|
70 | |
---|
71 | pylab.xlabel('d [A]') |
---|
72 | pylab.legend() |
---|
73 | |
---|
74 | |
---|
75 | |
---|
76 | def check_fixed(event_file, dspmap, do_plot=False): |
---|
77 | """ |
---|
78 | Test if file already fixed |
---|
79 | """ |
---|
80 | |
---|
81 | events_fd = open(event_file, 'rb') |
---|
82 | events_buf = events_fd.read(NeutronEvent.itemsize) |
---|
83 | if not events_buf: |
---|
84 | return True |
---|
85 | events = numpy.fromstring(events_buf, dtype=NeutronEvent) |
---|
86 | events_fd.close() |
---|
87 | |
---|
88 | if not (events[0]['pix'] & FIX_MASK): |
---|
89 | return False |
---|
90 | |
---|
91 | |
---|
92 | if do_plot: |
---|
93 | events_fd = open(event_file, 'rb') |
---|
94 | events_buf = events_fd.read(DataChunk) |
---|
95 | if not events_buf: |
---|
96 | return True |
---|
97 | events = numpy.fromstring(events_buf, dtype=NeutronEvent) |
---|
98 | events_fd.close() |
---|
99 | |
---|
100 | pix = events['pix'] |
---|
101 | tof = events['tof'] |
---|
102 | |
---|
103 | pix = pix & PIX_MASK |
---|
104 | d = 2*dspmap[pix]*tof |
---|
105 | plot_dspace([('Previously Fixed', d)], event_file) |
---|
106 | return True |
---|
107 | |
---|
108 | |
---|
109 | |
---|
110 | def fix_file_loop(event_file, dspace_file=None, verbose=True, do_force=False, do_plot=False): |
---|
111 | "Main function" |
---|
112 | |
---|
113 | log_msg("** processing file %s" % event_file) |
---|
114 | |
---|
115 | n_events = os.stat(event_file)[6]/NeutronEvent.itemsize # python 2.5 has better definition of os.stat |
---|
116 | if n_events < MIN_EVENTS : |
---|
117 | log_msg("!! event file too small, need at least %d events (got %d)" % (MIN_EVENTS, n_events)) |
---|
118 | return |
---|
119 | |
---|
120 | |
---|
121 | dspmap = numpy.memmap(dspace_file, dtype='float64' , mode='r' ) |
---|
122 | |
---|
123 | |
---|
124 | if check_fixed(event_file, dspmap, do_plot) and not do_force: |
---|
125 | log_msg('!! event file previously fixed; aborting fix of this file') |
---|
126 | return |
---|
127 | |
---|
128 | # copy files first |
---|
129 | # orig_file = event_file + '.orig' |
---|
130 | orig_file = os.path.join(os.path.dirname(event_file), os.path.basename(event_file).split('.')[0]) + '_orig.dat' |
---|
131 | |
---|
132 | |
---|
133 | if os.path.exists(orig_file) and do_force: |
---|
134 | log_msg('-- restoring original event file %s -> %s' % (orig_file, event_file)) |
---|
135 | os.remove(event_file) |
---|
136 | copy2(orig_file, event_file) |
---|
137 | else: |
---|
138 | log_msg('-- making backup copy of the original event file %s <- %s' % (orig_file, event_file)) |
---|
139 | copy2(event_file, orig_file) |
---|
140 | |
---|
141 | min_usperA = 15750.0 |
---|
142 | DifCr = 15000.0 |
---|
143 | LDifCr = math.log10(DifCr) |
---|
144 | K = 3.22 # gradient of lin |
---|
145 | sns_frame = 1e7/SNS_FREQ # convert to units of 100 ns |
---|
146 | |
---|
147 | usperAr = numpy.zeros(MAX_PIXEL, dtype='float') # microseconds per Agstrom |
---|
148 | fw_tof = numpy.zeros(MAX_PIXEL, dtype='float') # an array for the min_tof calculated for all pixels |
---|
149 | sqrtdmin = numpy.zeros(MAX_PIXEL, dtype='float') # sqrt(d_min) |
---|
150 | min_tofA = numpy.zeros(MAX_PIXEL, dtype='float') # A = array |
---|
151 | |
---|
152 | # two_sin_theta = 0.003 + (0.60507e-5/dspmap) + (0.00973e-10*dspmap**-2) |
---|
153 | usperAr = ((10.0*dspmap*(0.003 + (0.60507e-5/dspmap) + (0.00973e-10*dspmap**-2)))**-1)/min_usperA |
---|
154 | |
---|
155 | |
---|
156 | n_events = 0 |
---|
157 | n_events2fix_t = 0 |
---|
158 | n_events2fix_r = 0 |
---|
159 | |
---|
160 | first_buf = True |
---|
161 | events_inp = open(orig_file, 'rb') |
---|
162 | events_out = open(event_file, 'wb') |
---|
163 | while True: |
---|
164 | events_buf = events_inp.read(DataChunk) |
---|
165 | if not events_buf: |
---|
166 | break |
---|
167 | log_msg('-- read %.1f MB of data' % (len(events_buf)/1024.0/1024.0)) |
---|
168 | events = numpy.fromstring(events_buf, dtype=NeutronEvent) |
---|
169 | |
---|
170 | pix = events['pix'] |
---|
171 | tof = events['tof'] |
---|
172 | |
---|
173 | n_events += numpy.size(pix) # number of events |
---|
174 | |
---|
175 | if first_buf: # calculate things only on the first buffer |
---|
176 | min_tof = min(tof[MIN_EVENTS/2:MIN_EVENTS]) |
---|
177 | max_tof = max(tof[MIN_EVENTS/2:MIN_EVENTS]) |
---|
178 | d_tof = max_tof - min_tof |
---|
179 | time_offset = int(sns_frame*int(d_tof/sns_frame+0.49)) # get the time offset |
---|
180 | log_msg('-- %f < TOFs < %f & Delta_TOF = %f' %(0.1*min_tof,0.1*max_tof,0.1*d_tof)) |
---|
181 | log_msg('-- using time offset of %f usec' % time_offset) |
---|
182 | |
---|
183 | # calculate the min_tof for each pixel |
---|
184 | fw_tof = (usperAr * max_tof *0.99284) - d_tof #0.99284 is because end of lambda band is not at end of TOF band |
---|
185 | |
---|
186 | # calculate the min_tof for each pixel |
---|
187 | sqrtdminr = (0.1*min_tof/DifCr)**0.5 #reference point |
---|
188 | |
---|
189 | sqrtdmin = sqrtdminr + (K*(LDifCr-numpy.log10(0.1/dspmap))) #expression for calculating d_min allowed for all pixels |
---|
190 | |
---|
191 | # calculate min TOF for each pixel |
---|
192 | min_tofA = sqrtdmin/dspmap |
---|
193 | min_tofA *= (min_tofA > 0.0) |
---|
194 | |
---|
195 | min_tofA *= sqrtdmin |
---|
196 | |
---|
197 | if verbose: |
---|
198 | log_msg( "-- usperAr: %s" % usperAr[230000:230050] ) |
---|
199 | log_msg( "-- usperAr: %s" % usperAr[55001:55005] ) |
---|
200 | log_msg( "-- usperAr: %s" % usperAr[231051:231055] ) |
---|
201 | |
---|
202 | log_msg( "-- fw_tof: %s" % (0.1*fw_tof[55001:55005] )) |
---|
203 | log_msg( "-- fw_tof: %s" % (0.1*fw_tof[231051:231055])) |
---|
204 | |
---|
205 | log_msg( "-- min_tofA: %s" % (0.1*min_tofA[55001:55005] )) |
---|
206 | log_msg( "-- min_tofA: %s" % (0.1*min_tofA[231051:231055] )) |
---|
207 | |
---|
208 | |
---|
209 | |
---|
210 | # correct any bad pixel ids in case thre are errors |
---|
211 | pix = pix * (pix < PIX_MASK) |
---|
212 | |
---|
213 | #create mask: events2fix |
---|
214 | events2fix_t = (tof < fw_tof[pix]) |
---|
215 | events2fix_r = (tof < min_tofA[pix]) |
---|
216 | n_events2fix_t += events2fix_t.sum() |
---|
217 | n_events2fix_r += events2fix_r.sum() |
---|
218 | |
---|
219 | # ---------------------------------------------------------------------------------------- |
---|
220 | # fixing of events |
---|
221 | tof_new = tof + events2fix_t * time_offset |
---|
222 | pix_new = pix | events2fix_t * TOF_FLAG | events2fix_r*RES_FLAG |
---|
223 | # ---------------------------------------------------------------------------------------- |
---|
224 | |
---|
225 | if do_plot and first_buf: |
---|
226 | do = 2*dspmap[pix]*tof # original d-plot |
---|
227 | df = 2*dspmap[pix]*tof_new # fixed d-plot |
---|
228 | histograms = [ |
---|
229 | ('Original' , do ), |
---|
230 | ('Bad TOF' , do[events2fix_t] ), |
---|
231 | ('Fixed' , df ), |
---|
232 | ('Bad TOF (fixed)', df[events2fix_t] ), |
---|
233 | ] |
---|
234 | plot_dspace(histograms , event_file) |
---|
235 | |
---|
236 | |
---|
237 | events['tof'] = tof_new |
---|
238 | events['pix'] = pix_new |
---|
239 | if first_buf: |
---|
240 | events[0]['pix'] = events[0]['pix'] | FIX_MASK |
---|
241 | first_buf = False |
---|
242 | if verbose: |
---|
243 | log_msg( "-- event[0]: PIX=0x%08X TOF=%d" % (events[0]['pix'], events[0]['tof'])) |
---|
244 | events.tofile(events_out) |
---|
245 | |
---|
246 | |
---|
247 | log_msg('-- %8d/%d events [%5.3f %%] are flagged with TOF_FLAG=0x%08x' % (n_events2fix_t, n_events, 100.0*n_events2fix_t/n_events, TOF_FLAG)) |
---|
248 | log_msg('-- %8d/%d events [%5.3f %%] are flagged with RES_FLAG=0x%08x' % (n_events2fix_r, n_events, 100.0*n_events2fix_r/n_events, RES_FLAG)) |
---|
249 | log_msg('-- event file %s is fixed' % event_file) |
---|
250 | events_out.close() |
---|
251 | events_inp.close() |
---|
252 | |
---|
253 | |
---|
254 | |
---|
255 | |
---|
256 | if __name__ == "__main__": |
---|
257 | import optparse |
---|
258 | |
---|
259 | usage = "%prog [options] run_number" #% argv0 |
---|
260 | version = "%%prog %s" % __version__ |
---|
261 | |
---|
262 | opt = optparse.OptionParser(usage=usage, version=version) |
---|
263 | |
---|
264 | opt.add_option("-D", "--dspace", dest="dspace" , default=DSPACEMAP) |
---|
265 | opt.add_option("-d", "--dir" , dest="top_dir", default='.') |
---|
266 | opt.add_option("-p", "--plot", |
---|
267 | action="store_true", dest="plot", default=False, |
---|
268 | help="plot diagnostic plots") |
---|
269 | opt.add_option("-v", "--verbose", |
---|
270 | action="store_true", dest="verbose", default=False, |
---|
271 | help="don't print status messages") |
---|
272 | |
---|
273 | opt.add_option("-l", "--log", dest="logfile", |
---|
274 | help="print messages to a log file") |
---|
275 | opt.add_option("-f", "--force", action="store_true", dest="force" , default=False) |
---|
276 | |
---|
277 | |
---|
278 | |
---|
279 | options, args = opt.parse_args() |
---|
280 | if options.logfile: |
---|
281 | fix_log = open(options.logfile, 'at+') |
---|
282 | |
---|
283 | if options.plot: |
---|
284 | import pylab |
---|
285 | |
---|
286 | if not args: |
---|
287 | args = [''] |
---|
288 | for arg in args: |
---|
289 | file_type="*%s*neutron_event.dat" % arg |
---|
290 | for file_name in find_prenexus(top_dir=options.top_dir, file_type=file_type): |
---|
291 | fix_file_loop(file_name, dspace_file=options.dspace, |
---|
292 | verbose=options.verbose, do_plot=options.plot, |
---|
293 | do_force=options.force) |
---|
294 | if options.plot: |
---|
295 | pylab.show() |
---|