Support 4x2 and 2x4 mode in HDAWG QMI driver

[heevasti]

#Description

The HDAWG QMI driver now only supports 1x8 channel grouping mode. The reason for this is that there was never a use case for any of the other modes and in 2x4 or 4x2 mode you need to upload multiple sequencer codes and command tables (one for each core, 2 or 4 respectively).

Proposed changes

The affected functions are those that use the self._awg_module attribute. These are

• __init__()
• open()
• close()
• set_channel_grouping() (remove assertion on mode == 2)
• compile_and_upload() and sub functions called therein
• set_awg_module_enabled()
• other new functions introduced in the PR !78

The proposal is to:

1. change the _awg_module attribute to be a list of 4 modules (one for each core)
2. modify open() and close() to initialize/finalize the AWG module handles
3. modify compile_and_upload() and set_awg_module_enabled() to take a awg_index argument

TBD: In order not to break existing scripts, it would be good to have awg_index an optional argument defaulting to zero, but that breaks the convention that awg_index is the first argument. Since only two functions are affected, it may be acceptable to break the API for this.

Possible caveats:

• At this time it's not clear what happens when you try e.g. self._awg_module.set("awgModule/index", 3) when in 1x8 or 2x4 mode. It may be needed to query the device to get the current channel grouping mode before initializing AWG module handles. Also, when changing the channel grouping mode, it may be needed to terminate invalidated handles or initialize more handles.

EDIT 18-11-25:

For grouped modes other than 1x8, it is likely that the use of the base zhinst.core and zhinst.utils modules is complicated to adapt the code to work in all grouped modes. We should step over to use as much as possible the zhinst-toolkit package, which is developed by Zurich Instruments themselves. It does mean we step one step back from the core of the control of the device, but the simplification for grouped mode control would make it worth the exchange. It is also MIT licensed, so no problems there as well.

There is a grouped mode example in the repo which shows how the cores and channels can also be defined in the grouped modes etc. We can also make use of the advanced upload features described in the example.

The use of the zhinst-toolkit should also enable almost all the existing RPC methods as well. Here will be laid out how we can achieve control of the existing calls with the new package.

1. expand imports

# Lazy import of the zhinst module. See the function _import_modules() below.
if typing.TYPE_CHECKING:
    import zhinst.toolkit
    import zhinst.toolkit.driver
    import zhinst.toolkit.driver.devices
    import zhinst.toolkit.driver.modules
    import zhinst.toolkit.driver.nodes
    import zhinst.toolkit.driver.nodes.awg as awg
    import zhinst.toolkit.nodetree
    import zhinst.toolkit.session
    import zhinst.core
    import zhinst.utils
else:
    zhinst = None

# and in import_modules:
    global zhinst
    if zhinst is None:
        import zhinst.toolkit
        import zhinst.toolkit.driver
        import zhinst.toolkit.driver.devices
        import zhinst.toolkit.driver.nodes
        import zhinst.toolkit.driver.nodes.awg
        import zhinst.toolkit.nodetree
        import zhinst.toolkit.session
        import zhinst.core  # pylint: disable=W0621
        import zhinst.utils

2. In __init__:
   Extra input parameter for grouping: grouping: str = "1x8". These could be mapped as enums as well in the same order as in the grouped mode example, which then enable the use of the example code for setting the class variables related to number of AWGs, number of channels (per AWG) and channel-to-core mapping.

        # ZI HDAWG server, module and device
        self._daq_server: None | zhinst.core.ziDAQServer = None
        self._awg_module: None | zhinst.core.AwgModule = None
        self._device: None | zhinst.toolkit.driver.devices.HDAWG = None
        ...
        # Connect to Zurich Instruments Data Server
        self._session = zhinst.toolkit.Session(self._server_host, self._server_port)

3. property

    @property
    def device(self) -> zhinst.toolkit.driver.devices.HDAWG:
        assert self._device is not None
        return self._device

4. (Optional) AWG channel input for _get_... and _set_value (note that _set_int and _set_double do not seem to have implementation in the zhinst.toolkit):

    def _get_int(self, node_path: str, awg_channel: None | int = None) -> int:
        """
        Get an integer value from the nodetree.

        Parameters:
            node_path:      The path to the node to be queried.
            awg_channel:    Optional, an AWG channel to get the node for. Default is None.

        Returns:
            integer value from node tree
        """
        if awg_channel is not None:
            awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]
            return awg_node.root.connection.getInt(node_path)

        return self.daq_server.getInt('/' + self._device_name + '/' + node_path)

and same for double and string, and

    def _set_value(self, node_path: str, value: str | int | float, awg_channel: None | int = None) -> None:
        """
        Set an integer value from the nodetree.

        Parameters:
            node_path:      The path to the node to be queried.
            value:          Value to set for the node.
            awg_channel:    Optional, an AWG channel to get the node for. Default is None.
        """
        if awg_channel is not None:
            awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]
            awg_node.root.connection.set(node_path, value)

        else:
            self.daq_server.set('/' + self._device_name + '/' + node_path, value)

And _set_int and _set_double would remain unchanged. The respective RPC methods also need the extra awg_channel input parameter.
5. The internal functions related to compilation, upload and their checks might become obsolete. See if they can be removed. Their removal should not break anything from the user perspective.
6. The open function replacements and removals.

        # Remove: self._daq_server = zhinst.core.ziDAQServer(self._server_host, self._server_port, api_level=6)
        self._daq_server = self._session.daq_server

        # Connect to the device.
        self._device = typing.cast(
            zhinst.toolkit.driver.devices.HDAWG, self._session.connect_device(self._device_name)
        )
        # Remove this alternative? Connect to the device via Ethernet.
        # Remove? If the device is already connected, this is a no-op.
        # Remove? self.daq_server.connectDevice(self._device_name, "1GbE")

        # Remove: self._awg_module = self.daq_server.awgModule()
        self._awg_module = typing.cast(zhinst.core.AwgModule, self._session.modules.awg)

        # Remove? self.awg_module.set("device", self._device_name)
        # Remove: self.awg_module.set("index", 0)  # only support 1x8 mode, so only one AWG module
        # TBD: Set grouping based on init
         self.set_channel_grouping(self._grouping)

7. def set_channel_grouping(...)
   This function needs fixing based on the grouped mode example.
8. sync

       # Remove: self.daq_server.sync()
       self._session.sync()

9. Compiling

    @rpc_method
    def compile_sequencer_program(self, awg_channel: int, sequencer_program: str) -> bytes:
        """
        Compile the given sequencer program.

        Parameters:
            awg_channel:        The AWG channel to compile the program for.
            sequencer_program:  The sequencer program as a string.

        Returns:
            compiled program as bytes.
        """
        self._check_is_open()
        _logger.info("[%s] Compiling sequencer program", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        # Remove? Test compiling.
        # Remove? compilation_result = self._wait_compile(sequencer_program)
        # Remove? self._last_compilation_successful = self._interpret_compilation_result_is_ok(compilation_result)

        compiled_program, compiler_output = awg_node.compile_sequencer_program(sequencer_program)
        _logger.debug(f"Compilation Info:\n{compiler_output}")

        return compiled_program

and | or in compile_and_upload

    @rpc_method
    def compile_and_upload(
            self,
            awg_channel: int,
            sequencer_program: str | zhinst.toolkit.Sequence,
            replacements: None | dict[str, str| int | float] = None,
    ) -> None:
        """Compile and upload the sequencer program.

        This function combines compilation followed by upload to the AWG if compilation was successful. This is forced
        by the HDAWG API.

        Notes on parameter replacements:
         - Parameters must adhere to the following format: $[A-Za-z][A-Za-z0-9]+ (literal $, followed by at least one
            letter followed by zero or more alphanumeric characters or underscores). Both the key in the replacements
            dictionary and the parameter reference in the sequencer code must adhere to this format.
         - Replacement respects word boundaries. The inclusion of the '$' allows to concatenate values in the sequencer
            program code, e.g. "wave w = "$TYPE$LENGTH;" to achieve "wave w = "sin1024";", but be careful that you
            don't accidentally create new parameters by concatenation.
         - Replacement values must be of type str, int or float.

        Warning:
            After uploading the sequencer program one needs to wait for the awg core to become ready, before it can be
            enabled. The awg core indicates the ready state through its `ready` node.
            (device.awgs[0].ready() == True). Example:
            > compile_info = self.device.awgs[0].load_sequencer_program(seqc)
            > self.device.awgs[0].ready.wait_for_state_change(1)
            > self.device.awgs[0].enable(True)

        Parameters:
            awg_channel:       The AWG channel to compile the program for.
            sequencer_program: Sequencer program as a string or a Sequencer class.
            replacements:      Optional dictionary of (parameter, value) pairs. Every occurrence of the parameter in
                               the sequencer program will be replaced with the specified value.
        """
        self._check_is_open()
        _logger.info("[%s] Loading sequencer program", self._name)

        if replacements is not None:
            # Perform parameter replacements.
            sequencer_program = self._process_parameter_replacements(sequencer_program, replacements)

        self._check_program_not_empty(sequencer_program)

        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        # Load sequencer program. Equivalent to compiling and uploading.
        try:
            info = awg_node.load_sequencer_program(sequencer_program)
            _logger.debug(f"Loading sequencer program info:\n{info}")
            awg_node.ready.wait_for_state_change(1, timeout=self._UPLOAD_TIMEOUT)
            self._last_compilation_successful = True
        except RuntimeError as err:
            self._last_compilation_successful = False
            _logger.exception(f"Loading sequencer program %s failed: %s", sequencer_program, str(err))
            raise QMI_RuntimeException("Loading sequencer program failed.") from err

10. Addition of new or replacement functions:

    @staticmethod
    def get_sequence_snippet(waveforms: zhinst.toolkit.Waveforms) -> str:
        """
        Get a sequencer code snippet that defines the given waveforms.

        Parameters:
            waveforms:  Waveforms to generate snippet for.

        Returns:
            Sequencer code snippet as a string.
        """
        _logger.info("[ZurichInstruments_HDAWG]: Generating sequencer code snippet for waveforms")
        return waveforms.get_sequence_snippet()

    @rpc_method
    def wait_done(self, awg_channel: int) -> None:
        """
        Wait for AWG to finish.

        Parameters:
            awg_channel:   AWG channel to wait for.
        """
        self._check_is_open()
        _logger.info("[%s] Waiting for sequencer to finish", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        awg_node.wait_done(timeout=self.GENERATOR_WAIT_TIME_S)

    @rpc_method
    def enable_sequencer(self, awg_channel: int, disable_when_finished: bool = True) -> None:
        """
        Enable the sequencer.

        Parameters:
            awg_channel:            AWG channel to enable.
            disable_when_finished:  Flag to disable sequencer after it finishes execution. Default is True.
        """
        self._check_is_open()
        _logger.info("[%s] Enabling sequencer", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        awg_node.enable_sequencer(single=disable_when_finished)

    @rpc_method
    def upload_program(self, awg_channel: int, compiled_program: bytes) -> None:
        """
        Upload the given compiler program.

        Parameters:
            awg_channel:        The AWG channel to upload the program to.
            compiled_program:   The compiled program to upload.
        """
        self._check_is_open()
        _logger.info("[%s] Uploading sequencer program", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        upload_info = awg_node.elf.data(compiled_program)
        _logger.debug(f"Upload Info:\n{upload_info}")

    @rpc_method
    def get_command_table(self, awg_channel: int) -> zhinst.toolkit.CommandTable:
        """
        Get the command table from the device.

        Parameters:
            awg_channel:    The AWG channel to get the command table for.

        Returns:
            The command table.
        """
        self._check_is_open()
        _logger.info("[%s] Getting command table for channel [%d]", self._name, awg_channel)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        return awg_node.commandtable.load_from_device()

    @rpc_method
    def write_to_waveform_memory(
        self, awg_channel: int, waveforms: zhinst.toolkit.Waveforms, indexes: None | list = None
    ) -> None:
        """
        Write waveforms to the waveform memory. The waveforms must already be assigned in the sequencer program.

        Parameters:
            awg_channel:    The AWG channel to upload the waveforms to.
            waveforms:      Waveforms to write.
            indexes:        List of indexes to upload. Default is None, which uploads all waveforms.
        """
        self._check_is_open()
        _logger.info("[%s] Writing waveforms to waveform memory", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        awg_node.write_to_waveform_memory(waveforms, indexes) if indexes else awg_node.write_to_waveform_memory(
            waveforms
        )

    @rpc_method
    def read_from_waveform_memory(
        self, awg_channel: int, indexes: None | list[int] = None
    ) -> zhinst.toolkit.Waveforms:
        """
        Read waveforms from the waveform memory.

        Parameters:
            awg_channel:    The AWG channel to upload the waveforms to.
            indexes:        List of indexes to read. Default is None, which uploads all waveforms.

        Returns:
            Waveforms from waveform memory.
        """
        self._check_is_open()
        _logger.info("[%s] Reading waveforms from waveform memory", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        return awg_node.read_from_waveform_memory(indexes) if indexes else awg_node.read_from_waveform_memory()

    @rpc_method
    def validate_waveforms(
        self,
        awg_channel: int,
        waveforms: zhinst.toolkit.Waveforms,
        compiled_sequencer_program: None | bytes = None,
    ) -> None:
        """
        Validate if the waveforms match the sequencer program.

        Parameters:
            awg_channel:                AWG channel.
            waveforms:                  Waveforms to validate.
            compiled_sequencer_program: Optional sequencer program. If this is not provided then information from the device is used.
        """
        self._check_is_open()
        _logger.info("[%s] Validating waveforms", self._name)
        # Get the AWG node/core
        awg_node: awg.AWG = self.device.awgs[self.CHANNEL_TO_CORE_MAPPING[awg_channel]]

        waveforms.validate(
            compiled_sequencer_program if compiled_sequencer_program else awg_node.waveform.descriptors()
        )

11. Setting command table commands. Replace with toolkit equivalents?
12. Further calls based on zhinst-toolkit replacements:

        self.device.system.clocks.referenceclock.source(value)
        # self._set_int("system/clocks/referenceclock/source", value)

        return self.device.system.clocks.referenceclock.status()
        # return self._get_int("system/clocks/referenceclock/status")

        self.device.system.clocks.sampleclock.freq(frequency)
        # self._set_double("system/clocks/sampleclock/freq", frequency)

        return self.device.system.clocks.sampleclock.status()
        # return self._get_int("system/clocks/sampleclock/status")

        self.device.triggers[trigger].imp50(value)
        # self._set_int("triggers/in/{}/imp50".format(trigger), value)

And any further RPC methods should be checked if we can or need (due to grouping) change the calls to be based on the toolkit.

Don't forget to document the changes well.

Finally, the nicest part of all: Edit and create new unit-tests to test the module properly!

Affected components

QMI

Modules to be created

None

Modules to be modified

qmi.instruments.zurich_instruments.hdawg

Tests to be created/updated

Unit tests

Documentation to be updated

It should be made clear in the docstrings that when using 2x4 or 4x2 mode that the user is responsible for uploading correct sequencer code and command tables ~~, since it is impossible to check that a priori (e.g. don't try to access output 8 from sequencer code for core 2). ~~ But there probably can be some check introduced.

Hardware

Test on a HDAWG. Check with labs.

Acceptance criteria

• Functionality implemented in agreement with description
• Tests are updated and pass
• Code quality metrics ok
• Documentation updated