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", "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 ; " "e.g. mlabs:LV,HV means LV for IN1 and HV for IN2 on RP mlabs;\n" " 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: :,") # 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: 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()