Add find_zsync_worker_port and adapt examples

examples/active_qubit_reset.md

@@ -6,7 +6,7 @@ jupyter:
       extension: .md
       format_name: markdown
       format_version: '1.3'
-      jupytext_version: 1.14.1
+      jupytext_version: 1.14.4
   kernelspec:
     display_name: Python 3 (ipykernel)
     language: python
@@ -31,7 +31,6 @@ Requirements:
 
 ```python vscode={"languageId": "python"}
 from zhinst.toolkit import Session, SHFQAChannelMode, Waveforms, Sequence, CommandTable
-import pqsc_helpers
 import numpy as np
 import time, textwrap
 ```
@@ -97,7 +96,7 @@ with session.set_transaction():
     # PQSC - Replicate ZSync triggers on trigger output port
     pqsc.triggers.out[0].enable(True)
     pqsc.triggers.out[0].source("start_trigger")
-    pqsc.triggers.out[0].port(pqsc_helpers.find_zsync_worker_port(pqsc, shfsg))
+    pqsc.triggers.out[0].port(pqsc.find_zsync_worker_port(shfsg))
 ```
 
 ```python vscode={"languageId": "python"}
@@ -207,7 +206,7 @@ Here we configure the PQSC to do register forwarding of the result of the qubit
 
 ```python vscode={"languageId": "python"}
 # Find SHFSG ZSync port ID
-shfsg_zsync_port = pqsc_helpers.find_zsync_worker_port(pqsc, shfsg)
+shfsg_zsync_port = pqsc.find_zsync_worker_port(shfsg)
 
 with pqsc.set_transaction():
     # Enable Readout Register Bank forwarding to ZSync output

examples/bell_state_stabilization.md

@@ -6,9 +6,9 @@ jupyter:
       extension: .md
       format_name: markdown
       format_version: '1.3'
-      jupytext_version: 1.14.1
+      jupytext_version: 1.14.4
   kernelspec:
-    display_name: toolkit
+    display_name: Python 3 (ipykernel)
     language: python
     name: python3
 ---
@@ -50,7 +50,6 @@ Create a session with a running Data Server, and then connect the Data Server wi
 
 ```python
 from zhinst.toolkit import Session, SHFQAChannelMode, Waveforms, Sequence, CommandTable
-import pqsc_helpers
 import numpy as np
 import time, textwrap
 ```
@@ -111,7 +110,7 @@ with session.set_transaction():
     # PQSC - Replicate ZSync triggers on trigger output port
     pqsc.triggers.out[0].enable(True)
     pqsc.triggers.out[0].source("start_trigger")
-    pqsc.triggers.out[0].port(pqsc_helpers.find_zsync_worker_port(pqsc, shfsg))
+    pqsc.triggers.out[0].port(pqsc.find_zsync_worker_port(shfsg))
 ```
 
 ```python
@@ -260,7 +259,7 @@ And finally enable the Decoder Unit.
 
 ```python
 # Find SHFSG ZSync port ID
-shfsg_zsync_port = pqsc_helpers.find_zsync_worker_port(pqsc, shfsg)
+shfsg_zsync_port = pqsc.find_zsync_worker_port(shfsg)
 
 with pqsc.set_transaction():
     # Enable Decoder Unit

examples/repeat_until_success.md

@@ -6,9 +6,9 @@ jupyter:
       extension: .md
       format_name: markdown
       format_version: '1.3'
-      jupytext_version: 1.14.1
+      jupytext_version: 1.14.4
   kernelspec:
-    display_name: toolkit
+    display_name: Python 3 (ipykernel)
     language: python
     name: python3
 ---
@@ -33,7 +33,6 @@ Requirements:
 
 ```python
 from zhinst.toolkit import Session, SHFQAChannelMode, Waveforms, Sequence, CommandTable
-import pqsc_helpers
 from zhinst.toolkit.waveform import Wave
 import numpy as np
 import time, textwrap
@@ -100,7 +99,7 @@ with session.set_transaction():
     # PQSC - Replicate ZSync triggers on trigger output port
     pqsc.triggers.out[0].enable(True)
     pqsc.triggers.out[0].source("start_trigger")
-    pqsc.triggers.out[0].port(pqsc_helpers.find_zsync_worker_port(pqsc, shfsg))
+    pqsc.triggers.out[0].port(pqsc.find_zsync_worker_port(shfsg))
 ```
 
 ```python
@@ -365,10 +364,8 @@ Finally, enable the Decoder.
 
 ```python
 # Find SHFSG and SHFQA ZSync port IDs
-import pqsc_helpers
-
-shfsg_zsync_port = pqsc_helpers.find_zsync_worker_port(pqsc, shfsg)
-shfqa_zsync_port = pqsc_helpers.find_zsync_worker_port(pqsc, shfqa)
+shfsg_zsync_port = pqsc.find_zsync_worker_port(shfsg)
+shfqa_zsync_port = pqsc.find_zsync_worker_port(shfqa)
 
 with pqsc.set_transaction():
     for zsync_port in [shfsg_zsync_port, shfqa_zsync_port]:

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

@@ -6,6 +6,7 @@
 
 from zhinst.toolkit.driver.devices.base import BaseInstrument
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -220,3 +221,37 @@ def _check_zsync_connection(
                 f"on the port {port}."
             ) from error
         return status_node() == 2
+
+    def find_zsync_worker_port(self, device: BaseInstrument) -> int:
+        """Find the ID of the PQSC ZSync port connected to a given device.
+
+        Args:
+            pqsc: PQSC device over whose ports the research shall be done.
+            device: device for which the connected ZSync port shall be found.
+
+        Returns:
+            Integer value represent the ID of the searched PQSC Zsync port.
+
+        Raises:
+            RuntimeError: If the given device doesn't appear to be connected
+                to the PQSC via ZSync.
+
+        """
+        device_serial = device.serial[3:]
+        node_to_serial_dict = self.zsyncs["*"].connection.serial()
+        serial_to_node_dict = {
+            serial: node for node, serial in node_to_serial_dict.items()
+        }
+        # Get the node of the ZSync connected to the device
+        # (will have the form "/devXXXX/zsyncs/N/connection/serial")
+        try:
+            device_zsync_node = serial_to_node_dict[device_serial]
+        except KeyError:
+            raise RuntimeError(
+                "No ZSync connection found between the PQSC "
+                f"{self.serial} and the device {device.serial}."
+            )
+        # Just interested in knowing N: split in
+        # ['', 'devXXXX', 'zsyncs', 'N', 'connection', 'serial']
+        # and take fourth value
+        return int(device_zsync_node.split("/")[3])

tests/test_pqsc.py

@@ -208,3 +208,20 @@ def getInt_side_effect(path):
     assert [True, False, True] == pqsc.check_zsync_connection(
         [0, 2, 10], sleep_time=0.001
     )
+
+
+def test_find_zsync_worker_port(mock_connection, pqsc):
+    # device not connected
+    ret_val = {"random/path": [{"timestamp": 0, "flags": 0, "vector": ""}]}
+    mock_connection.return_value.get.return_value = ret_val
+    with pytest.raises(RuntimeError):
+        pqsc.find_zsync_worker_port(pqsc)
+
+    # correct use case
+    ret_val = {
+        "/dev1234/zsyncs/0/connection/serial": [
+            {"timestamp": 0, "flags": 0, "vector": "1234"}
+        ]
+    }
+    mock_connection.return_value.get.return_value = ret_val
+    assert pqsc.find_zsync_worker_port(pqsc) == 0