Update RP data collection/plotting scripts with argparse & asyncio

plot_sayma_data.py

@@ -1,19 +1,44 @@
 import numpy as np
 import matplotlib.pyplot as plot
 from scipy import signal, constants
+import argparse
+import os
+
 
 def rp_raw_to_numpy(rp_raw):
     # Convert raw buffer strings to numpy arrays
     buff_string = rp_raw.split(',')
     return np.array(list(map(float, buff_string)))
 
-def main():
-    y1_raw = None
-    y2_raw = None
 
-    with open('rp_y1_raw.bin', 'rb') as f:
+RP_IP_ADDRS = {
+    "creotech-1": "192.168.1.104",
+    "creotech-2": "192.168.1.105",
+    "mlabs": "rp-f05cc9",
+}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Data plotting tool for Sayma DAC/TTL")
+    parser.add_argument("dir", help="output directory", type=str)
+    parser.add_argument("rps", metavar="NAME:CHANNEL",
+        help="input <CHANNEL> of the RedPitaya at <NAME> where data is collected; "
+        "CHANNEL must be 1 or 2; "
+        "NAME must be any of: " + " ".join(list(RP_IP_ADDRS.keys())),
+        type=str, nargs=2)
+    args = parser.parse_args()
+
+    # Must only compare 2 data
+    name, channel = args.rps[0].split(':')
+    y1_filename = 'rp_{}_y{}_raw.bin'.format(name, channel)
+    name, channel = args.rps[1].split(':')
+    y2_filename = 'rp_{}_y{}_raw.bin'.format(name, channel)
+    if y1_filename == y2_filename:
+        raise ValueError("Both files are the same.")
+
+    with open(os.path.join(args.dir, y1_filename), 'rb') as f:
         y1_raw = f.read().decode('utf-8')
-    with open('rp_y2_raw.bin', 'rb') as f:
+    with open(os.path.join(args.dir, y2_filename), 'rb') as f:
         y2_raw = f.read().decode('utf-8')
     if None in [y1_raw, y2_raw]:
         raise IOError("Raw RP string files cannot be opened.")
@@ -21,9 +46,14 @@ def main():
     y1 = rp_raw_to_numpy(y1_raw)
     y2 = rp_raw_to_numpy(y2_raw)
 
+    # Note that RedPitaya's oscilloscope has a sampling rate @ 125MHz
     t = np.arange(y1.shape[0])/125e6
+    # Generate matrix Y by having arrays y1 and y2 as 2 rows
     y = np.c_[y1, y2].T
+    # Multiply an array of cos(2pi*9e6*t) with each row in Y;
+    # Then, downsample the array by 10 as Z.
     z = signal.decimate(y*np.exp(1j*2*np.pi*9e6*t), q=10, ftype="fir", zero_phase=True)[:, 10:]
+    # Downsample Z by 10 again.
     z = signal.decimate(z, q=10, ftype="fir", zero_phase=True)[:, 10:]
     angle = np.angle(np.mean(z[0]*z[1].conj()))
 

rp_get_sayma_data.py

@@ -1,10 +1,18 @@
 import socket
+from time import sleep
+import asyncio
+import argparse
+import os
 
 
 class RPSCPI:
-    def connect(self, host):
+    def __init__(self, rp_name, rp_host):
+        self.name = rp_name
+        self.host = rp_host
+
+    def connect(self):
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.sock.connect((host, 5000))
+        self.sock.connect((self.host, 5000))
         self.sock_f = self.sock.makefile()
 
     def close(self):
@@ -17,8 +25,11 @@ class RPSCPI:
     def recvmsg(self):
         return self.sock_f.readline().strip()
 
-    def trigger(self):
+    def trigger_prep(self):
+        self.sendmsg("ACQ:RST")
         self.sendmsg("ACQ:START")
+
+    async def trigger(self):
         self.sendmsg("ACQ:TRIG NOW")
         while True:
             self.sendmsg("ACQ:TRIG:STAT?")
@@ -29,35 +40,67 @@ class RPSCPI:
         self.sendmsg("ACQ:SOUR{}:DATA?".format(channel))
         return self.recvmsg()[1:-1]
 
+
+RP_IP_ADDRS = {
+    "creotech-1": "192.168.1.104",
+    "creotech-2": "192.168.1.105",
+    "mlabs": "rp-f05cc9",
+}
+
+
+async def gather_trigger(rp_list):
+    await asyncio.gather(*[rp.trigger() for rp in rp_list])
+
+
 def main():
-    rp = RPSCPI()
+    parser = argparse.ArgumentParser(description="Data collection tool for Sayma DAC/TTL at RedPitaya")
-    rp.connect("192.168.1.104")
+    parser.add_argument("dir", help="output directory", type=str)
+    parser.add_argument("rps", help="name(s) of the target RedPitayas where data is collected "
+        "simultaneously; any of: " + " ".join(list(RP_IP_ADDRS.keys())),
+        type=str, nargs='+')
+    parser.add_argument("--txt", help="save data as additional human-readable text files",
+        action="store_true")
+    args = parser.parse_args()
+
+    # Connect a socket to each RP
+    rps = []
+    for rp_name in args.rps:
+        rp = RPSCPI(rp_name, RP_IP_ADDRS[rp_name])
+        rp.connect()
+        rps.append(rp)
+
     try:
-        rp.trigger()
+        for rp in rps:
-        y1_raw = rp.get_data(1)
+            rp.trigger_prep()
-        y2_raw = rp.get_data(2)
 
-        with open('rp_y1_raw.bin', 'wb') as f:
+        asyncio.run(gather_trigger(rps))
-            f.write(y1_raw.encode('utf-8'))
-            print("Succesfully written y1 raw string from RP.")
-        with open('rp_y2_raw.bin', 'wb') as f:
-            f.write(y2_raw.encode('utf-8'))
-            print("Succesfully written y2 raw string from RP.")
 
-        # DEBUGGING
+        for rp in rps:
-        y1 = [float(i) for i in y1_raw.split(',')]
+            y1_raw = rp.get_data(1)
-        y2 = [float(i) for i in y2_raw.split(',')]
+            y2_raw = rp.get_data(2)
-        with open('rp_y1_raw.txt', 'w') as f:
+
-            for i in y1:
+            with open(os.path.join(args.dir, 'rp_{}_y1_raw.bin'.format(rp.name)), 'wb') as f:
-                f.write(str(i) + '\n')
+                f.write(y1_raw.encode('utf-8'))
-            print("[DEBUG] Succesfully written y1 human-readable data.")
+                print("Succesfully written y1 raw data from RP {}.".format(rp.name))
-        with open('rp_y2_raw.txt', 'w') as f:
+            with open(os.path.join(args.dir, 'rp_{}_y2_raw.bin'.format(rp.name)), 'wb') as f:
-            for i in y2:
+                f.write(y2_raw.encode('utf-8'))
-                f.write(str(i) + '\n')
+                print("Succesfully written y2 raw data from RP {}.".format(rp.name))
-            print("[DEBUG] Succesfully written y2 human-readable data.")
+
+            if args.txt:
+                y1 = [float(i) for i in y1_raw.split(',')]
+                y2 = [float(i) for i in y2_raw.split(',')]
+                with open(os.path.join(args.dir, 'rp_{}_y1_raw.txt'.format(rp.name)), 'w') as f:
+                    for i in y1:
+                        f.write(str(i) + '\n')
+                    print("Succesfully written y1 human-readable data from {}.".format(rp.name))
+                with open(os.path.join(args.dir, 'rp_{}_y2_raw.txt'.format(rp.name)), 'w') as f:
+                    for i in y2:
+                        f.write(str(i) + '\n')
+                    print("Succesfully written y2 human-readable data from {}.".format(rp.name))
 
     finally:
-        rp.close()
+        for rp in rps:
+            rp.close()
 
 if __name__ == "__main__":
     main()