From 072d4ed0508d8b7c49e8dc5bef2dc961e14b9f59 Mon Sep 17 00:00:00 2001
From: Sebastian Duetsch <sebastian.duetsch@zhinst.com>
Date: Tue, 20 Aug 2024 16:59:02 +0200
Subject: [PATCH] Implement check_ref_clock for HDAWG

---
 src/zhinst/toolkit/driver/devices/hdawg.py | 41 ++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/src/zhinst/toolkit/driver/devices/hdawg.py b/src/zhinst/toolkit/driver/devices/hdawg.py
index 245ec019..b92d9854 100644
--- a/src/zhinst/toolkit/driver/devices/hdawg.py
+++ b/src/zhinst/toolkit/driver/devices/hdawg.py
@@ -1,4 +1,6 @@
 """HDAWG Instrument Driver."""
+
+import logging
 import typing as t
 
 from zhinst.toolkit.driver.devices.base import BaseInstrument
@@ -10,6 +12,8 @@
 from zhinst.toolkit.nodetree.node import NodeList
 from zhinst.toolkit.exceptions import ToolkitError
 
+logger = logging.getLogger(__name__)
+
 
 class HDAWG(BaseInstrument):
     """High-level driver for the Zurich Instruments HDAWG."""
@@ -89,3 +93,40 @@ def awgs(self) -> t.Sequence[AWG]:
             self._root,
             self._tree + ("awgs",),
         )
+
+    def check_ref_clock(
+        self, *, timeout: float = 30.0, sleep_time: float = 1.0
+    ) -> bool:
+        """Check if reference clock is locked successfully.
+
+        Args:
+            timeout: Maximum time in seconds the program waits
+                (default: 30.0).
+            sleep_time: Time in seconds to wait between
+                requesting the reference clock status (default: 1)
+
+        Raises:
+            TimeoutError: If the process of locking to the reference clock
+                exceeds the specified timeout.
+        """
+        ref_clock_status = self.system.clocks.referenceclock.status
+        ref_clock = self.system.clocks.referenceclock.source
+        ref_clock_actual = self.system.clocks.referenceclock.sourceactual
+        try:
+            ref_clock_status.wait_for_state_change(
+                2, invert=True, timeout=timeout, sleep_time=sleep_time
+            )
+        except TimeoutError as error:
+            raise TimeoutError(
+                "Timeout during locking to reference clock signal"
+            ) from error
+        if ref_clock_status() == 0:
+            return True
+        if ref_clock_status() == 1 and ref_clock_actual() != ref_clock():
+            ref_clock("internal", deep=True)
+            logger.error(
+                f"There was an error locking the device({self.serial}) "
+                f"onto reference clock signal. Automatically switching to internal "
+                f"reference clock. Please try again."
+            )
+        return False