Ticket #2502: pg_fix.py

File pg_fix.py, 9.0 KB (added by Peter Peterson, 10 years ago)

The actual program

Line 
1#!/bin/env python
2"""
3Fix frame wrapping events in powgen
4
5"""
6import os, sys, struct
7import numpy, math, fnmatch
8from shutil import move, copy2
9
10
11__version__  = "0.3"
12
13INSTRUMENT   = 'PG3'
14TOP_DIR      = '.'
15DSPACEMAP    = 'powgen_dspacemap_d1370_2010_09_12.dat'
16
17PIX_MASK     = 0x000FFFFF  # PG3 "good" pixels 
18FIX_MASK     = 0x00F00000  # PG3 mask for fixed pixels
19
20TOF_FLAG     = 0x00100000  #
21RES_FLAG     = 0x00200000  #
22
23MAX_PIXEL    = 300000      #
24
25SNS_FREQ     = 60.0        # 60 Hz
26
27MIN_EVENTS   = 4000        # need at least this many neutrons
28
29NeutronEvent = numpy.dtype([('tof','uint32'),('pix','uint32')]) # event data format
30DataChunk    = NeutronEvent.itemsize*1024L*1024L*24L            # read 24 Mevents at once, otherwise a small DFS might choke
31
32
33fix_log = sys.stdout
34
35def 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
48def 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
56def 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
76def 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
110def 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
256if __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()