Update RP data collection/plotting scripts with argparse & asyncio

This commit is contained in:
Harry Ho 2021-03-09 12:40:51 +08:00
parent fa0b190234
commit 45367ac52f
2 changed files with 104 additions and 31 deletions

View File

@ -1,19 +1,44 @@
import numpy as np import numpy as np
import matplotlib.pyplot as plot import matplotlib.pyplot as plot
from scipy import signal, constants from scipy import signal, constants
import argparse
import os
def rp_raw_to_numpy(rp_raw): def rp_raw_to_numpy(rp_raw):
# Convert raw buffer strings to numpy arrays # Convert raw buffer strings to numpy arrays
buff_string = rp_raw.split(',') buff_string = rp_raw.split(',')
return np.array(list(map(float, buff_string))) 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') 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') y2_raw = f.read().decode('utf-8')
if None in [y1_raw, y2_raw]: if None in [y1_raw, y2_raw]:
raise IOError("Raw RP string files cannot be opened.") raise IOError("Raw RP string files cannot be opened.")
@ -21,9 +46,14 @@ def main():
y1 = rp_raw_to_numpy(y1_raw) y1 = rp_raw_to_numpy(y1_raw)
y2 = rp_raw_to_numpy(y2_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 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 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:] 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:] z = signal.decimate(z, q=10, ftype="fir", zero_phase=True)[:, 10:]
angle = np.angle(np.mean(z[0]*z[1].conj())) angle = np.angle(np.mean(z[0]*z[1].conj()))

View File

@ -1,10 +1,18 @@
import socket import socket
from time import sleep
import asyncio
import argparse
import os
class RPSCPI: 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 = 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() self.sock_f = self.sock.makefile()
def close(self): def close(self):
@ -17,8 +25,11 @@ class RPSCPI:
def recvmsg(self): def recvmsg(self):
return self.sock_f.readline().strip() return self.sock_f.readline().strip()
def trigger(self): def trigger_prep(self):
self.sendmsg("ACQ:RST")
self.sendmsg("ACQ:START") self.sendmsg("ACQ:START")
async def trigger(self):
self.sendmsg("ACQ:TRIG NOW") self.sendmsg("ACQ:TRIG NOW")
while True: while True:
self.sendmsg("ACQ:TRIG:STAT?") self.sendmsg("ACQ:TRIG:STAT?")
@ -29,34 +40,66 @@ class RPSCPI:
self.sendmsg("ACQ:SOUR{}:DATA?".format(channel)) self.sendmsg("ACQ:SOUR{}:DATA?".format(channel))
return self.recvmsg()[1:-1] 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(): 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: try:
rp.trigger() for rp in rps:
rp.trigger_prep()
asyncio.run(gather_trigger(rps))
for rp in rps:
y1_raw = rp.get_data(1) y1_raw = rp.get_data(1)
y2_raw = rp.get_data(2) y2_raw = rp.get_data(2)
with open('rp_y1_raw.bin', 'wb') as f: with open(os.path.join(args.dir, 'rp_{}_y1_raw.bin'.format(rp.name)), 'wb') as f:
f.write(y1_raw.encode('utf-8')) f.write(y1_raw.encode('utf-8'))
print("Succesfully written y1 raw string from RP.") print("Succesfully written y1 raw data from RP {}.".format(rp.name))
with open('rp_y2_raw.bin', 'wb') as f: with open(os.path.join(args.dir, 'rp_{}_y2_raw.bin'.format(rp.name)), 'wb') as f:
f.write(y2_raw.encode('utf-8')) f.write(y2_raw.encode('utf-8'))
print("Succesfully written y2 raw string from RP.") print("Succesfully written y2 raw data from RP {}.".format(rp.name))
# DEBUGGING if args.txt:
y1 = [float(i) for i in y1_raw.split(',')] y1 = [float(i) for i in y1_raw.split(',')]
y2 = [float(i) for i in y2_raw.split(',')] y2 = [float(i) for i in y2_raw.split(',')]
with open('rp_y1_raw.txt', 'w') as f: with open(os.path.join(args.dir, 'rp_{}_y1_raw.txt'.format(rp.name)), 'w') as f:
for i in y1: for i in y1:
f.write(str(i) + '\n') f.write(str(i) + '\n')
print("[DEBUG] Succesfully written y1 human-readable data.") print("Succesfully written y1 human-readable data from {}.".format(rp.name))
with open('rp_y2_raw.txt', 'w') as f: with open(os.path.join(args.dir, 'rp_{}_y2_raw.txt'.format(rp.name)), 'w') as f:
for i in y2: for i in y2:
f.write(str(i) + '\n') f.write(str(i) + '\n')
print("[DEBUG] Succesfully written y2 human-readable data.") print("Succesfully written y2 human-readable data from {}.".format(rp.name))
finally: finally:
for rp in rps:
rp.close() rp.close()
if __name__ == "__main__": if __name__ == "__main__":