More HDAWG improvements

qupulse/hardware/awgs/zihdawg.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import functools
-from typing import Tuple, Set, Callable, Optional, Mapping, Generator, Union, List, Dict
+from typing import Tuple, Set, Callable, Optional, Mapping, Generator, Union, Dict, Any, NamedTuple
 from enum import Enum
 import weakref
 import logging
@@ -47,6 +47,9 @@ class HDAWGRepresentation:
     """HDAWGRepresentation represents an HDAWG8 instruments and manages a LabOne data server api session. A data server
     must be running and the device be discoverable. Channels are per default grouped into pairs."""
 
+    _MarkerState = NamedTuple("_MarkerState", [('idle', 'HDAWGTriggerOutSource'),
+                                               ('playback', 'HDAWGTriggerOutSource')])
+
     def __init__(self, device_serial: str = None,
                  device_interface: str = '1GbE',
                  data_server_addr: str = 'localhost',
@@ -64,6 +67,9 @@ def __init__(self, device_serial: str = None,
         :param reset:             Reset device before initialization
         :param timeout:           Timeout in seconds for uploading
         """
+        # TODO: lookup method to find channel count
+        n_channels = 8
+
         self._api_session = zhinst.ziPython.ziDAQServer(data_server_addr, data_server_port, api_level_number)
         assert zhinst.utils.api_server_version_check(self.api_session)  # Check equal data server and api version.
         self.api_session.connectDevice(device_serial, device_interface)
@@ -74,14 +80,13 @@ def __init__(self, device_serial: str = None,
             # Create a base configuration: Disable all available outputs, awgs, demods, scopes,...
             zhinst.utils.disable_everything(self.api_session, self.serial)
 
-        self._initialize()
+        self._initialize(force_defaults=reset)
 
         waveform_path = pathlib.Path(self.api_session.awgModule().getString('directory'), 'awg', 'waves')
         self._waveform_file_system = WaveformFileSystem(waveform_path)
         self._channel_groups: Dict[HDAWGChannelGrouping, Tuple[HDAWGChannelGroup, ...]] = {}
-
-        # TODO: lookup method to find channel count
-        n_channels = 8
+        self._marker_state = [self._MarkerState(HDAWGTriggerOutSource.OUT_1_MARK_1,
+                                                HDAWGTriggerOutSource.OUT_1_MARK_2)] * n_channels
 
         for grouping in HDAWGChannelGrouping:
             group_size = grouping.group_size()
@@ -129,41 +134,56 @@ def api_session(self) -> zhinst.ziPython.ziDAQServer:
     def serial(self) -> str:
         return self._dev_ser
 
-    def _initialize(self) -> None:
-        settings = [(f'/{self.serial}/awgs/*/userregs/*', 0),  # Reset all user registers to 0.
-                    (f'/{self.serial}/*/single', 1)]  # Single execution mode of sequence.
-        for ch in range(0, 8):  # Route marker 1 signal for each channel to marker output.
-            if ch % 2 == 0:
-                output = HDAWGTriggerOutSource.OUT_1_MARK_1.value
-            else:
-                output = HDAWGTriggerOutSource.OUT_1_MARK_2.value
-            settings.append(['/{}/triggers/out/{}/source'.format(self.serial, ch), output])
+    def _default_settings(self) -> Dict[str, Any]:
+        """Sensible default settings that might be changed by the user"""
+        # TODO: de-hardcode channel count of 8
+        marker_default_values = [HDAWGTriggerOutSource.OUT_1_MARK_1.value, HDAWGTriggerOutSource.OUT_1_MARK_2.value] * 4
+        marker_default_values = {
+            f'/{self.serial}/triggers/out/{ch}/source': value for ch, value in enumerate(marker_default_values)
+        }
+
+        return {
+            f'/{self.serial}/awgs/*/time': 0,  # Maximum sampling rate.
+            f'/{self.serial}/sigouts/*/range': HDAWGVoltageRange.RNG_1V.value,
+            f'/{self.serial}/awgs/*/outputs/*/amplitude': 0,
+            f'/{self.serial}/awgs/*/outputs/*/modulation/mode': HDAWGModulationMode.OFF.value,
+            **marker_default_values
+        }
+
+    def _required_settings(self) -> Dict[str, Any]:
+        """Settings that are required for this driver to work properly"""
+        return {
+            f'/{self.serial}/awgs/*/userregs/*': 0,  # Reset all user registers to 0.
+            f'/{self.serial}/awgs/*/single': 1  # Single execution mode of sequence.
+        }
+
+    def _initialize(self, force_defaults=False) -> None:
+        """Initialize the required settings of AWG to work as expected (single execution mode and user registers)
+        Args:
+            force_defaults: If true: Additionally sets sensible default settings for not required settings
+            (like marker source)
+        """
+        settings = self._required_settings()
+        if force_defaults:
+            settings.update(self._default_settings())
 
         self.api_session.set(settings)
         self.api_session.sync()  # Global sync: Ensure settings have taken effect on the device.
 
     def reset(self) -> None:
         zhinst.utils.disable_everything(self.api_session, self.serial)
-        self._initialize()
-        for tuple in self.channel_tuples:
-            tuple.clear()
-        self.api_session.set([
-            (f'/{self.serial}/awgs/*/time', 0),
-            (f'/{self.serial}/sigouts/*/range', HDAWGVoltageRange.RNG_1V.value),
-            (f'/{self.serial}/awgs/*/outputs/*/amplitude', 1.0),
-            (f'/{self.serial}/outputs/*/modulation/mode', HDAWGModulationMode.OFF.value),
-        ])
-
-        # marker outputs
-        marker_settings = []
-        for ch in range(0, 8):  # Route marker 1 signal for each channel to marker output.
-            if ch % 2 == 0:
-                output = HDAWGTriggerOutSource.OUT_1_MARK_1.value
-            else:
-                output = HDAWGTriggerOutSource.OUT_1_MARK_2.value
-            marker_settings.append([f'/{self.serial}/triggers/out/{ch}/source', output])
-        self.api_session.set(marker_settings)
-        self.api_session.sync()
+        self._initialize(force_defaults=True)
+
+    def _mark_source_setting(self, marker_idx: int, source: 'HDAWGTriggerOutSource') -> (str, Any):
+        assert marker_idx in range(8)
+        if isinstance(source, str):
+            source = getattr(HDAWGTriggerOutSource, source.upper())
+        elif isinstance(source, int):
+            source = HDAWGTriggerOutSource(source)
+        return f'/{self.serial}/triggers/out/{marker_idx}/source', source.value
+
+    def set_mark_source(self, marker_idx: int, source: 'HDAWGTriggerOutSource'):
+        self.api_session.setInt(*self._mark_source_setting(marker_idx, source))
 
     def group_name(self, group_idx, group_size) -> str:
         return str(self.serial) + '_' + 'ABCDEFGH'[group_idx*group_size:][:group_size]
@@ -249,8 +269,8 @@ class HDAWGTriggerOutSource(Enum):
     TRIG_IN_6 = 13  # Trigger output assigned to trigger inout 6.
     TRIG_IN_7 = 14  # Trigger output assigned to trigger inout 7.
     TRIG_IN_8 = 15  # Trigger output assigned to trigger inout 8.
-    HIGH = 17 # Trigger output is set to high.
-    LOW = 18 # Trigger output is set to low.
+    HIGH = 17  # Trigger output is set to high.
+    LOW = 18  # Trigger output is set to low.
 
 
 class HDAWGChannelGrouping(Enum):