add ability to fetch waveforms in chunks when fastframe is turned on.

1 files changed, 43 insertions, 6 deletions

diff --git a/tektronix.py b/tektronix.py
index d8129d6..3af4e5c 100644
--- a/tektronix.py
+++ b/tektronix.py
@@ -51,6 +51,9 @@ class TekScope(object):
         self.sock.connect((host,port))
 
     def send(self, msg):
+        if not msg.endswith('\n'):
+            msg += '\n'
+
         self.sock.sendall(msg)
 
     def recv(self):
@@ -71,7 +74,7 @@ class TekScope(object):
 
         self.send(msg)
 
-        return self.recv()
+        return self.recv().strip()
 
     def set_sequence_mode(self):
         """Sets the oscilloscope in single acquisition mode."""
@@ -171,10 +174,22 @@ class TekScope(object):
 
         y = int(self.sock.recv(int(x[1])))
 
-        waveform = np.fromstring(self.sock.recv(y),dtype)
+        buffer = ''
+
+        while len(buffer) < y:
+            buffer += self.sock.recv(y-len(buffer))
+
+        waveform = np.fromstring(buffer,dtype)
 
         # messages end with a newline
-        assert self.sock.recv(1024) == '\n'
+        eom = self.sock.recv(1024)
+
+        if eom != '\n':
+            print 'x = ', x
+            print 'y = ', y
+            print eom
+
+        assert eom == '\n'
 
         self.send(header)
 
@@ -243,16 +258,38 @@ if __name__ == '__main__':
             # enable single acquisition mode
             scope.set_sequence_mode()
 
-            for i in range(options.nevents):
-                print '\rsaving event: %i' % i,
+            scope.send('header 0\n')
+
+            fastframe_state = int(scope.query('horizontal:fastframe:state?'))
+            fastframe_count = int(scope.query('horizontal:fastframe:count?'))
+
+            #print 'fastframe_state = ', fastframe_state
+
+            i = 0
+            while i < options.nevents:
+                print '\rsaving event: %i' % (i+1),
                 sys.stdout.flush()
 
                 scope.acquire()
 
+                if fastframe_state:
+                    n = min(fastframe_count,options.nevents-i)
+                else:
+                    n = 1
+
                 for channel in active_channels:
                     dataset = f['channel%i' % channel]
 
-                    dataset[i] = scope.get_waveform(channel, dataset.dtype)
+                    data = scope.get_waveform(channel, dataset.dtype)
+
+                    if fastframe_state:
+                        data = data.reshape((fastframe_count,-1))
+                        dataset[i:i+n] = data[:n]
+                    else:
+                        dataset[i] = data
+
+                i += n
+
             print
 
         elapsed = time.time() - t0