quantnet-agent 1.0.0__py3-none-any.whl

quantnet_agent/hal/driver/artiq.py

@@ -0,0 +1,124 @@
+import logging
+from sipyco.pc_rpc import Client
+from sipyco.broadcast import Receiver
+from quantnet_agent.hal.hwclasses import ExpFramework
+import time
+import asyncio
+import numpy
+
+log = logging.getLogger(__name__)
+TIMEOUT = 10
+
+
+class ArtiqClient(ExpFramework):
+    def __init__(self, property, node_config, mqhost, mqport):
+        self.expid_to_rid = {}
+        self.logs = {}
+        if "version" in property and int(property["version"]) == 7:
+            self.schedule, self.exps, self.datasets = [
+            Client(property["host"], property["port"], i) for i in "master_schedule master_experiment_db master_dataset_db".split()
+        ]
+        else:
+            self.schedule, self.exps, self.datasets = [
+                Client(property["host"], property["port"], i) for i in "schedule experiment_db dataset_db".split()
+            ]
+
+        asyncio.create_task(self._connect_and_receive_logs(property["host"], int(property["logging"])))
+
+        log.info(f"Initializing ARTIQ with {property['host']} : {property['logging']}")
+
+    async def _connect_and_receive_logs(self, host, port):
+        # Create the receiver instance
+        log_receiver = Receiver("log", [], None)
+
+        # Connect to the log service at the specified address and port
+        await log_receiver.connect(host, port)
+
+        # Register the callback function to receive log messages
+        log_receiver.notify_cbs.append(self._append_message)
+
+        while True:
+            # This loop keeps running to listen for incoming messages
+            await asyncio.sleep(1)
+
+    def _append_message(self, msg):
+        log.debug(f"Received a message from artiq: {msg}")
+        try:
+            if msg[1] == 'master':
+                rid = "master"
+            else:
+                rid = int(msg[1][msg[1].find("(") + 1: msg[1].find(",")])
+
+            if rid not in self.logs:
+                self.logs[rid] = ""
+            self.logs[rid] = self.logs[rid] + msg[3]
+        except Exception as e:
+            log.error(f"Error reading ARTIQ log: {e}")
+            return
+
+    def reset(self):
+        log.info("Force deleting scheduled experiments")
+        for rid in self.schedule.get_status():
+            log.info(f"exp : {rid}")
+            self.schedule.delete(rid)
+
+    def stop(self):
+        log.info("Request for stopping experiments")
+        for rid in self.schedule.get_status():
+            log.info(f"exp : {rid}")
+            self.schedule.request_termination(rid)
+
+    def cleanUp(self):
+        self.stop()
+
+    def logs(self):
+        pass
+
+    async def receive(self, *args, **kwargs):
+        exp_id = args[0]        
+        rid = self.expid_to_rid[exp_id]
+        await self._wait_for_completion(rid)
+        if len(args) > 1:
+            params = args[1]
+            try:
+                results = {}
+                for param in params:
+                    result = self.datasets.get(param)
+                    if type(result) == numpy.ndarray:
+                        result = result.tolist()
+                    results[param] = result 
+                return {f"status for rid {rid}": "done", "results": results}
+            except Exception as e:
+                logging.error(f"Failed to get result for exp {exp_id}: {e}")
+                return {f"status for rid {rid}": "done"}
+        else:
+            return {f"status for rid {rid}": "done"}
+
+    async def submit(self, exp_id, expName, classname, args=dict()):
+        args['use_db'] = True
+        expid = dict(
+            file=expName,
+            class_name=classname,
+            log_level=logging.WARNING,
+            arguments=args,
+        )
+        log.info(f"Submitting an experiment {expName}")
+        rid = self.schedule.submit(pipeline_name="main", expid=expid, priority=0, due_date=None, flush=False)
+        self.expid_to_rid[exp_id] = rid
+        if self._check_rid(rid):
+            log.info(f"experiment with id {rid} scheduled")
+            return 0
+        else:
+            log.error(f"submitted exp with rid {rid} not found")
+            raise Exception
+
+    async def _wait_for_completion(self, rid):
+        """Wait for the experiment to complete."""
+        log.debug(f"Waiting for experiment {rid} to finish...")
+        while rid in self.schedule.get_status():
+            await asyncio.sleep(0.1)
+        self.logs
+        return 0
+
+    def _check_rid(self, rid):
+        return rid in self.schedule.get_status()

quantnet_agent/hal/driver/dummy.py

@@ -0,0 +1,373 @@
+from quantnet_agent.hal.hwclasses import LightSrc, Filter, LightMeasurement, ExpFramework
+import logging
+import asyncio
+import random
+
+log = logging.getLogger(__name__)
+
+
+class DummyLightSrc(LightSrc):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(args, kwargs)
+        self._power = 0
+        self._wavelength = 0
+        self._status = 0
+        self._failrate = float(args[0].get("failrate", 0))
+        log.info("Initializing Dummy Lightsource")
+        return
+
+    @property
+    def wavelength(self):
+        return self._wavelength
+
+    @wavelength.setter
+    def wavelength(self, freq):
+        log.info(f"Setting wavelength to {freq}")
+        self._wavelength = freq
+
+    @property
+    def power(self):
+        return self._power
+
+    @power.setter
+    def power(self, power):
+        log.info(f"Setting power to {power}")
+        self._power = power
+
+    async def generate(self, pol):
+        log.info(f"Generating {pol} light with to power = {self._power}, wavelength = {self._wavelength}")
+        if random.random() < self._failrate:
+            log.error("Light source failed")
+            return 1
+        else:
+            log.info("Calibrating for 10 sec")
+            await asyncio.sleep(10)
+            return 0
+
+    async def cleanUp(self):
+        self._status = 0
+        log.info("Cleaning up for 10 sec")
+        await asyncio.sleep(10)
+        return 0
+
+    @property
+    def status(self):
+        return self._status
+
+
+class DummyEPC(Filter):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(args, kwargs)
+        log.info("EPC initialized")
+        return
+
+    async def polarize(self, _):
+        log.info("Calibrating link using EPC")
+        await asyncio.sleep(5)
+        log.info("Link calibrate")
+        return 0
+
+    async def attenuate(self, _):
+        log.info("Attenuate link using EPC")
+        await asyncio.sleep(5)
+        log.info("Link attenuated")
+        return 0
+
+    async def cleanUp(self):
+        self._status = 0
+        return 0
+
+
+class DummyPolarimeter(LightMeasurement):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(args, kwargs)
+        self._wavelength = 0
+        log.info(f"Initializing LightMeasurement with wavelength={self._wavelength}")
+        return
+
+    @property
+    def wavelength(self):
+        return self._wavelength
+
+    @wavelength.setter
+    def wavelength(self, freq):
+        log.info(f"Setting wavelength to {freq}")
+        self._wavelength = freq
+
+    @property
+    def power(self):
+        return self._power
+
+    def sweep(self):
+        pass
+
+    async def measure(self):
+        log.info("Measuring SOP")
+        for i in range(5):
+            log.info(f"Generating random SOPs :[{random.random()},{random.random()},{random.random()}]")
+            await asyncio.sleep(1)
+        log.info("reached threashold SOP :[0.995, 0.995, 0.995]")
+        return 0
+
+    async def cleanUp(self):
+        self._status = 0
+        return 0
+
+
+class DummyExpFramework(ExpFramework):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(args, kwargs)
+        log.info("Initializing ExpFramework")
+        self.result_ind = 0
+        return
+
+    async def submit(self, exp_id, expName, classname, args=dict()):
+        log.info(f"Submitting exp:{expName} with id {id} with args:{args} for 10 sec")
+        # await asyncio.sleep(10)
+        return 0
+
+    async def receive(self, *args, **kwargs):
+        exp_id = args[0]
+        log.info(f"Receiving result from exp:{exp_id}")
+        results = [
+            {
+                "current_scan.plots.x": [
+                    170000000.0,
+                    172068965.5172414,
+                    174137931.03448275,
+                    176206896.55172414,
+                    178275862.06896552,
+                    180344827.58620688,
+                    182413793.10344827,
+                    184482758.62068966,
+                    186551724.13793105,
+                    188620689.6551724,
+                    190689655.1724138,
+                    192758620.68965518,
+                    194827586.20689654,
+                    196896551.72413793,
+                    198965517.24137932,
+                    201034482.75862068,
+                    203103448.27586207,
+                    205172413.79310346,
+                    207241379.31034482,
+                    209310344.8275862,
+                    211379310.3448276,
+                    213448275.86206895,
+                    215517241.37931034,
+                    217586206.89655173,
+                    219655172.4137931,
+                    221724137.93103448,
+                    223793103.44827586,
+                    225862068.96551722,
+                    227931034.4827586,
+                    230000000.0,
+                ],
+                "current_scan.plots.y": [
+                    [1896.0],
+                    [1715.0],
+                    [1913.0],
+                    [1844.0],
+                    [2061.0],
+                    [1977.0],
+                    [1936.0],
+                    [2039.0],
+                    [1937.0],
+                    [2319.0],
+                    [2458.0],
+                    [2743.0],
+                    [2943.0],
+                    [3469.0],
+                    [4059.0],
+                    [4817.0],
+                    [5377.0],
+                    [5807.0],
+                    [6039.0],
+                    [6113.0],
+                    [4882.0],
+                    [3784.0],
+                    [3607.0],
+                    [4132.0],
+                    [4200.0],
+                    [4517.0],
+                    [4693.0],
+                    [5064.0],
+                    [5146.0],
+                    [2963.0],
+                ],
+            },
+            {
+                "current_scan.plots.x": [
+                    30.0,
+                    28.94736842105263,
+                    27.894736842105264,
+                    26.842105263157894,
+                    25.789473684210527,
+                    24.736842105263158,
+                    23.684210526315788,
+                    22.63157894736842,
+                    21.578947368421055,
+                    20.526315789473685,
+                    19.473684210526315,
+                    18.42105263157895,
+                    17.36842105263158,
+                    16.315789473684212,
+                    15.263157894736842,
+                    14.210526315789474,
+                    13.157894736842106,
+                    12.105263157894736,
+                    11.05263157894737,
+                    10.0,
+                ],
+                "current_scan.plots.y": [
+                    [-2.0],
+                    [-26.0],
+                    [15.0],
+                    [4.0],
+                    [-31.0],
+                    [-14.0],
+                    [69.0],
+                    [63.0],
+                    [149.0],
+                    [177.0],
+                    [494.0],
+                    [576.0],
+                    [1112.0],
+                    [2253.0],
+                    [2947.0],
+                    [3776.0],
+                    [4876.0],
+                    [5963.0],
+                    [7261.0],
+                    [8522.0],
+                ],
+            },
+            {
+                "current_scan.plots.x": [
+                    60000000.0,
+                    61379310.344827585,
+                    62758620.68965517,
+                    64137931.03448276,
+                    65517241.37931035,
+                    66896551.72413793,
+                    68275862.06896552,
+                    69655172.4137931,
+                    71034482.7586207,
+                    72413793.10344827,
+                    73793103.44827586,
+                    75172413.79310346,
+                    76551724.13793103,
+                    77931034.48275863,
+                    79310344.8275862,
+                    80689655.1724138,
+                    82068965.51724139,
+                    83448275.86206897,
+                    84827586.20689656,
+                    86206896.55172414,
+                    87586206.89655173,
+                    88965517.24137932,
+                    90344827.5862069,
+                    91724137.93103449,
+                    93103448.27586207,
+                    94482758.62068966,
+                    95862068.96551725,
+                    97241379.31034483,
+                    98620689.65517241,
+                    100000000.0,
+                ],
+                "current_scan.plots.y": [
+                    [3034.0],
+                    [3063.0],
+                    [3198.0],
+                    [3290.0],
+                    [3017.0],
+                    [3097.0],
+                    [3066.0],
+                    [3188.0],
+                    [3201.0],
+                    [3348.0],
+                    [3278.0],
+                    [3243.0],
+                    [3223.0],
+                    [3145.0],
+                    [3238.0],
+                    [3159.0],
+                    [3193.0],
+                    [3328.0],
+                    [3227.0],
+                    [3342.0],
+                    [3338.0],
+                    [3433.0],
+                    [3293.0],
+                    [3261.0],
+                    [3205.0],
+                    [3437.0],
+                    [3306.0],
+                    [3325.0],
+                    [3359.0],
+                    [3394.0],
+                ],
+            },
+            {
+                "current_scan.plots.x": [
+                    30.0,
+                    28.94736842105263,
+                    27.894736842105264,
+                    26.842105263157894,
+                    25.789473684210527,
+                    24.736842105263158,
+                    23.684210526315788,
+                    22.63157894736842,
+                    21.578947368421055,
+                    20.526315789473685,
+                    19.473684210526315,
+                    18.42105263157895,
+                    17.36842105263158,
+                    16.315789473684212,
+                    15.263157894736842,
+                    14.210526315789474,
+                    13.157894736842106,
+                    12.105263157894736,
+                    11.05263157894737,
+                    10.0,
+                ],
+                "current_scan.plots.y": [
+                    [3514.0],
+                    [3372.0],
+                    [3508.0],
+                    [3575.0],
+                    [3508.0],
+                    [3536.0],
+                    [3671.0],
+                    [3489.0],
+                    [3597.0],
+                    [3494.0],
+                    [3596.0],
+                    [3576.0],
+                    [3661.0],
+                    [3721.0],
+                    [3531.0],
+                    [3609.0],
+                    [3596.0],
+                    [3626.0],
+                    [3681.0],
+                    [3543.0],
+                ],
+            },
+        ]
+        index = self.result_ind % 4
+        self.result_ind += 1
+        return {"dummy status": "done", "results": results[index]}
+
+    @property
+    def logs(self):
+        log.info("Getting log from ExpFramework")
+        return "Dummy logs"
+
+    async def cleanUp(self):
+        self._status = 0
+        return 0

quantnet_agent/hal/driver/overlayepc.py

@@ -0,0 +1,27 @@
+from quantnet_agent.hal.hwclasses import Filter
+import logging
+import socket
+
+log = logging.getLogger(__name__)
+
+
+class OverlayEPC(Filter):
+
+    def __init__(self):
+        log.info("Initializing EPC hardware")
+        return 0
+
+    def polarize(self, _):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect(("10.0.0.7", 5006))
+        s.send(b"calibrate")
+        result = s.recv(1024)
+        if result == b"OK":
+            return
+        elif result == b"Not OK":
+            raise Exception
+        else:
+            raise Exception
+
+    def attenuate(self, strength):
+        return super().attenuate(strength)

quantnet_agent/hal/driver/overlayexpframework.py

@@ -0,0 +1,47 @@
+from quantnet_agent.hal.hwclasses import ExpFramework
+import logging
+import socket
+import asyncio
+
+
+log = logging.getLogger(__name__)
+
+
+class OverlayExpFramework(ExpFramework):
+
+    def __init__(self, *args, **kwargs):
+        props = args[0]
+        self.bsm_ip = props['bsm_ip']
+        self.bsm_port = int(props['bsm_port'])
+        super().__init__(*args, **kwargs)
+
+    @property
+    def status(self):
+        return self._status
+
+    async def submit(self, exp_id, expName, classname, args=dict()):
+        log.info(f"submitting bsm to {self.bsm_ip}")
+        await self._send_message_async(self.bsm_ip, self.bsm_port, b'bsm')
+
+    def receive(self, exp_id):
+        return {"result": ""}
+
+    @property
+    def logs(self):
+        pass
+
+    async def cleanUp(self):
+        pass
+
+    async def _send_message_async(self, ip, port, msg):
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self._send_message, ip, port, msg)
+
+    def _send_message(self, ip, port, msg):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect((ip, port))
+        s.sendall(msg)
+        result = s.recv(1024)
+        if result != b'0':
+            raise Exception(f"Failed to handle {msg}")
+        s.close()

quantnet_agent/hal/driver/overlaylightsrc.py

@@ -0,0 +1,125 @@
+from quantnet_agent.hal.hwclasses import LightSrc
+
+import serial
+import time
+import socket
+import asyncio
+
+import logging
+
+log = logging.getLogger(__name__)
+
+polH = "S 4\n"
+polD = "S 5\n"
+polV = "S 6\n"
+polA = "S 7\n"
+polLCP = "S 2\n"
+polRCP = "S 13\n"
+
+
+class OverlayLightSrc(LightSrc):
+    def __init__(self, *args, **kwargs):
+        log.info("Initializing overlay LightSource")
+        props = args[0]
+        self.serial_dev = props['device']
+        self.baud_rate = props['baud_rate']
+        self.bsm_ip = props['bsm_ip']
+        self.bsm_port = int(props['bsm_port'])
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.sock.bind(("0.0.0.0", 5005))
+        self.sock.setblocking(False)
+        asyncio.create_task(self.serve())
+        super().__init__(args, kwargs)
+
+    def src_init(self, request):
+        return 0
+
+    async def serve(self):
+        loop = asyncio.get_running_loop()
+        self.dev_PSG, self.msg_device = self._connect()
+        while True:
+            log.info("Waiting for data")
+            data = await loop.sock_recv(self.sock, 1024)
+            log.info(data)
+            msg_devflag, msg_write, msg_read, polset_error = self._polSET(data.decode())
+            complete_msg = msg_devflag + "\n" + msg_write + "\n" + msg_read + "\n" + polset_error
+            log.info(complete_msg)
+            self.sock.sendto(complete_msg.encode(), (self.bsm_ip, self.bsm_port))
+
+    async def generate(self, _):
+        return 0
+
+    def is_device_connected(self, port_name):
+        try:
+            ser = serial.Serial(port_name)
+            ser.close()
+            return True
+        except serial.SerialException:
+            return False
+
+    def _connect(self):
+        try:
+            dev = serial.Serial(port=self.serial_dev, baudrate=self.baud_rate, timeout=0.1)
+        except Exception as e:
+            log.error(f"Could not connect to serial port {self.serial_dev}: {e}")
+            return None, ""
+        log.info(dev.name)
+        msg_device = "####### PSG is connected ############# "
+        log.info(msg_device)
+        return dev, msg_device
+
+    def _polSET(self, PSGpol):
+        try:
+            dev_connected_flag = self.is_device_connected(self.serial_dev)
+
+            if dev_connected_flag:
+                msg_devflag = "device connected flag:1, PSG is already connected"
+                log.info(msg_devflag)
+            else:
+                msg_devflag = "device connected flag:0, PSG was not connected. Connecting now"
+                log.info(msg_devflag)
+                self.dev_PSG, self.msg_device = self._connect()
+
+            try:
+                self.dev_PSG.write(PSGpol.encode())
+                msg_write = "Written: Ok"
+                log.info(msg_write)
+            except Exception as e:
+                log.info(e)
+                msg_write = "Write ISSUE"
+                log.info(msg_write)
+
+            time.sleep(1)
+            try:
+                self.dev_PSG.read(20)
+                msg_read = "Read: Ok"
+                log.info(msg_read)
+            except Exception as e:
+                log.info(e)
+                msg_read = "Read Issue"
+            polset_error = "Polset error None"
+            log.info(polset_error)
+        except Exception as e:
+            polset_error = f"Polarization: Could NOT set: {e}"
+            log.info(polset_error)
+
+        return msg_devflag, msg_write, msg_read, polset_error
+
+    def _disconnect(self, dev):
+        try:
+            msg_disconnect = "####### PSG is disconnected ############# "
+            log.info(msg_disconnect)
+            dev.close
+        except Exception:
+            msg_disconnect = "PSG: Could NOT disconnect"
+
+        return msg_disconnect
+
+    def wavelength(self, freq):
+        pass
+
+    def power(self, power):
+        pass
+
+    async def cleanUp(self):
+        return 0