UserContributedScripts: dat2wav.py

File dat2wav.py, 12.9 KB (added by jayce, 4 years ago)

Convert a DRX file to a wav audio file.

Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4"""
5Given a DRX file,  create a wav audio file of its contents
6
7$Rev: 1375 $
8$LastChangedBy: jdowell $
9$LastChangedDate: 2013-07-11 10:26:08 -0600 (Thu, 11 Jul 2013) $
10"""
11
12import os
13import sys
14import h5py
15import math
16import numpy
17import ephem
18import getopt
19from datetime import datetime
20
21from scipy.io import wavfile
22from scipy.interpolate import interp1d
23
24import lsl.reader.drx as drx
25import lsl.reader.drspec as drspec
26import lsl.reader.errors as errors
27import lsl.statistics.robust as robust
28import lsl.correlator.fx as fxc
29from lsl.correlator import _core
30from lsl.astro import unix_to_utcjd, DJD_OFFSET
31from lsl.common import progress, stations
32from lsl.common import mcs, metabundle
33
34import matplotlib.pyplot as plt
35
36
37def usage(exitCode=None):
38        print """dat2wav.py - Read in a DRX file and create a wav audio file of its contents
39
40Usage: dat2wav.py [OPTIONS] file
41
42Options:
43-h, --help                  Display this help information
44-s, --skip                  Skip the specified number of seconds at the beginning
45                            of the file (default = 0)
46-a, --average               Number of seconds of data to average for spectra
47                            (default = 1)
48-d, --duration              Number of seconds to calculate the waterfall for
49                            (default = 0; run the entire file)
50-l, --fft-length            Set FFT length (default = 4096)
51-c, --circularize           Convert data to R/L before processing
52-b, --bandpass              Apply a crude bandpass correction to the data
53                            (default = no)
54-i, --itu-r468              Apply ITU-R 468 noise weighting (default = no)
55"""
56
57        if exitCode is not None:
58                sys.exit(exitCode)
59        else:
60                return True
61
62
63def parseOptions(args):
64        config = {}
65        # Command line flags - default values
66        config['offset'] = 0.0
67        config['average'] = 1.0
68        config['LFFT'] = 4096
69        config['freq1'] = 0
70        config['freq2'] = 0
71        config['maxFrames'] = 28000
72        config['duration'] = 0.0
73        config['circularize'] = False
74        config['bandpass'] = False
75        config['noiseWeight'] = False
76        config['args'] = []
77
78        # Read in and process the command line flags
79        try:
80                opts, args = getopt.getopt(args, "hql:s:a:d:l:cbi", ["help", "quiet", "fft-length=", "skip=", "average=", "duration=", "fft-length=", "circularize", "bandpass", "itu-r468"])
81        except getopt.GetoptError, err:
82                # Print help information and exit:
83                print str(err) # will print something like "option -a not recognized"
84                usage(exitCode=2)
85       
86        # Work through opts
87        for opt, value in opts:
88                if opt in ('-h', '--help'):
89                        usage(exitCode=0)
90                elif opt in ('-q', '--quiet'):
91                        config['verbose'] = False
92                elif opt in ('-l', '--fft-length'):
93                        config['LFFT'] = int(value)
94                elif opt in ('-s', '--skip'):
95                        config['offset'] = float(value)
96                elif opt in ('-a', '--average'):
97                        config['average'] = float(value)
98                elif opt in ('-d', '--duration'):
99                        config['duration'] = float(value)
100                elif opt in ('-c', '--circularize'):
101                        config['circularize'] = True
102                elif opt in ('-b', '--bandpass'):
103                        config['bandpass'] = True
104                elif opt in ('-i', '--itu-r468'):
105                        config['noiseWeight'] = True
106                else:
107                        assert False
108                       
109        # Add in arguments
110        config['args'] = args
111
112        # Return configuration
113        return config
114
115
116def mapFreqToAudio(freq, min=40, max=10000, log=False):
117        """
118        Given an array of observed frequencies, map them down to audio
119        frequencies between `min' and `max' either linearly (default)
120        or logarithmically.
121       
122        Return an array containing the mapped frequencies in Hz.
123        """
124       
125        freqPrime = 1.0*freq
126        if log:
127                freqPrime = numpy.log10(freqPrime)
128        freqPrime -= freqPrime.min()
129        freqPrime /= freqPrime.max()
130        freqPrime = freqPrime*(max-min) + min
131        return freqPrime
132
133
134def main(args):
135        # Parse command line options
136        config = parseOptions(args)
137
138        # Length of the FFT
139        LFFT = config['LFFT']
140
141        # Open the file and find good data (not spectreomter data)
142        filename = config['args'][0]
143        fh = open(filename, "rb")
144        nFramesFile = os.path.getsize(filename) / drx.FrameSize
145       
146        try:
147                for i in xrange(5):
148                        junkFrame = drspec.readFrame(fh)
149                raise RuntimeError("ERROR: '%s' appears to be a DR spectrometer file, not a raw DRX file" % filename)
150        except errors.syncError:
151                fh.seek(0)
152               
153        while True:
154                try:
155                        junkFrame = drx.readFrame(fh)
156                        try:
157                                srate = junkFrame.getSampleRate()
158                                t0 = junkFrame.getTime()
159                                break
160                        except ZeroDivisionError:
161                                pass
162                except errors.syncError:
163                        fh.seek(-drx.FrameSize+1, 1)
164                       
165        fh.seek(-drx.FrameSize, 1)
166       
167        beam,tune,pol = junkFrame.parseID()
168        beams = drx.getBeamCount(fh)
169        tunepols = drx.getFramesPerObs(fh)
170        tunepol = tunepols[0] + tunepols[1] + tunepols[2] + tunepols[3]
171        beampols = tunepol
172
173        # Offset in frames for beampols beam/tuning/pol. sets
174        offset = int(config['offset'] * srate / 4096 * beampols)
175        offset = int(1.0 * offset / beampols) * beampols
176        fh.seek(offset*drx.FrameSize, 1)
177       
178        # Iterate on the offsets until we reach the right point in the file.  This
179        # is needed to deal with files that start with only one tuning and/or a
180        # different sample rate. 
181        while True:
182                ## Figure out where in the file we are and what the current tuning/sample
183                ## rate is
184                junkFrame = drx.readFrame(fh)
185                srate = junkFrame.getSampleRate()
186                t1 = junkFrame.getTime()
187                tunepols = drx.getFramesPerObs(fh)
188                tunepol = tunepols[0] + tunepols[1] + tunepols[2] + tunepols[3]
189                beampols = tunepol
190                fh.seek(-drx.FrameSize, 1)
191               
192                ## See how far off the current frame is from the target
193                tDiff = t1 - (t0 + config['offset'])
194               
195                ## Half that to come up with a new seek parameter
196                tCorr = -tDiff / 2.0
197                cOffset = int(tCorr * srate / 4096 * beampols)
198                cOffset = int(1.0 * cOffset / beampols) * beampols
199                offset += cOffset
200               
201                ## If the offset is zero, we are done.  Otherwise, apply the offset
202                ## and check the location in the file again/
203                if cOffset is 0:
204                        break
205                fh.seek(cOffset*drx.FrameSize, 1)
206       
207        # Update the offset actually used
208        config['offset'] = t1 - t0
209        offset = int(round(config['offset'] * srate / 4096 * beampols))
210        offset = int(1.0 * offset / beampols) * beampols
211
212        # Make sure that the file chunk size contains is an integer multiple
213        # of the FFT length so that no data gets dropped.  This needs to
214        # take into account the number of beampols in the data, the FFT length,
215        # and the number of samples per frame.
216        maxFrames = int(1.0*config['maxFrames']/beampols*4096/float(LFFT))*LFFT/4096*beampols
217
218        # Number of frames to integrate over
219        nFramesAvg = int(config['average'] * srate / 4096 * beampols)
220        nFramesAvg = int(1.0 * nFramesAvg / beampols*4096/float(LFFT))*LFFT/4096*beampols
221        config['average'] = 1.0 * nFramesAvg / beampols * 4096 / srate
222        maxFrames = nFramesAvg
223
224        # Number of remaining chunks (and the correction to the number of
225        # frames to read in).
226        if config['duration'] == 0:
227                config['duration'] = 1.0 * nFramesFile / beampols * 4096 / srate
228        else:
229                config['duration'] = int(round(config['duration'] * srate * beampols / 4096) / beampols * 4096 / srate)
230        nChunks = int(round(config['duration'] / config['average']))
231        if nChunks == 0:
232                nChunks = 1
233        nFrames = nFramesAvg*nChunks
234       
235        # Date & Central Frequnecy
236        t1  = junkFrame.getTime()
237        beginDate = ephem.Date(unix_to_utcjd(junkFrame.getTime()) - DJD_OFFSET)
238        centralFreq1 = 0.0
239        centralFreq2 = 0.0
240        for i in xrange(4):
241                junkFrame = drx.readFrame(fh)
242                b,t,p = junkFrame.parseID()
243                if p == 0 and t == 1:
244                        try:
245                                centralFreq1 = junkFrame.getCentralFreq()
246                        except AttributeError:
247                                from lsl.common.dp import fS
248                                centralFreq1 = fS * ((junkFrame.data.flags>>32) & (2**32-1)) / 2**32
249                elif p == 0 and t == 2:
250                        try:
251                                centralFreq2 = junkFrame.getCentralFreq()
252                        except AttributeError:
253                                from lsl.common.dp import fS
254                                centralFreq2 = fS * ((junkFrame.data.flags>>32) & (2**32-1)) / 2**32
255                else:
256                        pass
257        fh.seek(-4*drx.FrameSize, 1)
258       
259        config['freq1'] = centralFreq1
260        config['freq2'] = centralFreq2
261
262        # File summary
263        print "Filename: %s" % filename
264        print "Date of First Frame: %s" % str(beginDate)
265        print "Beams: %i" % beams
266        print "Tune/Pols: %i %i %i %i" % tunepols
267        print "Sample Rate: %i Hz" % srate
268        print "Tuning Frequency: %.3f Hz (1); %.3f Hz (2)" % (centralFreq1, centralFreq2)
269        print "Frames: %i (%.3f s)" % (nFramesFile, 1.0 * nFramesFile / beampols * 4096 / srate)
270        print "---"
271        print "Offset: %.3f s (%i frames)" % (config['offset'], offset)
272        print "Integration: %.3f s (%i frames; %i frames per beam/tune/pol)" % (config['average'], nFramesAvg, nFramesAvg / beampols)
273        print "Duration: %.3f s (%i frames; %i frames per beam/tune/pol)" % (config['average']*nChunks, nFrames, nFrames / beampols)
274        print "Chunks: %i" % nChunks
275        print " "
276               
277        # Bandpass estimate (if requested)
278        bp = numpy.zeros((4,LFFT))
279        if config['bandpass']:
280                for i in xrange(nChunks/10):
281                        print "Starting bandpass chunk %i of %i" % (i+1, nChunks/10)
282                        data = numpy.zeros((4,4096*nFramesAvg/beampols), dtype=numpy.complex64)
283                        count = [0,0,0,0]
284                       
285                        for j in xrange(nFramesAvg):
286                                cFrame = drx.readFrame(fh)
287                                bm,tn,pl = cFrame.parseID()
288                                aStand = 2*(tn-1) + pl
289                               
290                                data[aStand, count[aStand]*4096:(count[aStand]+1)*4096] = cFrame.data.iq
291                                count[aStand] += 1
292                               
293                        ## Circularize, if needed
294                        if config['circularize']:
295                                l1 = (data[0,:] + j*data[1,:]) / numpy.sqrt(2)
296                                r1 = (data[0,:] - j*data[1,:]) / numpy.sqrt(2)
297                                l2 = (data[2,:] + j*data[3,:]) / numpy.sqrt(2)
298                                r2 = (data[2,:] - j*data[3,:]) / numpy.sqrt(2)
299                               
300                                data[0,:] = l1
301                                data[1,:] = r1
302                                data[2,:] = l2
303                                data[3,:] = r2
304                               
305                        junk, spec = fxc.SpecMaster(data, LFFT=LFFT)
306                        bp += spec
307                bp = numpy.sqrt(bp)
308                bp[0,:] /= bp[0,:].mean()
309                bp[1,:] /= bp[1,:].mean()
310                bp[2,:] /= bp[2,:].mean()
311                bp[3,:] /= bp[3,:].mean()
312        else:
313                bp += 1
314               
315        # Define variables to help setup, calculate, and save the audio output
316        t = numpy.arange(2*LFFT, dtype=numpy.float64) / 44100
317        s = numpy.zeros((4,t.size), dtype=numpy.float64)
318        output = numpy.zeros((2*LFFT*nChunks,4), dtype=numpy.int16)
319        outputCount = 0
320        outputPeak = None
321       
322        # Define variables to help calculate the complex spectra from the data
323        freq = numpy.fft.fftshift( numpy.fft.fftfreq(LFFT, d=1.0/srate) )
324        delay = numpy.zeros((4,freq.size))
325       
326        # Define the mapped audio frequeciens
327        freqPrime = numpy.zeros((4,freq.size), dtype=freq.dtype)
328        freqPrime[0,:] = mapFreqToAudio(freq+centralFreq1)
329        freqPrime[1,:] = mapFreqToAudio(freq+centralFreq1)
330        freqPrime[2,:] = mapFreqToAudio(freq+centralFreq2)
331        freqPrime[3,:] = mapFreqToAudio(freq+centralFreq2)
332       
333        # Calculate the ITU-R 486 weighting, if requested
334        weighting = numpy.zeros_like(freqPrime)
335        if config['noiseWeight']:
336                ituR468Data = numpy.array([[   31.5, -29.9], 
337                                                          [   63.0, -23.9], 
338                                                          [  100.0, -19.8], 
339                                                          [  200.0, -13.8], 
340                                                          [  400.0,  -7.8], 
341                                                          [  800.0,  -1.9], 
342                                                          [ 1000.0,   0.0],
343                                                          [ 2000.0,   5.6], 
344                                                          [ 3150.0,   9.0], 
345                                                          [ 4000.0,  10.5],
346                                                          [ 5000.0,  11.7], 
347                                                          [ 6300.0,  12.2], 
348                                                          [ 7100.0,  12.0],
349                                                          [ 8000.0,  11.4], 
350                                                          [ 9000.0,  10.1], 
351                                                          [10000.0,   8.1],
352                                                          [12500.0,   0.0],
353                                                          [14000.0,  -5.3],
354                                                          [16000.0, -11.7],
355                                                          [20000.0, -22.2],
356                                                          [31500.0, -42.7]])
357                ituR468Interp = interp1d(ituR468Data[:,0], ituR468Data[:,1])
358               
359                for i in xrange(4):
360                        weighting[i,:] = 10**(ituR468Interp(freqPrime[i,:])/10.0)
361                       
362        else:
363                weighting += 1
364       
365        for i in xrange(nChunks):
366                print "Starting audio chunk %i of %i" % (i+1, nChunks)
367                data = numpy.zeros((4,4096*nFramesAvg/beampols), dtype=numpy.complex64)
368                count = [0,0,0,0]
369               
370                ## Read
371                for j in xrange(nFramesAvg):
372                        cFrame = drx.readFrame(fh)
373                        bm,tn,pl = cFrame.parseID()
374                        aStand = 2*(tn-1) + pl
375                       
376                        data[aStand, count[aStand]*4096:(count[aStand]+1)*4096] = cFrame.data.iq
377                        count[aStand] += 1
378                       
379                ## Circularize, if needed
380                if config['circularize']:
381                        l1 = (data[0,:] + j*data[1,:]) / numpy.sqrt(2)
382                        r1 = (data[0,:] - j*data[1,:]) / numpy.sqrt(2)
383                        l2 = (data[2,:] + j*data[3,:]) / numpy.sqrt(2)
384                        r2 = (data[2,:] - j*data[3,:]) / numpy.sqrt(2)
385                       
386                        data[0,:] = l1
387                        data[1,:] = r1
388                        data[2,:] = l2
389                        data[3,:] = r2
390                       
391                ## Integrated complex spectra
392                camps, valid = _core.FEngineC2(data, freq+centralFreq1, delay, LFFT=LFFT, Overlap=1, SampleRate=srate, ClipLevel=0)
393                camps = camps.mean(axis=2)
394               
395                ## Bandpass correct
396                camps /= bp
397               
398                ## Build the audio signal
399                s *= 0
400                for j in xrange(4):
401                        for c,w,f in zip(camps[j,:], weighting[j,:], freqPrime[j,:]):
402                                a = numpy.abs(c)
403                                p = numpy.angle(c)
404                               
405                                s[j,:] += w*a*numpy.cos(2*numpy.pi*f*t + p)
406                               
407                ## Scale
408                if outputPeak is None:
409                        good = numpy.where( numpy.isfinite(s) )
410                        outputPeak = 4.0*(numpy.abs(s[good])).max()
411                s /= outputPeak
412                s[numpy.where( s < -1)] = -1
413                s[numpy.where( s >  1)] =  1
414                s *= 2**15-1
415               
416                ## Save to the output audio array
417                output[outputCount*2*LFFT:(outputCount+1)*2*LFFT,:] = s.T
418                outputCount += 1
419               
420        # Write out the audio files
421        basename = os.path.split(filename)[1]
422        basename = os.path.splitext(basename)[0]
423        wavfile.write('%s_tuning1.wav' % basename, 44100, output[:,:2])
424        wavfile.write('%s_tuning2.wav' % basename, 44100, output[:,2:])
425
426
427if __name__ == "__main__":
428        main(sys.argv[1:])
429