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 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()))

View File

@ -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()
rp.connect("192.168.1.104")
parser = argparse.ArgumentParser(description="Data collection tool for Sayma DAC/TTL at RedPitaya")
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()
y1_raw = rp.get_data(1)
y2_raw = rp.get_data(2)
for rp in rps:
rp.trigger_prep()
with open('rp_y1_raw.bin', 'wb') as f:
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.")
asyncio.run(gather_trigger(rps))
# DEBUGGING
y1 = [float(i) for i in y1_raw.split(',')]
y2 = [float(i) for i in y2_raw.split(',')]
with open('rp_y1_raw.txt', 'w') as f:
for i in y1:
f.write(str(i) + '\n')
print("[DEBUG] Succesfully written y1 human-readable data.")
with open('rp_y2_raw.txt', 'w') as f:
for i in y2:
f.write(str(i) + '\n')
print("[DEBUG] Succesfully written y2 human-readable data.")
for rp in rps:
y1_raw = rp.get_data(1)
y2_raw = rp.get_data(2)
with open(os.path.join(args.dir, 'rp_{}_y1_raw.bin'.format(rp.name)), 'wb') as f:
f.write(y1_raw.encode('utf-8'))
print("Succesfully written y1 raw data from RP {}.".format(rp.name))
with open(os.path.join(args.dir, 'rp_{}_y2_raw.bin'.format(rp.name)), 'wb') as f:
f.write(y2_raw.encode('utf-8'))
print("Succesfully written y2 raw data from RP {}.".format(rp.name))
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()