added beginnings of script to save waveforms to disk to tektronix.py.

1 files changed, 134 insertions, 16 deletions

diff --git a/tektronix.py b/tektronix.py
index a644a05..d8129d6 100644
--- a/tektronix.py
+++ b/tektronix.py
@@ -20,6 +20,24 @@ preamble_fields = {'BYT_NR': int, # data width for waveform
                    'YUNIT' : str,
                    'NR_FR' : int }
 
+def get_dtype(preamble):
+    """Returns the numpy dtype for the raw waveform data given the preamble."""
+    if preamble['BYT_OR'] == 'MSB':
+        byteorder = '>'
+    elif preamble['BYT_OR'] == 'LSB':
+        byteorder = '<'
+    else:
+        raise Exception('unknown byte order %s' % preamble['BYT_OR'])
+
+    if preamble['BN_FMT'] == 'RI':
+        signedchar = 'i'
+    elif preamble['BN_FMT'] == 'RP':
+        signedchar = 'u'
+    else:
+        raise Exception('unknown binary format string %s' % preamble['BN_FMT'])
+
+    return byteorder + signedchar + str(preamble['BYT_NR'])
+
 def convert_waveform(waveform, preamble):
     """Converts a waveform returned by the scope into voltage values."""
     return preamble['YZERO'] + (waveform - preamble['YOFF'])*preamble['YMULT']
@@ -55,21 +73,35 @@ class TekScope(object):
 
         return self.recv()
 
-    def get_preamble(self):
+    def set_sequence_mode(self):
+        """Sets the oscilloscope in single acquisition mode."""
+        self.send('acquire:stopafter sequence\n')
+
+    def acquire(self):
+        """Trigger and acquire a new waveform."""
+        self.send('acquire:state run\n')
+
+    def get_preamble(self, channel):
         """Returns a dictionary containing information about the waveform
-        format."""
+        format for a channel."""
         header = self.query('header?\n')
 
         # turn header on so that we know preamble field names
         self.send('header 1\n')
 
+        source = self.query('data:source?')
+
+        self.send('data:source ch%i\n' % channel)
+
         preamble = {}
         for s in self.query('wfmpre?\n').strip()[8:].split(';'):
             key, value = s.split(' ',1)
 
             preamble[key] = preamble_fields[key](value)
 
+        # reset header format and data:source
         self.send(header)
+        self.send(source)
 
         return preamble
 
@@ -93,7 +125,10 @@ class TekScope(object):
 
         return channels
 
-    def get_waveform(self, channel):
+    def get_waveform(self, channel, dtype=None):
+        """Returns the waveform from channel as a numpy array. If dtype is
+        specified, the function does not need to query the scope for the data
+        format which will be much quicker."""
         header = self.query('header?\n')
 
         self.send('header 0\n')
@@ -102,22 +137,26 @@ class TekScope(object):
         # in pt_fmt y format
         self.send('wfmpre:pt_fmt y\n')
 
-        encoding = self.query('data:encdg?')[:3]
+        self.send('data:source ch%i\n' % channel)
 
-        if encoding == 'RIB':
-            dtype = '>i'
-        elif encoding == 'RPB':
-            dtype = '>u'
-        elif encoding == 'SRI':
-            dtype = '<i'
-        elif encoding == 'SRP':
-            dtype = '<u'
-        else:
-            raise ValueError('unkown encoding: %s' % encoding)
+        if dtype is None:
+            dtype = get_dtype(self.get_preamble(channel))
+            # encoding = self.query('data:encdg?')[:3]
+
+            # if encoding == 'RIB':
+            #     dtype = '>i'
+            # elif encoding == 'RPB':
+            #     dtype = '>u'
+            # elif encoding == 'SRI':
+            #     dtype = '<i'
+            # elif encoding == 'SRP':
+            #     dtype = '<u'
+            # else:
+            #     raise ValueError('unkown encoding: %s' % encoding)
+
+            # dtype += self.query('data:width?').strip()
 
-        dtype += self.query('data:width?').strip()
 
-        self.send('data:source ch%i\n' % channel)
 
         self.send('curve?\n')
 
@@ -140,3 +179,82 @@ class TekScope(object):
         self.send(header)
 
         return waveform
+
+if __name__ == '__main__':
+    import h5py
+    import optparse
+    import setup
+    import sys
+    import os
+    import math
+    import time
+
+    usage = 'usage: %prog <filename> [-n]'
+    parser = optparse.OptionParser(usage)
+    parser.add_option('-n', type='int', dest='nevents',
+                      help='number of events per run', default=1000)
+    parser.add_option('-r', type='int', dest='nruns',
+                      help='number of runs', default=1)
+    options, args = parser.parse_args()
+
+    if len(args) < 1:
+        sys.exit(parser.format_help())
+
+    if options.nevents < 1 or options.nruns < 1:
+        sys.exit('nevents and nruns must be greater than 1.')
+
+    scope = TekScope(setup.scope_ip, setup.port)
+
+    root, ext = os.path.splitext(args[0])
+
+    if ext == '':
+        ext = '.hdf5'
+
+    #print (root, ext)
+
+    for run in range(options.nruns):
+        if options.nruns > 1:
+            fileid = str(run).zfill(int(log10(options.nruns))+1)
+        else:
+            fileid = ''
+
+        filename = root + fileid + ext
+
+        #print filename
+
+        t0 = time.time()
+
+        with h5py.File(filename, 'w') as f:
+            f.attrs['settings'] = scope.query('*lrn?')
+
+            active_channels = scope.get_active_channels()
+
+            dtypes = {}
+            for channel in active_channels:
+                preamble = scope.get_preamble(channel)
+
+                #dtypes[channel] = get_dtype(preamble)
+
+                dataset = f.create_dataset('channel%i' % channel, (options.nevents, preamble['NR_PT']), dtype=get_dtype(preamble), chunks=(max(1,min(100, options.nevents//100)), preamble['NR_PT']), compression='gzip')
+
+                for key, value in preamble.iteritems():
+                    dataset.attrs[key] = value
+
+            # enable single acquisition mode
+            scope.set_sequence_mode()
+
+            for i in range(options.nevents):
+                print '\rsaving event: %i' % i,
+                sys.stdout.flush()
+
+                scope.acquire()
+
+                for channel in active_channels:
+                    dataset = f['channel%i' % channel]
+
+                    dataset[i] = scope.get_waveform(channel, dataset.dtype)
+            print
+
+        elapsed = time.time() - t0
+
+        print 'saved %s. elapsed %f sec.' % (filename,elapsed)