Merge pull request #25 from qua-platform/feat/set-dc-offset

Feat/set dc offset for channels

quam/components/channels.py

@@ -1,5 +1,5 @@
 from dataclasses import field
-from typing import ClassVar, Dict, List, Optional, Tuple, Union
+from typing import ClassVar, Dict, List, Literal, Optional, Tuple, Union
 import warnings
 
 from quam.components.hardware import BaseFrequencyConverter, Mixer, LocalOscillator
@@ -17,6 +17,7 @@
     measure,
     dual_demod,
     declare,
+    set_dc_offset,
     fixed,
     frame_rotation,
 )
@@ -409,6 +410,18 @@ class SingleChannel(Channel):
     opx_output_offset: float = None
     intermediate_frequency: float = None
 
+    def set_dc_offset(self, offset: QuaNumberType):
+        """Set the DC offset of an element's input to the given value.
+        This value will remain the DC offset until changed or until the Quantum Machine
+        is closed.
+
+        Args:
+            offset (QuaNumberType): The DC offset to set the input to.
+                This is limited by the OPX output voltage range.
+                The number can be a QUA variable
+        """
+        set_dc_offset(element=self.name, element_input="single", offset=offset)
+
     def apply_to_config(self, config: dict):
         """Adds this SingleChannel to the QUA configuration.
 
@@ -558,6 +571,26 @@ def mixer(self) -> Optional[Mixer]:
     def rf_frequency(self):
         return self.frequency_converter_up.LO_frequency + self.intermediate_frequency
 
+    def set_dc_offset(self, offset: QuaNumberType, element_input: Literal["I", "Q"]):
+        """Set the DC offset of an element's input to the given value.
+        This value will remain the DC offset until changed or until the Quantum Machine
+        is closed.
+
+        Args:
+            offset (QuaNumberType): The DC offset to set the input to.
+                This is limited by the OPX output voltage range.
+                The number can be a QUA variable
+            element_input (Literal["I", "Q"]): The element input to set the offset for.
+
+        Raises:
+            ValueError: If element_input is not "I" or "Q"
+        """
+        if element_input not in ["I", "Q"]:
+            raise ValueError(
+                f"element_input should be either 'I' or 'Q', got {element_input}"
+            )
+        set_dc_offset(element=self.name, element_input=element_input, offset=offset)
+
     def apply_to_config(self, config: dict):
         """Adds this IQChannel to the QUA configuration.
 

tests/components/channels/test_IQ_channel.py

@@ -0,0 +1,27 @@
+import pytest
+from quam.components import IQChannel
+
+
+def test_IQ_channel_set_dc_offset(mocker):
+    mocker.patch("quam.components.channels.set_dc_offset")
+
+    channel = IQChannel(
+        id="channel",
+        opx_output_I=("con1", 1),
+        opx_output_Q=("con1", 2),
+        frequency_converter_up=None,
+    )
+
+    with pytest.raises(TypeError):
+        channel.set_dc_offset(0.5)
+
+    with pytest.raises(ValueError):
+        channel.set_dc_offset(0.5, "X")
+
+    channel.set_dc_offset(0.5, "I")
+
+    from quam.components.channels import set_dc_offset
+
+    set_dc_offset.assert_called_once_with(
+        element="channel", element_input="I", offset=0.5
+    )

tests/components/channels/test_single_channel.py

@@ -123,3 +123,16 @@ class QuamTest(QuamRoot):
     cfg = machine.generate_config()
     expected_cfg["controllers"]["con1"]["analog_outputs"][1]["offset"] = 0.1
     assert cfg == expected_cfg
+
+
+def test_single_channel_set_dc_offset(mocker):
+    mocker.patch("quam.components.channels.set_dc_offset")
+
+    channel = SingleChannel(id="channel", opx_output=("con1", 1))
+    channel.set_dc_offset(0.5)
+
+    from quam.components.channels import set_dc_offset
+
+    set_dc_offset.assert_called_once_with(
+        element="channel", element_input="single", offset=0.5
+    )