forked from harry/creotech-sayma-testsuite
107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
import socket
|
|
from time import sleep
|
|
import asyncio
|
|
import argparse
|
|
import os
|
|
|
|
|
|
class RPSCPI:
|
|
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((self.host, 5000))
|
|
self.sock_f = self.sock.makefile()
|
|
|
|
def close(self):
|
|
self.sock_f.close()
|
|
self.sock.close()
|
|
|
|
def sendmsg(self, msg):
|
|
self.sock.send(msg.encode() + b"\r\n")
|
|
|
|
def recvmsg(self):
|
|
return self.sock_f.readline().strip()
|
|
|
|
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?")
|
|
if self.recvmsg() == "TD":
|
|
break
|
|
|
|
def get_data(self, channel):
|
|
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():
|
|
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:
|
|
for rp in rps:
|
|
rp.trigger_prep()
|
|
|
|
asyncio.run(gather_trigger(rps))
|
|
|
|
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:
|
|
for rp in rps:
|
|
rp.close()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|