creotech-sayma-testsuite/rp_get_sayma_data.py

146 lines
5.2 KiB
Python
Raw Normal View History

import socket
import asyncio
import argparse
import os
import datetime
class RPSCPI:
def __init__(self, rp_name, rp_host, ch1_gain, ch2_gain):
self.name = rp_name
self.host = rp_host
self.ch1_gain = ch1_gain
self.ch2_gain = ch2_gain
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:SOUR1:GAIN {}".format(self.ch1_gain.upper()))
self.sendmsg("ACQ:SOUR2:GAIN {}".format(self.ch2_gain.upper()))
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",
2021-05-13 10:27:16 +08:00
"mlabs-1": "rp-f05cc9",
"mlabs-2": "rp-f0612e",
}
async def gather_trigger(rp_list):
await asyncio.gather(*[rp.trigger() for rp in rp_list])
def main():
# Get timestamp at script start
now = datetime.datetime.now().replace(
microsecond=0, # get rid of microseconds when printed
)
now_iso = now.isoformat(sep=' ')
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_gains", metavar="RP_NAME:CH1_GAIN,CH2_GAIN",
help="choose LV (+-1V) gain or HV (+-20V) gain on each channel for RP named <RP_NAME>; "
"e.g. mlabs:LV,HV means LV for IN1 and HV for IN2 on RP mlabs;\n"
"<RP_NAME> must be one 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()
gains = dict()
try:
for gain in args.rps_gains:
rp_name, ch1ch2_gains = gain.split(':')
ch1_gain, ch2_gain = ch1ch2_gains.split(',')
gains[rp_name] = ch1_gain, ch2_gain
if rp_name == "" or ch1_gain == "" or ch2_gain == "":
raise SyntaxError("Must specify IN1 and IN2 gains for each target RP.")
except ValueError:
raise SyntaxError("Must specify input gains in the format of: <RP_NAME>:<CH1_GAIN>,<CH2_GAIN>")
# Connect a socket to each RP
rps = []
for rp_name in gains:
rp = RPSCPI(rp_name, RP_IP_ADDRS[rp_name], *gains[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)
print("Data collection started at {}.".format(now_iso))
with open(os.path.join(args.dir, 'rp_{}_y1_raw.bin'.format(rp.name)), 'wb') as f:
f.write("{}\n".format(now_iso).encode('utf-8'))
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("{}\n".format(now_iso).encode('utf-8'))
f.write(y2_raw.encode('utf-8'))
print("Succesfully written y2 raw data from RP {}.".format(rp.name))
if args.txt:
header = '''\
RedPitaya Oscilloscope Reading
RP Hostname: {} ({})
RP Channel: Input {} ({})
Date & Time: {}
------
'''
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:
f.write(''.join([l.lstrip(' ') for l in header.format(
RP_IP_ADDRS[rp.name], rp.name, 1, ch1_gain.upper(), now_iso).splitlines(True)]))
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:
f.write(''.join([l.lstrip(' ') for l in header.format(
RP_IP_ADDRS[rp.name], rp.name, 2, ch2_gain.upper(), now_iso).splitlines(True)]))
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()