Add extra calls to HDAWG driver, with more utilization of zhinst-toolkit

Open
#185 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Stale
Tech stack
python
Domain
backend

Research direction

Start at the HDAWG driver entry points in qmi.instruments.zurich_instruments.hdawg and review the listed unit tests alongside the zhinst-toolkit AWG APIs. Done means all six RPC calls are implemented, their listed tests pass, and the new calls have been tested on the hardware.

Written by the indexing model from the issue text.

Description

Is your feature request related to a problem? Please describe.
In the addition of grouping mode handling some dependencies into the new zhinst-toolkit package have been introduced. At the same time several nice-to-have calls utilizing some of the new features in zhinst-toolkit were proposed for merging. It was decided to postpone the addition of most of these to another ticket to keep the grouping mode pull request a bit more contained. The postponed RPC calls with their respective unit-tests are listed below.

Describe the solution you'd like
Add the following RPC calls and respective unit-tests:

    @rpc_method
    def compile_sequencer_program(self, awg_core: int, sequencer_program: str) -> tuple[bytes, Any]:
        """Compile the given sequencer program for specific AWG channel, which is translated into the
        respective AWG core.

        Parameters:
            awg_core:                  AWG core number [0-3].
            sequencer_program: The sequencer program as a string.

        Returns:
            compiled_program: The compiled program as bytes.
            compiler_output:  Output of the compilation.
        """
        self._check_is_open()
        _logger.info("[%s] Compiling sequencer program", self._name)
        # Get the AWG node/core
        awg_node = self.device.awgs[awg_core]

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

        return compiled_program, compiler_output

    @rpc_method
    def upload_compiled_program(self, awg_core: int, compiled_program: bytes) -> Any:
        """Upload the given compiled program for specific AWG channel. The AWG channel is translated into the
        respective AWG core number.

        Parameters:
            awg_core:                AWG core number [0-3].
            compiled_program: The compiled program to upload.

        Returns:
            upload_info: The returned value from the upload command.
        """
        self._check_is_open()
        _logger.info("[%s] Uploading sequencer program", self._name)
        # Get the AWG node/core
        awg_node = self.device.awgs[awg_core]

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

        return upload_info

    @rpc_method
    def get_sequence_snippet(self, waveforms: "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 write_to_waveform_memory(self, awg_core: int, waveforms: "Waveforms", indexes: None | list = None) -> None:
        """Write waveforms to the waveform memory. The waveforms must already be assigned in the sequencer program.

        Parameters:
            awg_core:     AWG core number [0-3].
            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 = self.device.awgs[awg_core]
        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_core: int, indexes: None | list[int] = None) -> "Waveforms":
        """Read waveforms from the waveform memory for an AWG channel.

        Parameters:
            awg_core: AWG core number [0-3].
            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 = self.device.awgs[awg_core]

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

    @rpc_method
    def upload_waveforms_per_awg_core(
        self, unpacked_waveforms: list[tuple[int, int, np.ndarray, None | np.ndarray, None | np.ndarray]]
    ) -> None:
        """Upload a set of new waveform data to the AWG. Works as singular waveform uploading, but organizing waveforms
        into Waveforms objects per AWG channel pairs. Then upload is done per one AWG channel pair at a time.

        The loop that creates the unpacked waveforms loops such that:
        - outer loop: for waveform_index, sequence in enumerate(waveforms)
          - inner loop: for awg_core in range(4), where channels are paired per core:
              channel_a = 2 * awg_core + 1  # 1, 3, 5, 7
              channel_b = 2 * awg_core + 2  # 2, 4, 6, 8

        Parameters:
            unpacked_waveforms: List of tuples, each tuple is a collection of AWG core index, waveform sequence index,
                                wave1, wave2 and markers.
        """
        self._check_is_open()

        waveforms = {}
        for sequence in unpacked_waveforms:
            awg_core, waveform_index, wave1, wave2, markers = sequence
            markers, wave1, wave2 = self._control_waveform_inputs(markers, wave1, wave2)
            if awg_core not in waveforms:
                waveforms[awg_core] = Waveforms()

            waveforms[awg_core][waveform_index] = (wave1, wave2, markers)

        for awg_core in waveforms.keys():
            self.device.awgs[awg_core].write_to_waveform_memory(waveforms[awg_core])

    @rpc_method
    def get_command_table(self, awg_core: int) -> "zhinst.toolkit.CommandTable":
        """Get the command table for the respective core of the channel from the device.

        Parameters:
            awg_core:        AWG core number [0-3].

        Returns:
            command_table: 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 = self.device.awgs[awg_core]

        return awg_node.commandtable.load_from_device()

Unit-tests:

    def test_compile_sequencer_program(self):
        """Test compiling sequencer program."""
        channel = 4
        program = "Very short program"
        node_mock = self.hdawg.device.awgs[channel // 2]
        node_mock.compile_sequencer_program.return_value = (program.encode(), {})

        _, __ = self.hdawg.compile_sequencer_program(channel // 2, program)

        node_mock.compile_sequencer_program.assert_called_once_with(program)

    def test_upload_compiled_program(self):
        """Test uploading a compiled program.""" 
        channel = 6
        comp_prog = b"Compiled Very short program"
        node_mock = self.hdawg.device.awgs[channel // 2]

        _ = self.hdawg.upload_compiled_program(channel // 2, comp_prog)

        node_mock.elf.data.assert_called_once_with(comp_prog)

    def test_get_sequence_snippet(self):
        """Test getting a sequence snippet"""
        expected_snippet = "snippy-snip"
        waveforms_mock = WaveformMock()

        snippet = self.hdawg.get_sequence_snippet(waveforms_mock)

        self.assertEqual(expected_snippet, snippet)

    def test_write_to_waveform_memory(self):
        """Test writing to waveform memory with and without indexes"""
        channel = 5
        node_mock = self.hdawg.device.awgs[channel // 2]
        wf_mock = Mock()
        # No indexes test
        self.hdawg.write_to_waveform_memory(channel // 2, wf_mock, None)

        node_mock.write_to_waveform_memory.assert_called_once_with(wf_mock)
        node_mock.reset_mock()

        # With indexes test
        indexes = [0, 2]
        self.hdawg.write_to_waveform_memory(channel // 2, wf_mock, indexes)

        node_mock.write_to_waveform_memory.assert_called_once_with(wf_mock, indexes)

    def test_read_from_waveform_memory(self):
        """Test reading from waveform memory with and without indexes."""
        channel = 6
        node_mock = self.hdawg.device.awgs[channel // 2]
        # No indexes test
        _ = self.hdawg.read_from_waveform_memory(channel // 2, None)

        node_mock.read_from_waveform_memory.assert_called_once_with()
        node_mock.reset_mock()

        # With indexes test
        indexes = [0, 2]
        self.hdawg.read_from_waveform_memory(channel // 2, indexes)

        node_mock.read_from_waveform_memory.assert_called_once_with(indexes)

    def test_upload_waveforms_per_awg_core(self):
        """Test uploading a 'large' batch of waveforms (more than batch size limit but exact multiple of)."""
        unpacked_waveforms = []
        batch_size = 30
        for i in range(batch_size):
            awg_index = i % 4
            waveform_index = i
            wave1 = np.array([i + 1, i + 2, i + 3])
            wave2 = np.array([i + 4, i + 5, i + 6])
            markers = np.array([i + 7, i + 8, i + 9])
            unpacked_waveforms.append((awg_index, waveform_index, wave1, wave2, markers))

        core_mock_0 = self._device.awgs[0]
        core_mock_1 = self._device.awgs[1]
        core_mock_2 = self._device.awgs[2]
        core_mock_3 = self._device.awgs[3]
        with unittest.mock.patch(
            "qmi.instruments.zurich_instruments.hdawg.zhinst"
        ), unittest.mock.patch(
            "qmi.instruments.zurich_instruments.hdawg.zhinst.utils"
        ), unittest.mock.patch(
            "qmi.instruments.zurich_instruments.hdawg.Waveforms"
        ) as wf_patch:
            self.hdawg.upload_waveforms_per_awg_core(unpacked_waveforms)

        core_mock_0.write_to_waveform_memory.assert_called_with(wf_patch())
        core_mock_1.write_to_waveform_memory.assert_called_with(wf_patch())
        core_mock_2.write_to_waveform_memory.assert_called_with(wf_patch())
        core_mock_3.write_to_waveform_memory.assert_called_with(wf_patch())

    def test_get_command_table(self):
        """Test getting a command table for a channel."""
        channel = 0
        node_mock = self.hdawgdevice.awgs[channel // 2]

        _ = self.hdawg.get_command_table(channel // 2)

        node_mock.commandtable.load_from_device.assert_called_once_with()

Describe alternatives you've considered
N/A

Acceptance criteria
New calls have been tested on the HW.

Additional context
Add any other context or screenshots about the feature request here.

Dominant language
Python
Stars
25
Forks
9
PR merge metrics
No merged PRs in 30d

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from QuTech-Delft/QMI

All issues in QuTech-Delft/QMI

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.