add method to fetch waveforms. also add methods to convert waveform to voltage values and build a time array for the waveform.

2 files changed, 81 insertions, 1 deletions

diff --git a/setup.py b/setup.py
index 0c153ca..4a8cd3b 100644
--- a/setup.py
+++ b/setup.py
@@ -1 +1 @@
-scope_ip = ''
+scope_ip = '205.208.56.226'
diff --git a/tektronix.py b/tektronix.py
index e473d7e..a644a05 100644
--- a/tektronix.py
+++ b/tektronix.py
@@ -1,4 +1,6 @@
 import socket
+import re
+import numpy as np
 
 preamble_fields = {'BYT_NR': int, # data width for waveform
                    'BIT_NR': int, # number of bits per waveform point
@@ -18,6 +20,13 @@ preamble_fields = {'BYT_NR': int, # data width for waveform
                    'YUNIT' : str,
                    'NR_FR' : int }
 
+def convert_waveform(waveform, preamble):
+    """Converts a waveform returned by the scope into voltage values."""
+    return preamble['YZERO'] + (waveform - preamble['YOFF'])*preamble['YMULT']
+
+def build_time_array(preamble):
+    return preamble['XZERO'] + (np.arange(preamble['NR_PT']) - preamble['PT_OFF'])*preamble['XINCR']
+
 class TekScope(object):
     def __init__(self, host, port):
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -38,6 +47,7 @@ class TekScope(object):
         return buffer
 
     def query(self, msg):
+        """Sends a query to the oscilloscope and returns the response."""
         if not msg.endswith('\n'):
             msg += '\n'
 
@@ -46,6 +56,8 @@ class TekScope(object):
         return self.recv()
 
     def get_preamble(self):
+        """Returns a dictionary containing information about the waveform
+        format."""
         header = self.query('header?\n')
 
         # turn header on so that we know preamble field names
@@ -60,3 +72,71 @@ class TekScope(object):
         self.send(header)
 
         return preamble
+
+    def get_active_channels(self):
+        """Returns a list of the active (displayed) channel numbers."""
+        header = self.query('header?\n')
+
+        self.send('header 1\n')
+
+        channels = []
+        for s in self.query('select?').strip()[8:].split(';'):
+            m = re.match('CH(\d) (\d)', s)
+
+            if m is not None:
+                ch, state = map(int,m.groups())
+
+                if state != 0:
+                    channels.append(ch)
+
+        self.send(header)
+
+        return channels
+
+    def get_waveform(self, channel):
+        header = self.query('header?\n')
+
+        self.send('header 0\n')
+
+        # not sure what pt_fmt env is, so we'll always transmit
+        # in pt_fmt y format
+        self.send('wfmpre:pt_fmt y\n')
+
+        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()
+
+        self.send('data:source ch%i\n' % channel)
+
+        self.send('curve?\n')
+
+        # the initial response from curve looks like '#x<y>' where x is
+        # the number of y characters, and y is the number of bytes in the
+        # waveform.
+        # for example: '#41000' at the beginning of a curve? response means
+        # that there are 1000 bytes in the waveform
+        x = self.sock.recv(2)
+
+        assert x.startswith('#')
+
+        y = int(self.sock.recv(int(x[1])))
+
+        waveform = np.fromstring(self.sock.recv(y),dtype)
+
+        # messages end with a newline
+        assert self.sock.recv(1024) == '\n'
+
+        self.send(header)
+
+        return waveform