updated quaconfig (#17)

* updated quaconfig

* revised formatted and added tests

entropylab_qpudb/_quaconfig.py

@@ -1,6 +1,7 @@
 from collections import UserDict as _UserDict
 from copy import deepcopy as _deepcopy
 import json as _json
+from typing import Union, Tuple, List
 
 
 class QuaConfig(_UserDict):
@@ -75,7 +76,17 @@ def dump(self, filename):
         with open("config.json", "w") as fp:
             _json.dump(self.data, fp)
 
-    def get_waveforms_from_op(self, element, operation):
+    def get_waveforms_from_op(
+        self, element: str, operation: str
+    ) -> Union[Tuple[List[float], List[float]], List[float]]:
+        """
+        Get output waveforms associated with an operation on a quantum element.
+        For both arbitrary and constant pulses, the waveform returned will be the actual values played.
+
+        :param element: Name of the element
+        :param operation: Name of the operation
+        :return: Either a waveform entries list if element is of type singleInput, or a tuple of waveform entries list if element is of type mixedInputs.
+        """
         pulse = self.get_pulse_from_op(element, operation)
         if "mixInputs" in self.data["elements"][element]:
             waveform_i = self.data["waveforms"][pulse["waveforms"]["I"]]
@@ -133,18 +144,70 @@ def update_intermediate_frequency(
 
     def update_op_amp(self, element, operation, amp):
         pulse = self.get_pulse_from_op(element, operation)
-        try:
-            self.data["waveforms"][pulse["waveforms"]["single"]]["sample"] = amp
-        except KeyError:
+        if "sample" not in self.data["waveforms"][pulse["waveforms"]["single"]]:
             raise KeyError(
                 "Can only access amplitude for a constant pulse of a single element"
             )
+        else:
+            self.data["waveforms"][pulse["waveforms"]["single"]]["sample"] = amp
 
     def get_op_amp(self, element, operation):
+        """
+        DEPRECATED - DO NOT USE
+        use `get_waveforms_from_op` instead
+        """
         pulse = self.get_pulse_from_op(element, operation)
         try:
             return self.data["waveforms"][pulse["waveforms"]["single"]]["sample"]
         except KeyError:
             raise KeyError(
                 "Can only access amplitude for a constant pulse of a single element"
             )
+
+    def get_port_by_element_input(
+        self, element: str, element_input: str
+    ) -> Tuple[str, int]:
+        """
+        returns the ports of a quantum element.
+        :param element: Name of the element
+        :param element_input:
+            Name of the element port. Can be either 'single' if element has `singleInput`
+            or 'I', 'Q' if element has mixed inputs.
+        :return: a tuple of the form (con_name, port number)
+        """
+        element_data = self.data["elements"][element]
+        if element_input == "single":
+            if "singleInput" in element_data:
+                return element_data["singleInput"]["port"]
+            else:
+                raise ValueError(
+                    "can only use 'single' for a singleInput quantum element"
+                )
+        if "mixInputs" in element_data:
+            if element_input == "I" or element_input == "Q":
+                return element_data["mixInputs"][element_input]
+            else:
+                raise ValueError(
+                    "can only use 'I' or 'Q' for a mixInputs quantum element"
+                )
+        else:
+            raise ValueError(
+                f"element input is {element_input} but can only be I, Q for mixInputs or "
+                f"single for singleInput"
+            )
+
+    def set_output_dc_offset_by_element(
+        self, element: str, port: str, offset: float
+    ) -> None:
+        """
+        Set a DC offset value by element
+
+        :param element: Name of the element
+        :param port: Name of the port. Can be either 'single' if element has `singleInput` or 'I', 'Q' if element has mixed inputs.
+        :param offset: offset value to set
+        """
+        con, port = self.get_port_by_element_input(element, port)
+        if port in self.data["controllers"][con]["analog_outputs"]:
+            self.data["controllers"][con]["analog_outputs"][port]["offset"] = offset
+        else:
+            self.data["controllers"][con]["analog_outputs"][port] = {"offset": offset}

entropylab_qpudb/tests/test_quaconfig.py

@@ -0,0 +1,107 @@
+# todo
+import pytest
+from entropylab_qpudb import QuaConfig
+
+
+@pytest.fixture
+def config():
+    pulse_len = 1000
+    return QuaConfig(
+        {
+            "version": 1,
+            "controllers": {
+                "con1": {
+                    "type": "opx1",
+                    "analog_outputs": {
+                        1: {"offset": +0.0},
+                    },
+                    "analog_inputs": {
+                        1: {"offset": +0.0},
+                    },
+                }
+            },
+            "elements": {
+                "qe1": {
+                    "singleInput": {"port": ("con1", 1)},
+                    "outputs": {"output1": ("con1", 1)},
+                    "intermediate_frequency": 100e6,
+                    "operations": {
+                        "readoutOp": "readoutPulse",
+                        "readoutOp2": "readoutPulse2",
+                    },
+                    "time_of_flight": 180,
+                    "smearing": 0,
+                },
+                "qe2": {
+                    "mixInputs": {"I": ("con1", 2), "Q": ("con1", 3)},
+                    "intermediate_frequency": 100e6,
+                    "operations": {
+                        "readoutOp": "readoutPulse",
+                        "readoutOp2": "readoutPulse2",
+                    },
+                },
+            },
+            "pulses": {
+                "readoutPulse": {
+                    "operation": "measure",
+                    "length": pulse_len,
+                    "waveforms": {"single": "ramp_wf"},
+                    "digital_marker": "ON",
+                    "integration_weights": {"x": "xWeights", "y": "yWeights"},
+                },
+                "readoutPulse2": {
+                    "operation": "measure",
+                    "length": 2 * pulse_len,
+                    "waveforms": {"single": "ramp_wf2"},
+                    "digital_marker": "ON",
+                    "integration_weights": {"x": "xWeights2", "y": "yWeights"},
+                },
+            },
+            "waveforms": {
+                "const_wf": {"type": "constant", "sample": 0.2},
+                "ramp_wf": {
+                    "type": "arbitrary",
+                    "samples": [n / (2 * pulse_len) for n in range(pulse_len)],
+                },
+                "ramp_wf2": {
+                    "type": "arbitrary",
+                    "samples": [n / (2 * pulse_len) for n in range(pulse_len)],
+                },
+            },
+            "digital_waveforms": {
+                "ON": {"samples": [(1, 0)]},
+            },
+            "integration_weights": {
+                "xWeights": {
+                    "cosine": [1.0] * (pulse_len // 4),
+                    "sine": [0.0] * (pulse_len // 4),
+                },
+                "xWeights2": {
+                    "cosine": [1.0] * (2 * pulse_len // 4),
+                    "sine": [0.0] * (2 * pulse_len // 4),
+                },
+                "yWeights": {
+                    "cosine": [0.0] * (pulse_len // 4),
+                    "sine": [1.0] * (pulse_len // 4),
+                },
+            },
+        }
+    )
+
+
+def test_get_port_by_element_input(config):
+    # single
+    r = config.get_port_by_element_input("qe1", "single")
+    assert r == ("con1", 1)
+    # IQ pair
+    r = config.get_port_by_element_input("qe2", "I")
+    assert r == ("con1", 2)
+
+
+def test_set_output_dc_offset_by_element(config):
+    # single
+    config.set_output_dc_offset_by_element("qe1", "single", 0.3)
+    assert config["controllers"]["con1"]["analog_outputs"][1]["offset"] == 0.3
+    # IQ pair
+    config.set_output_dc_offset_by_element("qe2", "I", 0.2)
+    assert config["controllers"]["con1"]["analog_outputs"][2]["offset"] == 0.2