Add a paramater to specify the QCCS generation in the function `enabl…

…e_qccs_mode` in the HDAWG driver.

CHANGELOG.md

@@ -1,5 +1,10 @@
 # zhinst-toolkit Changelog
 
+## Version 0.6.2
+* The function `enable_qccs_mode` of the HDAWG driver now accept an optional argument to select the QCCS generation.
+  It's advised to set it to 2 (gen2) when the HDAWG is operated with PQSC together with SHF instruments, and set it
+  to 1 (gen1) when the HDAWG is operated with PQSC together with UHFQA instruments.
+
 ## Version 0.6.1
 * Deep gets on nodes with keywords returns an enum like the regular get.
 * Fix rare failures of `wait_for_state_change` function that resulted in early timeouts.

docs/source/spelling_wordlist.txt

@@ -138,3 +138,4 @@ LabOne
 lru
 func
 enums
+GSa

src/zhinst/toolkit/driver/devices/hdawg.py

@@ -8,32 +8,64 @@
     lazy_property,
 )
 from zhinst.toolkit.nodetree.node import NodeList
+from zhinst.toolkit.exceptions import ToolkitError
 
 
 class HDAWG(BaseInstrument):
     """High-level driver for the Zurich Instruments HDAWG."""
 
-    def enable_qccs_mode(self) -> None:
+    def enable_qccs_mode(self, gen: int = 1) -> None:
         """Configure the instrument to work with PQSC.
 
         This method sets the reference clock source to
         connect the instrument to the PQSC.
 
+        Args:
+            gen: The QCCS generation that is being configured.
+                 Use 1 for a gen1 system, when only HDAWG and UHFQA are used.
+                 In this case, the sample rate is set to 2.4 GSa/s and the DIO
+                 interface is configured for connection with the UHFQA.
+                 Use 2 for a gen2 system, when only HDAWG and SHFs are used.
+                 In this case, the sample rate is set to 2.0 GSa/s.
+                 (default: 1)
+
+            .. versionchanged:: 0.6.2: Added gen parameter.
+
+        Raises:
+            ToolkitError: If the gen argument is not correct.
+
         Info:
             Use ``factory_reset`` to reset the changes if necessary
         """
         with create_or_append_set_transaction(self._root):
+            # Enable QCCS mode, so the instrument is able
+            # to operate with ZSync
+            self.dios[0].mode("qccs")
+
+            # Gen1 specific elements
+            if gen == 1:
+                # Set the sample clock to 2.4 GSa/s
+                self.system.clocks.sampleclock.freq(2.4e9)
+                # Configure DIO
+                # Set interface standard to use on the 32-bit DIO to LVCMOS
+                self.dios[0].interface(0)
+                # Drive the two most significant bytes of the DIO port
+                # so the UHFQA receives triggers from PQSC
+                self.dios[0].drive(0b1100)
+
+            # Gen2 specific elements
+            elif gen == 2:
+                # Set the sample clock to 2.0 GSa/s
+                self.system.clocks.sampleclock.freq(2.0e9)
+                # Turn off DIO
+                self.dios[0].drive(0b0000)
+
+            else:
+                raise ToolkitError("Only gen1 or gen2 are supported!")
+
             # Set ZSync clock to be used as reference
             self.system.clocks.referenceclock.source("zsync")
-            # Configure DIO
-            # Set interface standard to use on the 32-bit DIO to LVCMOS
-            self.dios[0].interface(0)
-            # Set DIO output values to ZSync input values.
-            # Forward the ZSync input values to the AWG sequencer.
-            # Forward the DIO input values to the ZSync output.
-            self.dios[0].mode("qccs")
-            # Drive the two most significant bytes of the DIO port
-            self.dios[0].drive(0b1100)
+
             # Disable DIO triggering on the AWGs,
             # since it's not needed for ZSync messages
             self.awgs["*"].dio.strobe.slope("off")

tests/test_hdawg.py

@@ -3,6 +3,7 @@
 
 from zhinst.toolkit.driver.devices.hdawg import AWG, HDAWG
 from zhinst.toolkit.nodetree import Node
+from zhinst.toolkit.exceptions import ToolkitError
 
 
 @pytest.fixture()
@@ -29,31 +30,70 @@ def test_repr(hdawg):
 
 
 def test_enable_qccs_mode(mock_connection, hdawg):
+    # No argument, equal to gen1
     hdawg.enable_qccs_mode()
     mock_connection.return_value.set.assert_called_with(
         [
-            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
-            ("/dev1234/dios/0/interface", 0),
             ("/dev1234/dios/0/mode", "qccs"),
+            ("/dev1234/system/clocks/sampleclock/freq", 2.4e9),
+            ("/dev1234/dios/0/interface", 0),
             ("/dev1234/dios/0/drive", 12),
+            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
             ("/dev1234/awgs/*/dio/strobe/slope", "off"),
             ("/dev1234/awgs/*/dio/valid/polarity", "none"),
         ]
     )
+
+    # Set transaction check
     mock_connection.reset_mock()
     with hdawg.set_transaction():
         hdawg.enable_qccs_mode()
     mock_connection.return_value.set.assert_called_with(
         [
-            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
+            ("/dev1234/dios/0/mode", "qccs"),
+            ("/dev1234/system/clocks/sampleclock/freq", 2.4e9),
             ("/dev1234/dios/0/interface", 0),
+            ("/dev1234/dios/0/drive", 12),
+            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
+            ("/dev1234/awgs/*/dio/strobe/slope", "off"),
+            ("/dev1234/awgs/*/dio/valid/polarity", "none"),
+        ]
+    )
+
+    # Gen1 test
+    mock_connection.reset_mock()
+    hdawg.enable_qccs_mode(gen=1)
+    mock_connection.return_value.set.assert_called_with(
+        [
             ("/dev1234/dios/0/mode", "qccs"),
+            ("/dev1234/system/clocks/sampleclock/freq", 2.4e9),
+            ("/dev1234/dios/0/interface", 0),
             ("/dev1234/dios/0/drive", 12),
+            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
             ("/dev1234/awgs/*/dio/strobe/slope", "off"),
             ("/dev1234/awgs/*/dio/valid/polarity", "none"),
         ]
     )
 
+    # Gen2 test
+    mock_connection.reset_mock()
+    hdawg.enable_qccs_mode(gen=2)
+    mock_connection.return_value.set.assert_called_with(
+        [
+            ("/dev1234/dios/0/mode", "qccs"),
+            ("/dev1234/system/clocks/sampleclock/freq", 2.0e9),
+            ("/dev1234/dios/0/drive", 0),
+            ("/dev1234/system/clocks/referenceclock/source", "zsync"),
+            ("/dev1234/awgs/*/dio/strobe/slope", "off"),
+            ("/dev1234/awgs/*/dio/valid/polarity", "none"),
+        ]
+    )
+
+    # wrong argument
+    mock_connection.reset_mock()
+    with pytest.raises(ToolkitError):
+        hdawg.enable_qccs_mode(gen=1234)
+
 
 def test_awg(hdawg):
     assert len(hdawg.awgs) == 4