From 4aaa4163d97e9ad48563829ba2c83b0f06d77326 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 7 Feb 2024 11:57:13 +0400
Subject: [PATCH 01/45] remove emulation option

---
 src/qibolab/instruments/zhinst.py | 37 ++++---------------------------
 1 file changed, 4 insertions(+), 33 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index ef7e5e22e..9f55fabd4 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -9,7 +9,6 @@
 import laboneq._token
 import laboneq.simple as lo
 import numpy as np
-from laboneq.contrib.example_helpers.plotting.plot_helpers import plot_simulation
 from laboneq.dsl.experiment.pulse_library import (
     sampled_pulse_complex,
     sampled_pulse_real,
@@ -291,13 +290,11 @@ class Zurich(Controller):
     PortType = ZhPort
 
     def __init__(
-        self, name, device_setup, use_emulation=False, time_of_flight=0.0, smearing=0.0
+        self, name, device_setup, time_of_flight=0.0, smearing=0.0
     ):
         self.name = name
         "Setup name (str)"
 
-        self.emulation = use_emulation
-        "Enable emulation mode (bool)"
         self.is_connected = False
         "Is the device connected ? (bool)"
 
@@ -355,7 +352,7 @@ def connect(self):
             # To fully remove logging #configure_logging=False
             # I strongly advise to set it to 20 to have time estimates of the experiment duration!
             self.session = lo.Session(self.device_setup, log_level=20)
-            self.device = self.session.connect(do_emulation=self.emulation)
+            self.device = self.session.connect()
             self.is_connected = True
 
     def disconnect(self):
@@ -1404,31 +1401,5 @@ def sweep_recursion_nt(
             else:
                 self.define_exp(qubits, couplers, options, exp, exp_calib)
 
-    def play_sim(self, qubits, sequence, options, sim_time):
-        """Play pulse sequence."""
-
-        self.experiment_flow(qubits, sequence, options)  # missing couplers?
-        self.run_sim(sim_time)
-
-    def run_sim(self, sim_time):
-        """Run the simulation.
-
-        Args:
-            sim_time (float): Time[s] to simulate starting from 0
-        """
-        # create a session
-        self.sim_session = lo.Session(self.device_setup)
-        # connect to session
-        self.sim_device = self.sim_session.connect(do_emulation=True)
-        self.exp = self.sim_session.compile(
-            self.experiment, compiler_settings=COMPILER_SETTINGS
-        )
-
-        # Plot simulated output signals with helper function
-        plot_simulation(
-            self.exp,
-            start_time=0,
-            length=sim_time,
-            plot_width=10,
-            plot_height=3,
-        )
+    def split_batches(self, sequences):
+        return batch_max_sequences(sequences, MAX_SEQUENCES)

From 3e9ca7152007edcafefd428694a3aec75efb5f7f Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 7 Feb 2024 12:11:51 +0400
Subject: [PATCH 02/45] remove redundant comments

---
 src/qibolab/instruments/zhinst.py | 41 ++++---------------------------
 1 file changed, 5 insertions(+), 36 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 9f55fabd4..4a9691bc9 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -117,17 +117,6 @@ def select_pulse(pulse, pulse_type):
             can_compress=True,
         )
 
-    # Implement Slepian shaped flux pulse https://arxiv.org/pdf/0909.5368.pdf
-
-    # """
-    # Typically, the sampler function should discard ``length`` and ``amplitude``, and
-    # instead assume that the pulse extends from -1 to 1, and that it has unit
-    # amplitude. LabOne Q will automatically rescale the sampler's output to the correct
-    # amplitude and length.
-
-    # They don't even do that on their notebooks
-    # and just use lenght and amplitude but we have to check
-
 
 @dataclass
 class ZhPort(Port):
@@ -218,17 +207,12 @@ def add_sweeper(self, sweeper, qubit):
 class ZhSweeperLine:
     """Zurich sweeper from qibolab sweeper for non pulse parameters Bias, Delay
     (, power_range, local_oscillator frequency, offset ???)
-
-    For now Parameter.bias sweepers are implemented as
-    Parameter.Amplitude on a flux pulse. We may want to keep this class
-    separate for future Near Time sweeps
     """
 
     def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
         self.sweeper = sweeper
         """Qibolab sweeper."""
 
-        # Do something with the pulse coming here
         if sweeper.parameter is Parameter.bias:
             if isinstance(qubit, Qubit):
                 pulse = FluxPulse(
@@ -332,7 +316,6 @@ def __init__(
         self.sequence = defaultdict(list)
         "Zurich pulse sequence"
         self.sequence_qibo = None
-        # Remove if able
         self.sub_sequences = {}
         "Sub sequences between each measurement"
 
@@ -524,13 +507,11 @@ def run_exp(self):
         self.exp = self.session.compile(
             self.experiment, compiler_settings=COMPILER_SETTINGS
         )
-        # self.exp.save_compiled_experiment("saved_exp")
         self.results = self.session.run(self.exp)
 
     @staticmethod
     def frequency_from_pulses(qubits, sequence):
         """Gets the frequencies from the pulses to the qubits."""
-        # Implement Dual drive frequency experiments, we don't have any for now
         for pulse in sequence:
             qubit = qubits[pulse.qubit]
             if pulse.type is PulseType.READOUT:
@@ -625,8 +606,6 @@ def play(self, qubits, couplers, sequence, options):
                     qubit = ropulse.pulse.qubit
                     results[serial] = results[qubit] = options.results_type(data)
 
-        # html containing the pulse sequence schedule
-        # lo.show_pulse_sheet("pulses", self.exp)
         return results
 
     def sequence_zh(self, sequence, qubits, couplers):
@@ -639,8 +618,6 @@ def sequence_zh(self, sequence, qubits, couplers):
         for pulse in sequence:
             zhsequence[f"{pulse.type.name.lower()}{pulse.qubit}"].append(ZhPulse(pulse))
 
-        # Mess that gets the sweeper and substitutes the pulse it sweeps in the right place
-
         def nt_loop(sweeper):
             if not self.nt_sweeps:
                 self.nt_sweeps = [sweeper]
@@ -809,8 +786,7 @@ def play_sweep_select_single(exp, qubit, pulse, section, parameters, partial_swe
             exp.play(
                 signal=f"{section}{qubit.name}",
                 pulse=pulse.zhpulse,
-                phase=pulse.zhsweeper,  # I believe this is the global phase sweep
-                # increment_oscillator_phase=pulse.zhsweeper, # I believe this is the relative phase sweep
+                phase=pulse.zhsweeper,
             )
         elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
             exp.play(
@@ -1111,8 +1087,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                     if i == len(self.sequence[f"readout{q}"]) - 1:
                         reset_delay = relaxation_time * NANO_TO_SECONDS
                     else:
-                        # Here time of flight or not ?
-                        reset_delay = 0  # self.time_of_flight * NANO_TO_SECONDS
+                        reset_delay = 0
 
                     exp.measure(
                         acquire_signal=f"acquire{q}",
@@ -1219,8 +1194,6 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
                     qubit = ropulse.pulse.qubit
                     results[serial] = results[qubit] = options.results_type(data)
 
-        # html containing the pulse sequence schedule
-        # lo.show_pulse_sheet("pulses", self.exp)
         return results
 
     def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
@@ -1250,8 +1223,6 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 pulse = pulse.copy()
                 pulse.amplitude *= max(abs(sweeper.values))
 
-                # Proper copy(sweeper) here if we want to keep the sweepers
-                # sweeper_aux = copy.copy(sweeper)
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
@@ -1281,9 +1252,9 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
             ).zhsweeper
 
         with exp.sweep(
-            uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",  # This uid trouble double freq ???
+            uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
             parameter=parameter,
-            reset_oscillator_phase=True,  # Should we reset this phase ???
+            reset_oscillator_phase=True,
         ):
             if len(self.sweepers) > 0:
                 self.sweep_recursion(qubits, couplers, exp, exp_calib, exp_options)
@@ -1362,8 +1333,6 @@ def sweep_recursion_nt(
                 pulse = pulse.copy()
                 pulse.amplitude *= max(abs(sweeper.values))
 
-                # Proper copy(sweeper) here
-                # sweeper_aux = copy.copy(sweeper)
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
@@ -1381,7 +1350,7 @@ def sweep_recursion_nt(
                     f"/{path}/qachannels/*/oscs/0/gain"  # Hardcoded SHFQA device
                 )
 
-        elif parameter is None:  # can it be accessed?
+        elif parameter is None:
             parameter = ZhSweeper(
                 sweeper.pulses[0], sweeper, qubits[sweeper.pulses[0].qubit]
             ).zhsweeper

From 4d19afa2087b8a450cb1c9249c9d33e950253409 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 7 Feb 2024 12:15:57 +0400
Subject: [PATCH 03/45] remove fast reset (to reimplement it somewhere else in
 the future)

---
 src/qibolab/instruments/zhinst.py | 31 -------------------------------
 1 file changed, 31 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 4a9691bc9..cfd8321ca 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -368,13 +368,6 @@ def calibration_step(self, qubits, couplers, options):
                     - qubit.readout.local_oscillator.frequency,
                     options=options,
                 )
-                if options.fast_reset is not False:
-                    if len(self.sequence[f"drive{qubit.name}"]) == 0:
-                        self.register_drive_line(
-                            qubit=qubit,
-                            intermediate_frequency=qubit.drive_frequency
-                            - qubit.drive.local_oscillator.frequency,
-                        )
         self.device_setup.set_calibration(self.calibration)
 
     def register_readout_line(self, qubit, intermediate_frequency, options):
@@ -698,9 +691,6 @@ def create_exp(self, qubits, couplers, options):
             if len(self.sequence[f"readout{q}"]) != 0:
                 signals.append(lo.ExperimentSignal(f"measure{q}"))
                 signals.append(lo.ExperimentSignal(f"acquire{q}"))
-                if options.fast_reset is not False:
-                    if len(self.sequence[f"drive{q}"]) == 0:
-                        signals.append(lo.ExperimentSignal(f"drive{q}"))
 
         exp = lo.Experiment(
             uid="Sequence",
@@ -760,8 +750,6 @@ def select_exp(self, exp, qubits, couplers, exp_options):
             exp_options.relaxation_time,
             exp_options.acquisition_type,
         )
-        if exp_options.fast_reset is not False:
-            self.fast_reset(exp, qubits, exp_options.fast_reset)
 
     @staticmethod
     def play_sweep_select_single(exp, qubit, pulse, section, parameters, partial_sweep):
@@ -1106,25 +1094,6 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                         reset_delay=reset_delay,
                     )
 
-    def fast_reset(self, exp, qubits, fast_reset):
-        """
-        Conditional fast reset after readout - small delay for signal processing
-        This is a very naive approach that can be improved by repeating this step until
-        we reach non fast reset fidelity
-        https://quantum-computing.ibm.com/lab/docs/iql/manage/systems/reset/backend_reset
-        """
-        log.warning("Im fast resetting")
-        for qubit_name in self.sequence_qibo.qubits:
-            qubit = qubits[qubit_name]
-            q = qubit.name  # pylint: disable=C0103
-            with exp.section(uid=f"fast_reset{q}", play_after=f"sequence_measure"):
-                with exp.match_local(handle=f"sequence{q}"):
-                    with exp.case(state=0):
-                        pass
-                    with exp.case(state=1):
-                        pulse = ZhPulse(qubit.native_gates.RX.pulse(0, 0))
-                        exp.play(signal=f"drive{q}", pulse=pulse.zhpulse)
-
     @staticmethod
     def rearrange_sweepers(sweepers: List[Sweeper]) -> Tuple[np.ndarray, List[Sweeper]]:
         """Rearranges sweepers from qibocal based on device hardware

From 4c247c9b936a26744edcdaa243da7a1520790058 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 7 Feb 2024 17:44:45 +0400
Subject: [PATCH 04/45] remove unused instance variables self.device,
 self.chip, self.exp_options, self.exp_calib, self._ports, and self.settings

---
 src/qibolab/instruments/zhinst.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index cfd8321ca..927aee52d 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -289,18 +289,14 @@ def __init__(
 
         self.device_setup = device_setup
         self.session = None
-        self.device = None
         "Zurich device parameters for connection"
 
         self.time_of_flight = time_of_flight
         self.smearing = smearing
-        self.chip = "iqm5q"
         "Parameters read from the runcard not part of ExecutionParameters"
 
         self.exp = None
         self.experiment = None
-        self.exp_options = ExecutionParameters()
-        self.exp_calib = lo.Calibration()
         self.results = None
         "Zurich experiment definitions"
 
@@ -323,8 +319,6 @@ def __init__(
         self.nt_sweeps = None
         "Storing sweepers"
         # Improve the storing of multiple sweeps
-        self._ports = {}
-        self.settings = None
 
     @property
     def sampling_rate(self):
@@ -335,12 +329,12 @@ def connect(self):
             # To fully remove logging #configure_logging=False
             # I strongly advise to set it to 20 to have time estimates of the experiment duration!
             self.session = lo.Session(self.device_setup, log_level=20)
-            self.device = self.session.connect()
+            _ = self.session.connect()
             self.is_connected = True
 
     def disconnect(self):
         if self.is_connected:
-            self.device = self.session.disconnect()
+            _ = self.session.disconnect()
             self.is_connected = False
 
     def calibration_step(self, qubits, couplers, options):

From 2b53646b5daf81c56c6ce492c632176f534631c4 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 10:20:18 +0400
Subject: [PATCH 05/45] remove laboneq token related code

---
 src/qibolab/instruments/zhinst.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 927aee52d..a22c70019 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -1,12 +1,10 @@
 """Instrument for using the Zurich Instruments (Zhinst) devices."""
 
 import copy
-import os
 from collections import defaultdict
 from dataclasses import dataclass, replace
 from typing import Dict, List, Tuple, Union
 
-import laboneq._token
 import laboneq.simple as lo
 import numpy as np
 from laboneq.dsl.experiment.pulse_library import (
@@ -25,9 +23,6 @@
 from .abstract import Controller
 from .port import Port
 
-# this env var just needs to be set
-os.environ["LABONEQ_TOKEN"] = "not required"
-laboneq._token.is_valid_token = lambda _token: True  # pylint: disable=W0212
 
 SAMPLING_RATE = 2
 NANO_TO_SECONDS = 1e-9

From 772ba44a4219cc207f6900cb6f7b4ce8f4cf4c00 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 11:23:51 +0400
Subject: [PATCH 06/45] call super constructor for instrument

---
 src/qibolab/instruments/zhinst.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index a22c70019..9cbb192e6 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -271,11 +271,7 @@ class Zurich(Controller):
     def __init__(
         self, name, device_setup, time_of_flight=0.0, smearing=0.0
     ):
-        self.name = name
-        "Setup name (str)"
-
-        self.is_connected = False
-        "Is the device connected ? (bool)"
+        super().__init__(name, None)
 
         self.signal_map = {}
         "Signals to lines mapping"

From 77c7a3a9eb587ff4feb79ba36de085902f9bd127 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 11:40:30 +0400
Subject: [PATCH 07/45] minimize ZhPort (preparing to remove)

---
 src/qibolab/instruments/zhinst.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 9cbb192e6..32ee09911 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -113,13 +113,6 @@ def select_pulse(pulse, pulse_type):
         )
 
 
-@dataclass
-class ZhPort(Port):
-    name: Tuple[str, str]
-    offset: float = 0.0
-    power_range: int = 0
-
-
 class ZhPulse:
     """Zurich pulse from qibolab pulse translation."""
 
@@ -263,6 +256,12 @@ def select_sweeper(sweeper):
             )
 
 
+# FIXME: not needed for any logic inside the driver. Only needed to meet the expectations set by parent class Controller.
+@dataclass
+class ZhPort:
+    name: str
+
+
 class Zurich(Controller):
     """Zurich driver main class."""
 

From 390d1e1ef47843761cf6444ce78d4e2eb7699f79 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 12:55:52 +0400
Subject: [PATCH 08/45] Revert "minimize ZhPort (preparing to remove)"

This reverts commit 5cb192d09af1e0659bbe5e5dbc9eff76367e46ea.
---
 src/qibolab/instruments/zhinst.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 32ee09911..9cbb192e6 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -113,6 +113,13 @@ def select_pulse(pulse, pulse_type):
         )
 
 
+@dataclass
+class ZhPort(Port):
+    name: Tuple[str, str]
+    offset: float = 0.0
+    power_range: int = 0
+
+
 class ZhPulse:
     """Zurich pulse from qibolab pulse translation."""
 
@@ -256,12 +263,6 @@ def select_sweeper(sweeper):
             )
 
 
-# FIXME: not needed for any logic inside the driver. Only needed to meet the expectations set by parent class Controller.
-@dataclass
-class ZhPort:
-    name: str
-
-
 class Zurich(Controller):
     """Zurich driver main class."""
 

From 7b86f95b24b6bd79ad44a7c2652e452f0506be48 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 14:45:44 +0400
Subject: [PATCH 09/45] simplify creation of list of logical signals

---
 src/qibolab/instruments/zhinst.py | 14 +-------------
 1 file changed, 1 insertion(+), 13 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 9cbb192e6..d485608dd 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -663,19 +663,7 @@ def create_exp(self, qubits, couplers, options):
         """Zurich experiment initialization using their Experiment class."""
 
         # Setting experiment signal lines
-        signals = []
-        for coupler in couplers.values():
-            signals.append(lo.ExperimentSignal(f"couplerflux{coupler.name}"))
-
-        for qubit in qubits.values():
-            q = qubit.name  # pylint: disable=C0103
-            if len(self.sequence[f"drive{q}"]) != 0:
-                signals.append(lo.ExperimentSignal(f"drive{q}"))
-            if qubit.flux is not None:
-                signals.append(lo.ExperimentSignal(f"flux{q}"))
-            if len(self.sequence[f"readout{q}"]) != 0:
-                signals.append(lo.ExperimentSignal(f"measure{q}"))
-                signals.append(lo.ExperimentSignal(f"acquire{q}"))
+        signals = [lo.ExperimentSignal(name) for name in self.signal_map.keys()]
 
         exp = lo.Experiment(
             uid="Sequence",

From 0074cc22a6995d8ff3ffa4b281934093a8ab87d4 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 8 Feb 2024 12:32:30 +0000
Subject: [PATCH 10/45] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/qibolab/instruments/zhinst.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index d485608dd..a0997becf 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -23,7 +23,6 @@
 from .abstract import Controller
 from .port import Port
 
-
 SAMPLING_RATE = 2
 NANO_TO_SECONDS = 1e-9
 COMPILER_SETTINGS = {
@@ -201,8 +200,7 @@ def add_sweeper(self, sweeper, qubit):
 
 class ZhSweeperLine:
     """Zurich sweeper from qibolab sweeper for non pulse parameters Bias, Delay
-    (, power_range, local_oscillator frequency, offset ???)
-    """
+    (, power_range, local_oscillator frequency, offset ???)"""
 
     def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
         self.sweeper = sweeper
@@ -268,9 +266,7 @@ class Zurich(Controller):
 
     PortType = ZhPort
 
-    def __init__(
-        self, name, device_setup, time_of_flight=0.0, smearing=0.0
-    ):
+    def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
         super().__init__(name, None)
 
         self.signal_map = {}

From 20530f4d661a078f693f4c80ed49e3c319f52ac9 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 8 Feb 2024 17:01:45 +0400
Subject: [PATCH 11/45] use channel names as signal names, instead of inventing
 new name

---
 src/qibolab/instruments/zhinst.py | 242 ++++++++++++++++++------------
 1 file changed, 147 insertions(+), 95 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index a0997becf..3e2c94648 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -48,6 +48,24 @@
 SWEEPER_START = {"start"}
 
 
+def measure_signal_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's measure signal (logical) line.
+
+    FIXME: We cannot use channel name directly, because currently channels are named after wires, and due to multiplexed readout
+    multiple qubits have the same channel name for their readout. Should be fixed once channels are refactored.
+    """
+    return f"{qubit.readout.name}_{qubit.name}"
+
+
+def acquire_signal_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's acquire signal (logical) line.
+
+    FIXME: We cannot use acquire channel name, because qibolab does not have a concept of acquire channel. This function shall be removed
+    once all channel refactoring is done.
+    """
+    return f"acquire{qubit.name}"
+
+
 def select_pulse(pulse, pulse_type):
     """Pulse translation."""
 
@@ -216,7 +234,7 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
                     channel=qubit.flux.name,
                     qubit=qubit.name,
                 )
-                self.signal = f"flux{qubit.name}"
+                self.signal = qubit.flux.name
             if isinstance(qubit, Coupler):
                 pulse = CouplerFluxPulse(
                     start=0,
@@ -226,7 +244,7 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
                     channel=qubit.flux.name,
                     qubit=qubit.name,
                 )
-                self.signal = f"couplerflux{qubit.name}"
+                self.signal = qubit.flux.name
 
             self.pulse = pulse
 
@@ -239,7 +257,7 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
         elif sweeper.parameter is Parameter.start:
             if pulse:
                 self.pulse = pulse
-                self.signal = f"flux{qubit}"
+                self.signal = qubit.flux.name
 
                 self.zhpulse = ZhPulse(pulse).zhpulse
 
@@ -336,13 +354,13 @@ def calibration_step(self, qubits, couplers, options):
         for qubit in qubits.values():
             if qubit.flux is not None:
                 self.register_flux_line(qubit)
-            if len(self.sequence[f"drive{qubit.name}"]) != 0:
+            if len(self.sequence[qubit.drive.name]) != 0:
                 self.register_drive_line(
                     qubit=qubit,
                     intermediate_frequency=qubit.drive_frequency
                     - qubit.drive.local_oscillator.frequency,
                 )
-            if len(self.sequence[f"readout{qubit.name}"]) != 0:
+            if len(self.sequence[measure_signal_name(qubit)]) != 0:
                 self.register_readout_line(
                     qubit=qubit,
                     intermediate_frequency=qubit.readout_frequency
@@ -372,9 +390,11 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
         """
 
         q = qubit.name  # pylint: disable=C0103
-        self.signal_map[f"measure{q}"] = self.device_setup.logical_signal_groups[
-            f"q{q}"
-        ].logical_signals["measure_line"]
+        self.signal_map[measure_signal_name(qubit)] = (
+            self.device_setup.logical_signal_groups[f"q{q}"].logical_signals[
+                "measure_line"
+            ]
+        )
         self.calibration[f"/logical_signal_groups/q{q}/measure_line"] = (
             lo.SignalCalibration(
                 oscillator=lo.Oscillator(
@@ -391,9 +411,11 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
             )
         )
 
-        self.signal_map[f"acquire{q}"] = self.device_setup.logical_signal_groups[
-            f"q{q}"
-        ].logical_signals["acquire_line"]
+        self.signal_map[acquire_signal_name(qubit)] = (
+            self.device_setup.logical_signal_groups[f"q{q}"].logical_signals[
+                "acquire_line"
+            ]
+        )
 
         oscillator = lo.Oscillator(
             frequency=intermediate_frequency,
@@ -421,7 +443,7 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
     def register_drive_line(self, qubit, intermediate_frequency):
         """Registers qubit drive line to calibration and signal map."""
         q = qubit.name  # pylint: disable=C0103
-        self.signal_map[f"drive{q}"] = self.device_setup.logical_signal_groups[
+        self.signal_map[qubit.drive.name] = self.device_setup.logical_signal_groups[
             f"q{q}"
         ].logical_signals["drive_line"]
         self.calibration[f"/logical_signal_groups/q{q}/drive_line"] = (
@@ -443,7 +465,7 @@ def register_drive_line(self, qubit, intermediate_frequency):
     def register_flux_line(self, qubit):
         """Registers qubit flux line to calibration and signal map."""
         q = qubit.name  # pylint: disable=C0103
-        self.signal_map[f"flux{q}"] = self.device_setup.logical_signal_groups[
+        self.signal_map[qubit.flux.name] = self.device_setup.logical_signal_groups[
             f"q{q}"
         ].logical_signals["flux_line"]
         self.calibration[f"/logical_signal_groups/q{q}/flux_line"] = (
@@ -458,7 +480,7 @@ def register_flux_line(self, qubit):
     def register_couplerflux_line(self, coupler):
         """Registers qubit flux line to calibration and signal map."""
         c = coupler.name  # pylint: disable=C0103
-        self.signal_map[f"couplerflux{c}"] = self.device_setup.logical_signal_groups[
+        self.signal_map[coupler.flux.name] = self.device_setup.logical_signal_groups[
             f"qc{c}"
         ].logical_signals["flux_line"]
         self.calibration[f"/logical_signal_groups/qc{c}/flux_line"] = (
@@ -505,9 +527,24 @@ def create_sub_sequence(
             quantum_elements (dict[str, Qubit]|dict[str, Coupler]): qubits or couplers for the platform.
         """
         for quantum_element in quantum_elements.values():
-            q = quantum_element.name  # pylint: disable=C0103
-            measurements = self.sequence[f"readout{q}"]
-            pulses = self.sequence[f"{line_name}{q}"]
+            if isinstance(quantum_element, Coupler):
+                # Find arbitrary measurement from sequence and split coupler sequence based on that
+                # FIXME: neither the previous nor this way of doing this is not right - division into subsequences should happen
+                # for the entire sequence as a whole and not channel by channel (which is the source of prblem here, that we need to
+                # decide which qubit's measurement to take as baseline in splitting)
+                ch = next(
+                    iter(
+                        ch
+                        for ch, pulses in self.sequence.items()
+                        if any(p.pulse.type == PulseType.READOUT for p in pulses)
+                    )
+                )
+                measurements = self.sequence[ch]
+            else:
+                measurements = self.sequence[measure_signal_name(quantum_element)]
+
+            channel_name = getattr(quantum_element, line_name).name
+            pulses = self.sequence[channel_name]
             pulse_sequences = [[] for _ in measurements]
             pulse_sequences.append([])
             measurement_index = 0
@@ -516,7 +553,7 @@ def create_sub_sequence(
                     if pulse.pulse.finish > measurements[measurement_index].pulse.start:
                         measurement_index += 1
                 pulse_sequences[measurement_index].append(pulse)
-            self.sub_sequences[f"{line_name}{q}"] = pulse_sequences
+            self.sub_sequences[channel_name] = pulse_sequences
 
     def create_sub_sequences(
         self, qubits: Dict[str, Qubit], couplers: Dict[str, Coupler]
@@ -530,7 +567,7 @@ def create_sub_sequences(
         self.sub_sequences = {}
         self.create_sub_sequence("drive", qubits)
         self.create_sub_sequence("flux", qubits)
-        self.create_sub_sequence("couplerflux", couplers)
+        self.create_sub_sequence("flux", couplers)
 
     def experiment_flow(
         self,
@@ -569,8 +606,8 @@ def play(self, qubits, couplers, sequence, options):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            if len(self.sequence[f"readout{q}"]) != 0:
-                for i, ropulse in enumerate(self.sequence[f"readout{q}"]):
+            if len(self.sequence[measure_signal_name(qubit)]) != 0:
+                for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
                     data = np.array(self.results.get_data(f"sequence{q}_{i}"))
                     if options.acquisition_type is AcquisitionType.DISCRIMINATION:
                         data = (
@@ -590,7 +627,12 @@ def sequence_zh(self, sequence, qubits, couplers):
 
         # Fill the sequences with pulses according to their lines in temporal order
         for pulse in sequence:
-            zhsequence[f"{pulse.type.name.lower()}{pulse.qubit}"].append(ZhPulse(pulse))
+            if pulse.type == PulseType.READOUT:
+                zhsequence[measure_signal_name(qubits[pulse.qubit])].append(
+                    ZhPulse(pulse)
+                )
+            else:
+                zhsequence[pulse.channel].append(ZhPulse(pulse))
 
         def nt_loop(sweeper):
             if not self.nt_sweeps:
@@ -602,7 +644,10 @@ def nt_loop(sweeper):
         for sweeper in self.sweepers.copy():
             if sweeper.parameter.name in SWEEPER_SET:
                 for pulse in sweeper.pulses:
-                    aux_list = zhsequence[f"{pulse.type.name.lower()}{pulse.qubit}"]
+                    if pulse.type == PulseType.READOUT:
+                        aux_list = zhsequence[measure_signal_name(qubits[pulse.qubit])]
+                    else:
+                        aux_list = zhsequence[pulse.channel]
                     if (
                         sweeper.parameter is Parameter.frequency
                         and pulse.type is PulseType.READOUT
@@ -643,13 +688,19 @@ def nt_loop(sweeper):
             # This may not place the Zhsweeper when the start occurs among different sections or lines
             if sweeper.parameter.name in SWEEPER_START:
                 pulse = sweeper.pulses[0]
-                aux_list = zhsequence[f"{pulse.type.name.lower()}{pulse.qubit}"]
+
+                if pulse.type == PulseType.READOUT:
+                    aux_list = zhsequence[measure_signal_name(qubits[pulse.qubit])]
+                else:
+                    aux_list = zhsequence[pulse.channel]
                 for element in aux_list:
                     if pulse == element.pulse:
                         if isinstance(aux_list[aux_list.index(element)], ZhPulse):
                             aux_list.insert(
                                 aux_list.index(element),
-                                ZhSweeperLine(sweeper, pulse.qubit, sequence, pulse),
+                                ZhSweeperLine(
+                                    sweeper, qubits[pulse.qubit], sequence, pulse
+                                ),
                             )
                             break
 
@@ -702,16 +753,9 @@ def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
 
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
-        if "coupler" in str(self.sequence):
-            self.couplerflux(exp, couplers)
-        if "drive" in str(self.sequence):
-            if "flux" in str(self.sequence):
-                self.flux(exp, qubits)
-                self.drive(exp, qubits)
-            else:
-                self.drive(exp, qubits)
-        elif "flux" in str(self.sequence):
-            self.flux(exp, qubits)
+        self.drive(exp, qubits)
+        self.flux(exp, qubits)
+        self.couplerflux(exp, couplers)
         self.measure_relax(
             exp,
             qubits,
@@ -721,40 +765,42 @@ def select_exp(self, exp, qubits, couplers, exp_options):
         )
 
     @staticmethod
-    def play_sweep_select_single(exp, qubit, pulse, section, parameters, partial_sweep):
+    def play_sweep_select_single(
+        exp, qubit, pulse, channel_name, parameters, partial_sweep
+    ):
         """Play Zurich pulse when a single sweeper is involved."""
         if any("amplitude" in param for param in parameters):
             pulse.zhpulse.amplitude *= max(pulse.zhsweeper.values)
             pulse.zhsweeper.values /= max(pulse.zhsweeper.values)
             exp.play(
-                signal=f"{section}{qubit.name}",
+                signal=channel_name,
                 pulse=pulse.zhpulse,
                 amplitude=pulse.zhsweeper,
                 phase=pulse.pulse.relative_phase,
             )
         elif any("duration" in param for param in parameters):
             exp.play(
-                signal=f"{section}{qubit.name}",
+                signal=channel_name,
                 pulse=pulse.zhpulse,
                 length=pulse.zhsweeper,
                 phase=pulse.pulse.relative_phase,
             )
         elif any("relative_phase" in param for param in parameters):
             exp.play(
-                signal=f"{section}{qubit.name}",
+                signal=channel_name,
                 pulse=pulse.zhpulse,
                 phase=pulse.zhsweeper,
             )
         elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
             exp.play(
-                signal=f"{section}{qubit.name}",
+                signal=channel_name,
                 pulse=pulse.zhpulse,
                 phase=pulse.pulse.relative_phase,
             )
 
     # Hardcoded for the flux pulse for 2q gates
     @staticmethod
-    def play_sweep_select_dual(exp, qubit, pulse, section, parameters):
+    def play_sweep_select_dual(exp, qubit, pulse, channel_name, parameters):
         """Play Zurich pulse when two sweepers are involved on the same
         pulse."""
         if "amplitude" in parameters and "duration" in parameters:
@@ -768,7 +814,7 @@ def play_sweep_select_dual(exp, qubit, pulse, section, parameters):
                     sweeper_dur_index = pulse.zhsweepers.index(sweeper)
 
             exp.play(
-                signal=f"{section}{qubit.name}",
+                signal=channel_name,
                 pulse=pulse.zhpulse,
                 amplitude=pulse.zhsweepers[sweeper_amp_index],
                 length=pulse.zhsweepers[sweeper_dur_index],
@@ -777,11 +823,14 @@ def play_sweep_select_dual(exp, qubit, pulse, section, parameters):
     def play_sweep(self, exp, qubit, pulse, section):
         """Takes care of playing the sweepers and involved pulses for different
         options."""
-
+        if section == "readout":
+            channel_name = measure_signal_name(qubit)
+        else:
+            channel_name = getattr(qubit, section).name
         if isinstance(pulse, ZhSweeperLine):
             if pulse.zhsweeper.uid == "bias":
                 exp.play(
-                    signal=f"{section}{qubit.name}",
+                    signal=channel_name,
                     pulse=pulse.zhpulse,
                     amplitude=pulse.zhsweeper,
                 )
@@ -791,10 +840,10 @@ def play_sweep(self, exp, qubit, pulse, section):
                 parameters.append(partial_sweep.uid)
             # Recheck partial sweeps
             if len(parameters) == 2:
-                self.play_sweep_select_dual(exp, qubit, pulse, section, parameters)
+                self.play_sweep_select_dual(exp, qubit, pulse, channel_name, parameters)
             else:
                 self.play_sweep_select_single(
-                    exp, qubit, pulse, section, parameters, partial_sweep
+                    exp, qubit, pulse, channel_name, parameters, partial_sweep
                 )
 
     def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
@@ -805,27 +854,27 @@ def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
             couplers (dict[str, Coupler]): coupler on which pulses are played.
         """
         for coupler in couplers.values():
-            c = coupler.name  # pylint: disable=C0103
+            channel_name = coupler.flux.name
             time = 0
             previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[f"couplerflux{c}"]):
-                section_uid = f"sequence_couplerflux{c}_{i}"
+            for i, sequence in enumerate(self.sub_sequences[channel_name]):
+                section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
                         pulse.zhpulse.uid += f"{i}_{j}"
                         exp.delay(
-                            signal=f"couplerflux{c}",
+                            signal=channel_name,
                             time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
                         )
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
                         if isinstance(pulse, ZhSweeperLine):
-                            self.play_sweep(exp, coupler, pulse, section="couplerflux")
+                            self.play_sweep(exp, coupler, pulse, section="flux")
                         elif isinstance(pulse, ZhSweeper):
-                            self.play_sweep(exp, coupler, pulse, section="couplerflux")
+                            self.play_sweep(exp, coupler, pulse, section="flux")
                         elif isinstance(pulse, ZhPulse):
-                            exp.play(signal=f"couplerflux{c}", pulse=pulse.zhpulse)
+                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
                 previous_section = section_uid
 
     def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
@@ -837,16 +886,17 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
         """
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
+            channel_name = qubit.flux.name
             time = 0
             previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[f"flux{q}"]):
-                section_uid = f"sequence_flux{q}_{i}"
+            for i, sequence in enumerate(self.sub_sequences[qubit.flux.name]):
+                section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
                         if not isinstance(pulse, ZhSweeperLine):
                             pulse.zhpulse.uid += f"{i}_{j}"
                             exp.delay(
-                                signal=f"flux{q}",
+                                signal=qubit.flux.name,
                                 time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                                 - time,
                             )
@@ -858,7 +908,7 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                         elif isinstance(pulse, ZhSweeper):
                             self.play_sweep(exp, qubit, pulse, section="flux")
                         elif isinstance(pulse, ZhPulse):
-                            exp.play(signal=f"flux{q}", pulse=pulse.zhpulse)
+                            exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
                 previous_section = section_uid
 
     def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
@@ -869,16 +919,16 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
             qubits (dict[str, Qubit]): qubits on which pulses are played.
         """
         for qubit in qubits.values():
-            q = qubit.name  # pylint: disable=C0103
+            channel_name = qubit.drive.name
             time = 0
             previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[f"drive{q}"]):
-                section_uid = f"sequence_drive{q}_{i}"
+            for i, sequence in enumerate(self.sub_sequences[qubit.drive.name]):
+                section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
                         if not isinstance(pulse, ZhSweeperLine):
                             exp.delay(
-                                signal=f"drive{q}",
+                                signal=qubit.drive.name,
                                 time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                                 - time,
                             )
@@ -890,22 +940,24 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                                 self.play_sweep(exp, qubit, pulse, section="drive")
                             elif isinstance(pulse, ZhPulse):
                                 exp.play(
-                                    signal=f"drive{q}",
+                                    signal=qubit.drive.name,
                                     pulse=pulse.zhpulse,
                                     phase=pulse.pulse.relative_phase,
                                 )
                         elif isinstance(pulse, ZhSweeperLine):
-                            exp.delay(signal=f"drive{q}", time=pulse.zhsweeper)
+                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweeper)
 
-                    if len(self.sequence[f"readout{q}"]) > 0 and isinstance(
-                        self.sequence[f"readout{q}"][0], ZhSweeperLine
+                    if len(
+                        self.sequence[measure_signal_name(qubit)]
+                    ) > 0 and isinstance(
+                        self.sequence[measure_signal_name(qubit)][0], ZhSweeperLine
                     ):
                         exp.delay(
-                            signal=f"drive{q}",
-                            time=self.sequence[f"readout{q}"][0].zhsweeper,
+                            signal=qubit.drive.name,
+                            time=self.sequence[measure_signal_name(qubit)][0].zhsweeper,
                         )
-                        self.sequence[f"readout{q}"].remove(
-                            self.sequence[f"readout{q}"][0]
+                        self.sequence[measure_signal_name(qubit)].remove(
+                            self.sequence[measure_signal_name(qubit)][0]
                         )
 
                 previous_section = section_uid
@@ -914,16 +966,15 @@ def find_subsequence_finish(
         self,
         measurement_number: int,
         line: str,
-        quantum_elements: Union[Dict[str, Qubit], Dict[str, Coupler]],
+        quantum_elements: Union[Qubit, Coupler],
     ) -> Tuple[int, str]:
         """Find the finishing time and qubit for a given sequence.
 
         Args:
             measurement_number (int): number of the measure pulse.
             line (str): line from which measure the finishing time.
-                e.g.: "drive", "flux", "couplerflux"
-            quantum_elements (dict[str, Qubit]|dict[str, Coupler]): qubits or couplers from
-                which measure the finishing time.
+                e.g.: "drive", "flux"
+            quantum_elements: qubits or couplers from which to measure the finishing time.
 
         Returns:
             time_finish (int): Finish time of the last pulse of the subsequence
@@ -934,17 +985,13 @@ def find_subsequence_finish(
         time_finish = 0
         sequence_finish = "None"
         for quantum_element in quantum_elements:
-            if (
-                len(self.sub_sequences[f"{line}{quantum_element}"])
-                <= measurement_number
-            ):
+            channel_name = getattr(quantum_element, line).name
+            if len(self.sub_sequences[channel_name]) <= measurement_number:
                 continue
-            for pulse in self.sub_sequences[f"{line}{quantum_element}"][
-                measurement_number
-            ]:
+            for pulse in self.sub_sequences[channel_name][measurement_number]:
                 if pulse.pulse.finish > time_finish:
                     time_finish = pulse.pulse.finish
-                    sequence_finish = f"{line}{quantum_element}"
+                    sequence_finish = channel_name
         return time_finish, sequence_finish
 
     # For pulsed spectroscopy, set integration_length and either measure_pulse or measure_pulse_length.
@@ -956,12 +1003,12 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
         qubit_readout_schedule = defaultdict(list)
         iq_angle_readout_schedule = defaultdict(list)
         for qubit in qubits.values():
-            q = qubit.name  # pylint: disable=C0103
             iq_angle = qubit.iq_angle
-            if len(self.sequence[f"readout{q}"]) != 0:
-                for i, pulse in enumerate(self.sequence[f"readout{q}"]):
+            channel_name = measure_signal_name(qubit)
+            if len(self.sequence[channel_name]) != 0:
+                for i, pulse in enumerate(self.sequence[channel_name]):
                     readout_schedule[i].append(pulse)
-                    qubit_readout_schedule[i].append(q)
+                    qubit_readout_schedule[i].append(qubit)
                     iq_angle_readout_schedule[i].append(iq_angle)
 
         weights = {}
@@ -974,7 +1021,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
         ):
             qd_finish = self.find_subsequence_finish(i, "drive", qubits_readout)
             qf_finish = self.find_subsequence_finish(i, "flux", qubits_readout)
-            cf_finish = self.find_subsequence_finish(i, "couplerflux", couplers)
+            cf_finish = self.find_subsequence_finish(i, "flux", couplers.values())
             finish_times = np.array(
                 [
                     qd_finish,
@@ -990,11 +1037,12 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                 play_after = f"sequence_{latest_sequence['line']}_{i}"
             # Section on the outside loop allows for multiplex
             with exp.section(uid=f"sequence_measure_{i}", play_after=play_after):
-                for pulse, q, iq_angle in zip(pulses, qubits_readout, iq_angles):
+                for pulse, qubit, iq_angle in zip(pulses, qubits_readout, iq_angles):
+                    q = qubit.name
                     pulse.zhpulse.uid += str(i)
 
                     exp.delay(
-                        signal=f"acquire{q}",
+                        signal=acquire_signal_name(qubit),
                         time=self.smearing * NANO_TO_SECONDS,
                     )
 
@@ -1041,18 +1089,18 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
 
                     measure_pulse_parameters = {"phase": 0}
 
-                    if i == len(self.sequence[f"readout{q}"]) - 1:
+                    if i == len(self.sequence[measure_signal_name(qubit)]) - 1:
                         reset_delay = relaxation_time * NANO_TO_SECONDS
                     else:
                         reset_delay = 0
 
                     exp.measure(
-                        acquire_signal=f"acquire{q}",
+                        acquire_signal=acquire_signal_name(qubit),
                         handle=f"sequence{q}_{i}",
                         integration_kernel=weight,
                         integration_kernel_parameters=None,
                         integration_length=None,
-                        measure_signal=f"measure{q}",
+                        measure_signal=measure_signal_name(qubit),
                         measure_pulse=pulse.zhpulse,
                         measure_pulse_length=round(
                             pulse.pulse.duration * NANO_TO_SECONDS, 9
@@ -1111,8 +1159,8 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            if len(self.sequence[f"readout{q}"]) != 0:
-                for i, ropulse in enumerate(self.sequence[f"readout{q}"]):
+            if len(self.sequence[measure_signal_name(qubit)]) != 0:
+                for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
                     exp_res = self.results.get_data(f"sequence{q}_{i}")
                     # if using singleshot, the first axis contains shots,
                     # i.e.: (nshots, sweeper_1, sweeper_2)
@@ -1145,12 +1193,16 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
 
         if sweeper.parameter is Parameter.frequency:
             for pulse in sweeper.pulses:
-                line = "drive" if pulse.type is PulseType.DRIVE else "measure"
+                line = "drive" if pulse.type is PulseType.DRIVE else "readout"
                 zhsweeper = ZhSweeper(
                     pulse, sweeper, qubits[sweeper.pulses[0].qubit]
                 ).zhsweeper
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
-                exp_calib[f"{line}{pulse.qubit}"] = lo.SignalCalibration(
+                if line == "readout":
+                    channel_name = measure_signal_name(qubits[pulse.qubit])
+                else:
+                    channel_name = getattr(qubits[pulse.qubit], line).name
+                exp_calib[channel_name] = lo.SignalCalibration(
                     oscillator=lo.Oscillator(
                         frequency=zhsweeper,
                         modulation_type=lo.ModulationType.HARDWARE,

From 7c26555c2c99507cab191aee80b78a775c5467b0 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 13 Feb 2024 12:53:02 +0400
Subject: [PATCH 12/45] remove unnecessary if statements

---
 src/qibolab/instruments/zhinst.py | 65 ++++++++++++++-----------------
 1 file changed, 30 insertions(+), 35 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 3e2c94648..c9dbfc72f 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -606,16 +606,15 @@ def play(self, qubits, couplers, sequence, options):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            if len(self.sequence[measure_signal_name(qubit)]) != 0:
-                for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
-                    data = np.array(self.results.get_data(f"sequence{q}_{i}"))
-                    if options.acquisition_type is AcquisitionType.DISCRIMINATION:
-                        data = (
-                            np.ones(data.shape) - data.real
-                        )  # Probability inversion patch
-                    serial = ropulse.pulse.serial
-                    qubit = ropulse.pulse.qubit
-                    results[serial] = results[qubit] = options.results_type(data)
+            for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
+                data = np.array(self.results.get_data(f"sequence{q}_{i}"))
+                if options.acquisition_type is AcquisitionType.DISCRIMINATION:
+                    data = (
+                        np.ones(data.shape) - data.real
+                    )  # Probability inversion patch
+                serial = ropulse.pulse.serial
+                qubit = ropulse.pulse.qubit
+                results[serial] = results[qubit] = options.results_type(data)
 
         return results
 
@@ -1005,11 +1004,10 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
         for qubit in qubits.values():
             iq_angle = qubit.iq_angle
             channel_name = measure_signal_name(qubit)
-            if len(self.sequence[channel_name]) != 0:
-                for i, pulse in enumerate(self.sequence[channel_name]):
-                    readout_schedule[i].append(pulse)
-                    qubit_readout_schedule[i].append(qubit)
-                    iq_angle_readout_schedule[i].append(iq_angle)
+            for i, pulse in enumerate(self.sequence[channel_name]):
+                readout_schedule[i].append(pulse)
+                qubit_readout_schedule[i].append(qubit)
+                iq_angle_readout_schedule[i].append(iq_angle)
 
         weights = {}
         for i, (pulses, qubits_readout, iq_angles) in enumerate(
@@ -1159,26 +1157,23 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            if len(self.sequence[measure_signal_name(qubit)]) != 0:
-                for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
-                    exp_res = self.results.get_data(f"sequence{q}_{i}")
-                    # if using singleshot, the first axis contains shots,
-                    # i.e.: (nshots, sweeper_1, sweeper_2)
-                    # if using integration: (sweeper_1, sweeper_2)
-                    if options.averaging_mode is AveragingMode.SINGLESHOT:
-                        rearranging_axes += 1
-                    # Reorder dimensions
-                    data = np.moveaxis(
-                        exp_res, rearranging_axes[0], rearranging_axes[1]
-                    )
-                    if options.acquisition_type is AcquisitionType.DISCRIMINATION:
-                        data = (
-                            np.ones(data.shape) - data.real
-                        )  # Probability inversion patch
-
-                    serial = ropulse.pulse.serial
-                    qubit = ropulse.pulse.qubit
-                    results[serial] = results[qubit] = options.results_type(data)
+            for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
+                exp_res = self.results.get_data(f"sequence{q}_{i}")
+                # if using singleshot, the first axis contains shots,
+                # i.e.: (nshots, sweeper_1, sweeper_2)
+                # if using integration: (sweeper_1, sweeper_2)
+                if options.averaging_mode is AveragingMode.SINGLESHOT:
+                    rearranging_axes += 1
+                # Reorder dimensions
+                data = np.moveaxis(exp_res, rearranging_axes[0], rearranging_axes[1])
+                if options.acquisition_type is AcquisitionType.DISCRIMINATION:
+                    data = (
+                        np.ones(data.shape) - data.real
+                    )  # Probability inversion patch
+
+                serial = ropulse.pulse.serial
+                qubit = ropulse.pulse.qubit
+                results[serial] = results[qubit] = options.results_type(data)
 
         return results
 

From fc87919255174131ef9d6bc81cc0466f97b2331e Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 13 Feb 2024 13:43:49 +0400
Subject: [PATCH 13/45] Fix bug when rearrange axes are incorrectly increased
 multiple times

---
 src/qibolab/instruments/zhinst.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index c9dbfc72f..6e5f1d644 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -1147,6 +1147,11 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         sweepers = list(sweepers)
         rearranging_axes, sweepers = self.rearrange_sweepers(sweepers)
         self.sweepers = sweepers
+        # if using singleshot, the first axis contains shots,
+        # i.e.: (nshots, sweeper_1, sweeper_2)
+        # if using integration: (sweeper_1, sweeper_2)
+        if options.averaging_mode is AveragingMode.SINGLESHOT:
+            rearranging_axes += 1
 
         self.frequency_from_pulses(qubits, sequence)
 
@@ -1159,11 +1164,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
             q = qubit.name  # pylint: disable=C0103
             for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
                 exp_res = self.results.get_data(f"sequence{q}_{i}")
-                # if using singleshot, the first axis contains shots,
-                # i.e.: (nshots, sweeper_1, sweeper_2)
-                # if using integration: (sweeper_1, sweeper_2)
-                if options.averaging_mode is AveragingMode.SINGLESHOT:
-                    rearranging_axes += 1
+
                 # Reorder dimensions
                 data = np.moveaxis(exp_res, rearranging_axes[0], rearranging_axes[1])
                 if options.acquisition_type is AcquisitionType.DISCRIMINATION:

From 7cab4b9ddbf39c343fe90e9123b46c3ed53203c8 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 14 Feb 2024 10:36:28 +0400
Subject: [PATCH 14/45] access iq_angle directly from qubit instead of
 collecting them into another dict

---
 src/qibolab/instruments/zhinst.py | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 6e5f1d644..3cc45ebf8 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -1000,21 +1000,17 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
         """Qubit readout pulse, data acquisition and qubit relaxation."""
         readout_schedule = defaultdict(list)
         qubit_readout_schedule = defaultdict(list)
-        iq_angle_readout_schedule = defaultdict(list)
         for qubit in qubits.values():
-            iq_angle = qubit.iq_angle
             channel_name = measure_signal_name(qubit)
             for i, pulse in enumerate(self.sequence[channel_name]):
                 readout_schedule[i].append(pulse)
                 qubit_readout_schedule[i].append(qubit)
-                iq_angle_readout_schedule[i].append(iq_angle)
 
         weights = {}
-        for i, (pulses, qubits_readout, iq_angles) in enumerate(
+        for i, (pulses, qubits_readout) in enumerate(
             zip(
                 readout_schedule.values(),
                 qubit_readout_schedule.values(),
-                iq_angle_readout_schedule.values(),
             )
         ):
             qd_finish = self.find_subsequence_finish(i, "drive", qubits_readout)
@@ -1035,7 +1031,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                 play_after = f"sequence_{latest_sequence['line']}_{i}"
             # Section on the outside loop allows for multiplex
             with exp.section(uid=f"sequence_measure_{i}", play_after=play_after):
-                for pulse, qubit, iq_angle in zip(pulses, qubits_readout, iq_angles):
+                for pulse, qubit in zip(pulses, qubits_readout):
                     q = qubit.name
                     pulse.zhpulse.uid += str(i)
 
@@ -1045,13 +1041,12 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                     )
 
                     if (
-                        qubits[q].kernel is not None
+                        qubit.kernel is not None
                         and acquisition_type == lo.AcquisitionType.DISCRIMINATION
                     ):
-                        kernel = qubits[q].kernel
                         weight = lo.pulse_library.sampled_pulse_complex(
                             uid="weight" + str(q),
-                            samples=kernel * np.exp(1j * iq_angle),
+                            samples=qubit.kernel * np.exp(1j * qubit.iq_angle),
                         )
 
                     else:
@@ -1066,7 +1061,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                                             )
                                         ]
                                     )
-                                    * np.exp(1j * iq_angle),
+                                    * np.exp(1j * qubit.iq_angle),
                                     uid="weights" + str(q),
                                 )
                                 weights[q] = weight

From cb1293d0fe901486e890f24e411fd1401ac0777a Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 14 Feb 2024 12:33:24 +0400
Subject: [PATCH 15/45] rename functions

---
 src/qibolab/instruments/zhinst.py | 52 ++++++++++++++++---------------
 1 file changed, 27 insertions(+), 25 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 3cc45ebf8..f864b1c64 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -48,8 +48,8 @@
 SWEEPER_START = {"start"}
 
 
-def measure_signal_name(qubit: Qubit) -> str:
-    """Construct and return a name for qubit's measure signal (logical) line.
+def measure_channel_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's measure channel.
 
     FIXME: We cannot use channel name directly, because currently channels are named after wires, and due to multiplexed readout
     multiple qubits have the same channel name for their readout. Should be fixed once channels are refactored.
@@ -57,8 +57,8 @@ def measure_signal_name(qubit: Qubit) -> str:
     return f"{qubit.readout.name}_{qubit.name}"
 
 
-def acquire_signal_name(qubit: Qubit) -> str:
-    """Construct and return a name for qubit's acquire signal (logical) line.
+def acquire_channel_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's acquire channel.
 
     FIXME: We cannot use acquire channel name, because qibolab does not have a concept of acquire channel. This function shall be removed
     once all channel refactoring is done.
@@ -360,7 +360,7 @@ def calibration_step(self, qubits, couplers, options):
                     intermediate_frequency=qubit.drive_frequency
                     - qubit.drive.local_oscillator.frequency,
                 )
-            if len(self.sequence[measure_signal_name(qubit)]) != 0:
+            if len(self.sequence[measure_channel_name(qubit)]) != 0:
                 self.register_readout_line(
                     qubit=qubit,
                     intermediate_frequency=qubit.readout_frequency
@@ -390,7 +390,7 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
         """
 
         q = qubit.name  # pylint: disable=C0103
-        self.signal_map[measure_signal_name(qubit)] = (
+        self.signal_map[measure_channel_name(qubit)] = (
             self.device_setup.logical_signal_groups[f"q{q}"].logical_signals[
                 "measure_line"
             ]
@@ -411,7 +411,7 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
             )
         )
 
-        self.signal_map[acquire_signal_name(qubit)] = (
+        self.signal_map[acquire_channel_name(qubit)] = (
             self.device_setup.logical_signal_groups[f"q{q}"].logical_signals[
                 "acquire_line"
             ]
@@ -541,7 +541,7 @@ def create_sub_sequence(
                 )
                 measurements = self.sequence[ch]
             else:
-                measurements = self.sequence[measure_signal_name(quantum_element)]
+                measurements = self.sequence[measure_channel_name(quantum_element)]
 
             channel_name = getattr(quantum_element, line_name).name
             pulses = self.sequence[channel_name]
@@ -606,7 +606,7 @@ def play(self, qubits, couplers, sequence, options):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
+            for i, ropulse in enumerate(self.sequence[measure_channel_name(qubit)]):
                 data = np.array(self.results.get_data(f"sequence{q}_{i}"))
                 if options.acquisition_type is AcquisitionType.DISCRIMINATION:
                     data = (
@@ -627,7 +627,7 @@ def sequence_zh(self, sequence, qubits, couplers):
         # Fill the sequences with pulses according to their lines in temporal order
         for pulse in sequence:
             if pulse.type == PulseType.READOUT:
-                zhsequence[measure_signal_name(qubits[pulse.qubit])].append(
+                zhsequence[measure_channel_name(qubits[pulse.qubit])].append(
                     ZhPulse(pulse)
                 )
             else:
@@ -644,7 +644,7 @@ def nt_loop(sweeper):
             if sweeper.parameter.name in SWEEPER_SET:
                 for pulse in sweeper.pulses:
                     if pulse.type == PulseType.READOUT:
-                        aux_list = zhsequence[measure_signal_name(qubits[pulse.qubit])]
+                        aux_list = zhsequence[measure_channel_name(qubits[pulse.qubit])]
                     else:
                         aux_list = zhsequence[pulse.channel]
                     if (
@@ -689,7 +689,7 @@ def nt_loop(sweeper):
                 pulse = sweeper.pulses[0]
 
                 if pulse.type == PulseType.READOUT:
-                    aux_list = zhsequence[measure_signal_name(qubits[pulse.qubit])]
+                    aux_list = zhsequence[measure_channel_name(qubits[pulse.qubit])]
                 else:
                     aux_list = zhsequence[pulse.channel]
                 for element in aux_list:
@@ -823,7 +823,7 @@ def play_sweep(self, exp, qubit, pulse, section):
         """Takes care of playing the sweepers and involved pulses for different
         options."""
         if section == "readout":
-            channel_name = measure_signal_name(qubit)
+            channel_name = measure_channel_name(qubit)
         else:
             channel_name = getattr(qubit, section).name
         if isinstance(pulse, ZhSweeperLine):
@@ -947,16 +947,18 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             exp.delay(signal=qubit.drive.name, time=pulse.zhsweeper)
 
                     if len(
-                        self.sequence[measure_signal_name(qubit)]
+                        self.sequence[measure_channel_name(qubit)]
                     ) > 0 and isinstance(
-                        self.sequence[measure_signal_name(qubit)][0], ZhSweeperLine
+                        self.sequence[measure_channel_name(qubit)][0], ZhSweeperLine
                     ):
                         exp.delay(
                             signal=qubit.drive.name,
-                            time=self.sequence[measure_signal_name(qubit)][0].zhsweeper,
+                            time=self.sequence[measure_channel_name(qubit)][
+                                0
+                            ].zhsweeper,
                         )
-                        self.sequence[measure_signal_name(qubit)].remove(
-                            self.sequence[measure_signal_name(qubit)][0]
+                        self.sequence[measure_channel_name(qubit)].remove(
+                            self.sequence[measure_channel_name(qubit)][0]
                         )
 
                 previous_section = section_uid
@@ -1001,7 +1003,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
         readout_schedule = defaultdict(list)
         qubit_readout_schedule = defaultdict(list)
         for qubit in qubits.values():
-            channel_name = measure_signal_name(qubit)
+            channel_name = measure_channel_name(qubit)
             for i, pulse in enumerate(self.sequence[channel_name]):
                 readout_schedule[i].append(pulse)
                 qubit_readout_schedule[i].append(qubit)
@@ -1036,7 +1038,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                     pulse.zhpulse.uid += str(i)
 
                     exp.delay(
-                        signal=acquire_signal_name(qubit),
+                        signal=acquire_channel_name(qubit),
                         time=self.smearing * NANO_TO_SECONDS,
                     )
 
@@ -1082,18 +1084,18 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
 
                     measure_pulse_parameters = {"phase": 0}
 
-                    if i == len(self.sequence[measure_signal_name(qubit)]) - 1:
+                    if i == len(self.sequence[measure_channel_name(qubit)]) - 1:
                         reset_delay = relaxation_time * NANO_TO_SECONDS
                     else:
                         reset_delay = 0
 
                     exp.measure(
-                        acquire_signal=acquire_signal_name(qubit),
+                        acquire_signal=acquire_channel_name(qubit),
                         handle=f"sequence{q}_{i}",
                         integration_kernel=weight,
                         integration_kernel_parameters=None,
                         integration_length=None,
-                        measure_signal=measure_signal_name(qubit),
+                        measure_signal=measure_channel_name(qubit),
                         measure_pulse=pulse.zhpulse,
                         measure_pulse_length=round(
                             pulse.pulse.duration * NANO_TO_SECONDS, 9
@@ -1157,7 +1159,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         results = {}
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
-            for i, ropulse in enumerate(self.sequence[measure_signal_name(qubit)]):
+            for i, ropulse in enumerate(self.sequence[measure_channel_name(qubit)]):
                 exp_res = self.results.get_data(f"sequence{q}_{i}")
 
                 # Reorder dimensions
@@ -1190,7 +1192,7 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 ).zhsweeper
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
                 if line == "readout":
-                    channel_name = measure_signal_name(qubits[pulse.qubit])
+                    channel_name = measure_channel_name(qubits[pulse.qubit])
                 else:
                     channel_name = getattr(qubits[pulse.qubit], line).name
                 exp_calib[channel_name] = lo.SignalCalibration(

From a9ad17c9ff3ca4b4cebfbdf0be08b562af3443ad Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 14 Feb 2024 15:28:04 +0400
Subject: [PATCH 16/45] remove unused attribute .signal

---
 src/qibolab/instruments/zhinst.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index f864b1c64..245e5dadb 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -144,8 +144,6 @@ def __init__(self, pulse):
         """Zurich pulse from qibolab pulse."""
         self.pulse = pulse
         """Qibolab pulse."""
-        self.signal = f"{pulse.type.name.lower()}{pulse.qubit}"
-        """Line associated with the pulse."""
         self.zhpulse = select_pulse(pulse, pulse.type.name.lower())
         """Zurich pulse."""
 
@@ -160,8 +158,7 @@ def __init__(self, pulse, sweeper, qubit):
 
         self.pulse = pulse
         """Qibolab pulse associated to the sweeper."""
-        self.signal = f"{pulse.type.name.lower()}{pulse.qubit}"
-        """Line associated with the pulse."""
+
         self.zhpulse = ZhPulse(pulse).zhpulse
         """Zurich pulse associated to the sweeper."""
 
@@ -234,7 +231,6 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
                     channel=qubit.flux.name,
                     qubit=qubit.name,
                 )
-                self.signal = qubit.flux.name
             if isinstance(qubit, Coupler):
                 pulse = CouplerFluxPulse(
                     start=0,
@@ -244,7 +240,6 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
                     channel=qubit.flux.name,
                     qubit=qubit.name,
                 )
-                self.signal = qubit.flux.name
 
             self.pulse = pulse
 
@@ -257,8 +252,6 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
         elif sweeper.parameter is Parameter.start:
             if pulse:
                 self.pulse = pulse
-                self.signal = qubit.flux.name
-
                 self.zhpulse = ZhPulse(pulse).zhpulse
 
         # Need something better to store multiple sweeps on the same pulse

From 66e4cbc93cc7aca5b58416cf9e1951fd11629db4 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 14 Feb 2024 16:15:25 +0400
Subject: [PATCH 17/45] simplify handling of sweepers and sweep parameters

---
 src/qibolab/instruments/zhinst.py | 132 +++++++++++++-----------------
 1 file changed, 58 insertions(+), 74 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 245e5dadb..00c5d9a76 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -130,6 +130,23 @@ def select_pulse(pulse, pulse_type):
         )
 
 
+def classify_sweepers(sweepers):
+    """"""
+    nt_sweepers = []
+    rt_sweepers = []
+    for sweeper in sweepers:
+        if (
+            sweeper.parameter is Parameter.amplitude
+            and sweeper.pulses[0].type is PulseType.READOUT
+        ):
+            nt_sweepers.append(sweeper)
+        elif sweeper.parameter.name in SWEEPER_BIAS:
+            nt_sweepers.append(sweeper)
+        else:
+            rt_sweepers.append(sweeper)
+    return nt_sweepers, rt_sweepers
+
+
 @dataclass
 class ZhPort(Port):
     name: Tuple[str, str]
@@ -153,25 +170,16 @@ class ZhSweeper:
     Duration, Frequency (and maybe Phase)"""
 
     def __init__(self, pulse, sweeper, qubit):
-        self.sweeper = sweeper
-        """Qibolab sweeper."""
+        p = ZhPulse(pulse)
 
-        self.pulse = pulse
+        self.pulse = p.pulse
         """Qibolab pulse associated to the sweeper."""
 
-        self.zhpulse = ZhPulse(pulse).zhpulse
+        self.zhpulse = p.zhpulse
         """Zurich pulse associated to the sweeper."""
 
-        self.zhsweeper = self.select_sweeper(pulse.type, sweeper, qubit)
-        """Zurich sweeper."""
-
         self.zhsweepers = [self.select_sweeper(pulse.type, sweeper, qubit)]
-        """Zurich sweepers, Need something better to store multiple sweeps on
-        the same pulse.
-
-        Not properly implemented as it was only used on Rabi amplitude
-        vs lenght and it was an unused routine.
-        """
+        """Zurich sweepers."""
 
     @staticmethod  # pylint: disable=R0903
     def select_sweeper(ptype, sweeper, qubit):
@@ -218,9 +226,6 @@ class ZhSweeperLine:
     (, power_range, local_oscillator frequency, offset ???)"""
 
     def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
-        self.sweeper = sweeper
-        """Qibolab sweeper."""
-
         if sweeper.parameter is Parameter.bias:
             if isinstance(qubit, Qubit):
                 pulse = FluxPulse(
@@ -255,7 +260,7 @@ def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
                 self.zhpulse = ZhPulse(pulse).zhpulse
 
         # Need something better to store multiple sweeps on the same pulse
-        self.zhsweeper = self.select_sweeper(sweeper)
+        self.zhsweepers = [self.select_sweeper(sweeper)]
 
     @staticmethod  # pylint: disable=R0903
     def select_sweeper(sweeper):
@@ -314,7 +319,7 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
         "Sub sequences between each measurement"
 
         self.sweepers = []
-        self.nt_sweeps = None
+        self.nt_sweeps = []
         "Storing sweepers"
         # Improve the storing of multiple sweeps
 
@@ -626,14 +631,7 @@ def sequence_zh(self, sequence, qubits, couplers):
             else:
                 zhsequence[pulse.channel].append(ZhPulse(pulse))
 
-        def nt_loop(sweeper):
-            if not self.nt_sweeps:
-                self.nt_sweeps = [sweeper]
-            else:
-                self.nt_sweeps.append(sweeper)
-            self.sweepers.remove(sweeper)
-
-        for sweeper in self.sweepers.copy():
+        for sweeper in self.sweepers:
             if sweeper.parameter.name in SWEEPER_SET:
                 for pulse in sweeper.pulses:
                     if pulse.type == PulseType.READOUT:
@@ -650,7 +648,7 @@ def nt_loop(sweeper):
                         and pulse.type is PulseType.READOUT
                     ):
                         self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
-                        nt_loop(sweeper)
+
                     for element in aux_list:
                         if pulse == element.pulse:
                             if isinstance(aux_list[aux_list.index(element)], ZhPulse):
@@ -674,9 +672,6 @@ def nt_loop(sweeper):
                                         sweeper, qubits[pulse.qubit]
                                     )
 
-            if sweeper.parameter.name in SWEEPER_BIAS:
-                nt_loop(sweeper)
-
             # This may not place the Zhsweeper when the start occurs among different sections or lines
             if sweeper.parameter.name in SWEEPER_START:
                 pulse = sweeper.pulses[0]
@@ -697,6 +692,7 @@ def nt_loop(sweeper):
                             break
 
         self.sequence = zhsequence
+        self.nt_sweeps, self.sweepers = classify_sweepers(self.sweepers)
 
     def create_exp(self, qubits, couplers, options):
         """Zurich experiment initialization using their Experiment class."""
@@ -721,7 +717,7 @@ def create_exp(self, qubits, couplers, options):
 
         exp_calib = lo.Calibration()
         # Near Time recursion loop or directly to Real Time recursion loop
-        if self.nt_sweeps is not None:
+        if self.nt_sweeps:
             self.sweep_recursion_nt(qubits, couplers, exp_options, exp, exp_calib)
         else:
             self.define_exp(qubits, couplers, exp_options, exp, exp_calib)
@@ -762,26 +758,26 @@ def play_sweep_select_single(
     ):
         """Play Zurich pulse when a single sweeper is involved."""
         if any("amplitude" in param for param in parameters):
-            pulse.zhpulse.amplitude *= max(pulse.zhsweeper.values)
-            pulse.zhsweeper.values /= max(pulse.zhsweeper.values)
+            pulse.zhpulse.amplitude *= max(pulse.zhsweepers[0].values)
+            pulse.zhsweepers[0].values /= max(pulse.zhsweepers[0].values)
             exp.play(
                 signal=channel_name,
                 pulse=pulse.zhpulse,
-                amplitude=pulse.zhsweeper,
+                amplitude=pulse.zhsweepers[0],
                 phase=pulse.pulse.relative_phase,
             )
         elif any("duration" in param for param in parameters):
             exp.play(
                 signal=channel_name,
                 pulse=pulse.zhpulse,
-                length=pulse.zhsweeper,
+                length=pulse.zhsweepers[0],
                 phase=pulse.pulse.relative_phase,
             )
         elif any("relative_phase" in param for param in parameters):
             exp.play(
                 signal=channel_name,
                 pulse=pulse.zhpulse,
-                phase=pulse.zhsweeper,
+                phase=pulse.zhsweepers[0],
             )
         elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
             exp.play(
@@ -820,11 +816,11 @@ def play_sweep(self, exp, qubit, pulse, section):
         else:
             channel_name = getattr(qubit, section).name
         if isinstance(pulse, ZhSweeperLine):
-            if pulse.zhsweeper.uid == "bias":
+            if pulse.zhsweepers[0].uid == "bias":
                 exp.play(
                     signal=channel_name,
                     pulse=pulse.zhpulse,
-                    amplitude=pulse.zhsweeper,
+                    amplitude=pulse.zhsweepers[0],
                 )
         else:
             parameters = []
@@ -861,12 +857,11 @@ def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
-                        if isinstance(pulse, ZhSweeperLine):
-                            self.play_sweep(exp, coupler, pulse, section="flux")
-                        elif isinstance(pulse, ZhSweeper):
-                            self.play_sweep(exp, coupler, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse):
+                        if isinstance(pulse, ZhPulse):
                             exp.play(signal=channel_name, pulse=pulse.zhpulse)
+                        else:
+                            self.play_sweep(exp, coupler, pulse, section="flux")
+
                 previous_section = section_uid
 
     def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
@@ -895,12 +890,10 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             time = round(
                                 pulse.pulse.duration * NANO_TO_SECONDS, 9
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                        if isinstance(pulse, ZhSweeperLine):
-                            self.play_sweep(exp, qubit, pulse, section="flux")
-                        elif isinstance(pulse, ZhSweeper):
-                            self.play_sweep(exp, qubit, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse):
+                        if isinstance(pulse, ZhPulse):
                             exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
+                        else:
+                            self.play_sweep(exp, qubit, pulse, section="flux")
                 previous_section = section_uid
 
     def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
@@ -936,8 +929,8 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                                     pulse=pulse.zhpulse,
                                     phase=pulse.pulse.relative_phase,
                                 )
-                        elif isinstance(pulse, ZhSweeperLine):
-                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweeper)
+                        else:
+                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweepers[0])
 
                     if len(
                         self.sequence[measure_channel_name(qubit)]
@@ -948,7 +941,7 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             signal=qubit.drive.name,
                             time=self.sequence[measure_channel_name(qubit)][
                                 0
-                            ].zhsweeper,
+                            ].zhsweepers[0],
                         )
                         self.sequence[measure_channel_name(qubit)].remove(
                             self.sequence[measure_channel_name(qubit)][0]
@@ -1133,7 +1126,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
-        self.nt_sweeps = None
+        self.nt_sweeps = []
         sweepers = list(sweepers)
         rearranging_axes, sweepers = self.rearrange_sweepers(sweepers)
         self.sweepers = sweepers
@@ -1180,9 +1173,9 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
         if sweeper.parameter is Parameter.frequency:
             for pulse in sweeper.pulses:
                 line = "drive" if pulse.type is PulseType.DRIVE else "readout"
-                zhsweeper = ZhSweeper(
-                    pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                ).zhsweeper
+                zhsweeper = ZhSweeper.select_sweeper(
+                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
+                )
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
                 if line == "readout":
                     channel_name = measure_channel_name(qubits[pulse.qubit])
@@ -1202,30 +1195,21 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                parameter = ZhSweeper(
-                    pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                ).zhsweeper
+                parameter = ZhSweeper.select_sweeper(
+                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
+                )
                 sweeper.values *= aux_max
 
         if sweeper.parameter is Parameter.bias:
-            if sweeper.qubits:
-                for qubit in sweeper.qubits:
-                    parameter = ZhSweeperLine(
-                        sweeper, qubit, self.sequence_qibo
-                    ).zhsweeper
-            if sweeper.couplers:
-                for qubit in sweeper.couplers:
-                    parameter = ZhSweeperLine(
-                        sweeper, qubit, self.sequence_qibo
-                    ).zhsweeper
+            parameter = ZhSweeperLine.select_sweeper(sweeper)
 
         elif sweeper.parameter is Parameter.start:
-            parameter = ZhSweeperLine(sweeper).zhsweeper
+            parameter = ZhSweeperLine.select_sweeper(sweeper)
 
         elif parameter is None:
-            parameter = ZhSweeper(
-                sweeper.pulses[0], sweeper, qubits[sweeper.pulses[0].qubit]
-            ).zhsweeper
+            parameter = ZhSweeper.select_sweeper(
+                sweeper.pulses[0].type, sweeper, qubits[sweeper.pulses[0].qubit]
+            )
 
         with exp.sweep(
             uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
@@ -1294,9 +1278,9 @@ def sweep_recursion_nt(
         if sweeper.parameter is Parameter.bias:
             if sweeper.qubits:
                 for qubit in sweeper.qubits:
-                    zhsweeper = ZhSweeperLine(
+                    zhsweeper = ZhSweeperLine.select_sweeper(
                         sweeper, qubit, self.sequence_qibo
-                    ).zhsweeper
+                    )
                     zhsweeper.uid = "bias"
                     path = self.find_instrument_address(qubit, "bias")
 

From 3b46a7c5b532ee869b3f625e6aef596530b3d934 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Fri, 16 Feb 2024 13:30:11 +0400
Subject: [PATCH 18/45] get rid of ZhSweeper

---
 src/qibolab/instruments/zhinst.py | 86 +++++++++++--------------------
 1 file changed, 30 insertions(+), 56 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 00c5d9a76..dea6745d0 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -163,23 +163,11 @@ def __init__(self, pulse):
         """Qibolab pulse."""
         self.zhpulse = select_pulse(pulse, pulse.type.name.lower())
         """Zurich pulse."""
+        self.zhsweepers = []
 
-
-class ZhSweeper:
-    """Zurich sweeper from qibolab sweeper for pulse parameters Amplitude,
-    Duration, Frequency (and maybe Phase)"""
-
-    def __init__(self, pulse, sweeper, qubit):
-        p = ZhPulse(pulse)
-
-        self.pulse = p.pulse
-        """Qibolab pulse associated to the sweeper."""
-
-        self.zhpulse = p.zhpulse
-        """Zurich pulse associated to the sweeper."""
-
-        self.zhsweepers = [self.select_sweeper(pulse.type, sweeper, qubit)]
-        """Zurich sweepers."""
+    def add_sweeper(self, sweeper, qubit):
+        """Add sweeper to list of sweepers associated with this pulse."""
+        self.zhsweepers.append(self.select_sweeper(self.pulse.type, sweeper, qubit))
 
     @staticmethod  # pylint: disable=R0903
     def select_sweeper(ptype, sweeper, qubit):
@@ -216,10 +204,6 @@ def select_sweeper(ptype, sweeper, qubit):
                 count=len(sweeper.values),
             )
 
-    def add_sweeper(self, sweeper, qubit):
-        """Add sweeper to list of sweepers."""
-        self.zhsweepers.append(self.select_sweeper(self.pulse.type, sweeper, qubit))
-
 
 class ZhSweeperLine:
     """Zurich sweeper from qibolab sweeper for non pulse parameters Bias, Delay
@@ -625,11 +609,10 @@ def sequence_zh(self, sequence, qubits, couplers):
         # Fill the sequences with pulses according to their lines in temporal order
         for pulse in sequence:
             if pulse.type == PulseType.READOUT:
-                zhsequence[measure_channel_name(qubits[pulse.qubit])].append(
-                    ZhPulse(pulse)
-                )
+                ch = measure_channel_name(qubits[pulse.qubit])
             else:
-                zhsequence[pulse.channel].append(ZhPulse(pulse))
+                ch = pulse.channel
+            zhsequence[ch].append(ZhPulse(pulse))
 
         for sweeper in self.sweepers:
             if sweeper.parameter.name in SWEEPER_SET:
@@ -652,17 +635,6 @@ def sequence_zh(self, sequence, qubits, couplers):
                     for element in aux_list:
                         if pulse == element.pulse:
                             if isinstance(aux_list[aux_list.index(element)], ZhPulse):
-                                if isinstance(pulse, CouplerFluxPulse):
-                                    aux_list[aux_list.index(element)] = ZhSweeper(
-                                        pulse, sweeper, couplers[pulse.qubit]
-                                    )
-                                else:
-                                    aux_list[aux_list.index(element)] = ZhSweeper(
-                                        pulse, sweeper, qubits[pulse.qubit]
-                                    )
-                            elif isinstance(
-                                aux_list[aux_list.index(element)], ZhSweeper
-                            ):
                                 if isinstance(pulse, CouplerFluxPulse):
                                     aux_list[aux_list.index(element)].add_sweeper(
                                         sweeper, couplers[pulse.qubit]
@@ -857,10 +829,12 @@ def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
-                        if isinstance(pulse, ZhPulse):
-                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
-                        else:
+                        if isinstance(pulse, ZhSweeperLine):
+                            self.play_sweep(exp, coupler, pulse, section="flux")
+                        elif isinstance(pulse, ZhPulse) and pulse.zhsweepers:
                             self.play_sweep(exp, coupler, pulse, section="flux")
+                        elif isinstance(pulse, ZhPulse):
+                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
 
                 previous_section = section_uid
 
@@ -890,10 +864,12 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             time = round(
                                 pulse.pulse.duration * NANO_TO_SECONDS, 9
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                        if isinstance(pulse, ZhPulse):
-                            exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
-                        else:
+                        if isinstance(pulse, ZhSweeperLine):
+                            self.play_sweep(exp, qubit, pulse, section="flux")
+                        elif isinstance(pulse, ZhPulse) and pulse.zhsweepers:
                             self.play_sweep(exp, qubit, pulse, section="flux")
+                        elif isinstance(pulse, ZhPulse):
+                            exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
                 previous_section = section_uid
 
     def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
@@ -911,7 +887,9 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                 section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
-                        if not isinstance(pulse, ZhSweeperLine):
+                        if isinstance(pulse, ZhSweeperLine):
+                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweepers[0])
+                        else:
                             exp.delay(
                                 signal=qubit.drive.name,
                                 time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
@@ -921,16 +899,14 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                                 pulse.pulse.duration * NANO_TO_SECONDS, 9
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                             pulse.zhpulse.uid += f"{i}_{j}"
-                            if isinstance(pulse, ZhSweeper):
+                            if pulse.zhsweepers:
                                 self.play_sweep(exp, qubit, pulse, section="drive")
-                            elif isinstance(pulse, ZhPulse):
+                            else:
                                 exp.play(
                                     signal=qubit.drive.name,
                                     pulse=pulse.zhpulse,
                                     phase=pulse.pulse.relative_phase,
                                 )
-                        else:
-                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweepers[0])
 
                     if len(
                         self.sequence[measure_channel_name(qubit)]
@@ -1173,7 +1149,7 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
         if sweeper.parameter is Parameter.frequency:
             for pulse in sweeper.pulses:
                 line = "drive" if pulse.type is PulseType.DRIVE else "readout"
-                zhsweeper = ZhSweeper.select_sweeper(
+                zhsweeper = ZhPulse.select_sweeper(
                     pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
                 )
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
@@ -1195,7 +1171,7 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                parameter = ZhSweeper.select_sweeper(
+                parameter = ZhPulse.select_sweeper(
                     pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
                 )
                 sweeper.values *= aux_max
@@ -1207,7 +1183,7 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
             parameter = ZhSweeperLine.select_sweeper(sweeper)
 
         elif parameter is None:
-            parameter = ZhSweeper.select_sweeper(
+            parameter = ZhPulse.select_sweeper(
                 sweeper.pulses[0].type, sweeper, qubits[sweeper.pulses[0].qubit]
             )
 
@@ -1278,9 +1254,7 @@ def sweep_recursion_nt(
         if sweeper.parameter is Parameter.bias:
             if sweeper.qubits:
                 for qubit in sweeper.qubits:
-                    zhsweeper = ZhSweeperLine.select_sweeper(
-                        sweeper, qubit, self.sequence_qibo
-                    )
+                    zhsweeper = ZhSweeperLine.select_sweeper(sweeper)
                     zhsweeper.uid = "bias"
                     path = self.find_instrument_address(qubit, "bias")
 
@@ -1296,9 +1270,9 @@ def sweep_recursion_nt(
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                zhsweeper = ZhSweeper(
+                zhsweeper = ZhPulse.select_sweeper(
                     pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                ).zhsweeper
+                )
                 sweeper.values *= aux_max
 
                 zhsweeper.uid = "amplitude"
@@ -1311,9 +1285,9 @@ def sweep_recursion_nt(
                 )
 
         elif parameter is None:
-            parameter = ZhSweeper(
+            parameter = ZhPulse.select_sweeper(
                 sweeper.pulses[0], sweeper, qubits[sweeper.pulses[0].qubit]
-            ).zhsweeper
+            )
             device_path = f"/{path}/qachannels/*/oscs/0/gain"  # Hardcoded SHFQA device
 
         with exp.sweep(

From 15b05f6d27d1aadf84794932daa4d6ef07348c84 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:09:57 +0400
Subject: [PATCH 19/45] simplify usage of ZhSweeperLine

---
 src/qibolab/instruments/zhinst.py | 200 +++++++++++-------------------
 1 file changed, 71 insertions(+), 129 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index dea6745d0..4e9c6315d 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -15,7 +15,9 @@
 
 from qibolab import AcquisitionType, AveragingMode, ExecutionParameters
 from qibolab.couplers import Coupler
-from qibolab.pulses import CouplerFluxPulse, FluxPulse, PulseSequence, PulseType
+
+from qibolab.instruments.unrolling import batch_max_sequences
+from qibolab.pulses import CouplerFluxPulse, PulseSequence, PulseType
 from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
@@ -73,7 +75,7 @@ def select_pulse(pulse, pulse_type):
         if str(pulse.shape) == "Rectangular()":
             can_compress = pulse.type is not PulseType.READOUT
             return lo.pulse_library.const(
-                uid=(f"{pulse_type}_{pulse.qubit}_"),
+                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 can_compress=can_compress,
@@ -81,7 +83,7 @@ def select_pulse(pulse, pulse_type):
         if "Gaussian" in str(pulse.shape):
             sigma = pulse.shape.rel_sigma
             return lo.pulse_library.gaussian(
-                uid=(f"{pulse_type}_{pulse.qubit}_"),
+                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 sigma=2 / sigma,
@@ -93,7 +95,7 @@ def select_pulse(pulse, pulse_type):
             width = pulse.shape.width
             can_compress = pulse.type is not PulseType.READOUT
             return lo.pulse_library.gaussian_square(
-                uid=(f"{pulse_type}_{pulse.qubit}_"),
+                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
                 amplitude=pulse.amplitude,
@@ -106,7 +108,7 @@ def select_pulse(pulse, pulse_type):
             sigma = pulse.shape.rel_sigma
             beta = pulse.shape.beta
             return lo.pulse_library.drag(
-                uid=(f"{pulse_type}_{pulse.qubit}_"),
+                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 sigma=2 / sigma,
@@ -116,20 +118,54 @@ def select_pulse(pulse, pulse_type):
 
     if np.all(pulse.envelope_waveform_q(SAMPLING_RATE).data == 0):
         return sampled_pulse_real(
-            uid=(f"{pulse_type}_{pulse.qubit}_"),
+            uid=f"{pulse_type}_{pulse.qubit}_",
             samples=pulse.envelope_waveform_i(SAMPLING_RATE).data,
             can_compress=True,
         )
     else:
         # Test this when we have pulses that use it
         return sampled_pulse_complex(
-            uid=(f"{pulse_type}_{pulse.qubit}_"),
+            uid=f"{pulse_type}_{pulse.qubit}_",
             samples=pulse.envelope_waveform_i(SAMPLING_RATE).data
             + (1j * pulse.envelope_waveform_q(SAMPLING_RATE).data),
             can_compress=True,
         )
 
 
+def select_sweeper(sweeper, ptype=None, qubit=None):
+    """Sweeper translation."""
+    if sweeper.parameter in (Parameter.duration, Parameter.start):
+        return lo.SweepParameter(
+            uid=sweeper.parameter.name,
+            values=sweeper.values * NANO_TO_SECONDS,
+        )
+    if sweeper.parameter is Parameter.frequency:
+        if ptype is None or qubit is None:
+            raise ValueError(
+                "For frequency sweep ptype and qubit arguments are mandatory"
+            )  # FIXME
+        if ptype is PulseType.READOUT:
+            intermediate_frequency = (
+                qubit.readout_frequency - qubit.readout.local_oscillator.frequency
+            )
+        elif ptype is PulseType.DRIVE:
+            intermediate_frequency = (
+                qubit.drive_frequency - qubit.drive.local_oscillator.frequency
+            )
+        else:
+            raise ValueError(
+                f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
+            )
+        return lo.LinearSweepParameter(
+            uid=sweeper.parameter.name,
+            start=sweeper.values[0] + intermediate_frequency,
+            stop=sweeper.values[-1] + intermediate_frequency,
+            count=len(sweeper.values),
+        )
+
+    return lo.SweepParameter(uid=sweeper.parameter.name, values=sweeper.values)
+
+
 def classify_sweepers(sweepers):
     """"""
     nt_sweepers = []
@@ -165,100 +201,24 @@ def __init__(self, pulse):
         """Zurich pulse."""
         self.zhsweepers = []
 
+    # pylint: disable=R0903
     def add_sweeper(self, sweeper, qubit):
         """Add sweeper to list of sweepers associated with this pulse."""
-        self.zhsweepers.append(self.select_sweeper(self.pulse.type, sweeper, qubit))
-
-    @staticmethod  # pylint: disable=R0903
-    def select_sweeper(ptype, sweeper, qubit):
-        """Sweeper translation."""
-
-        if sweeper.parameter is Parameter.amplitude:
-            return lo.SweepParameter(
-                uid=sweeper.parameter.name,
-                values=copy.copy(sweeper.values),
-            )
-        if sweeper.parameter is Parameter.duration:
-            return lo.SweepParameter(
-                uid=sweeper.parameter.name,
-                values=sweeper.values * NANO_TO_SECONDS,
-            )
-        if sweeper.parameter is Parameter.relative_phase:
-            return lo.SweepParameter(
-                uid=sweeper.parameter.name,
-                values=sweeper.values,
-            )
-        if sweeper.parameter is Parameter.frequency:
-            if ptype is PulseType.READOUT:
-                intermediate_frequency = (
-                    qubit.readout_frequency - qubit.readout.local_oscillator.frequency
-                )
-            elif ptype is PulseType.DRIVE:
-                intermediate_frequency = (
-                    qubit.drive_frequency - qubit.drive.local_oscillator.frequency
-                )
-            return lo.LinearSweepParameter(
-                uid=sweeper.parameter.name,
-                start=sweeper.values[0] + intermediate_frequency,
-                stop=sweeper.values[-1] + intermediate_frequency,
-                count=len(sweeper.values),
-            )
+        self.zhsweepers.append(select_sweeper(sweeper, self.pulse.type, qubit))
 
 
 class ZhSweeperLine:
     """Zurich sweeper from qibolab sweeper for non pulse parameters Bias, Delay
     (, power_range, local_oscillator frequency, offset ???)"""
 
-    def __init__(self, sweeper, qubit=None, sequence=None, pulse=None):
-        if sweeper.parameter is Parameter.bias:
-            if isinstance(qubit, Qubit):
-                pulse = FluxPulse(
-                    start=0,
-                    duration=sequence.duration + sequence.start,
-                    amplitude=1,
-                    shape="Rectangular",
-                    channel=qubit.flux.name,
-                    qubit=qubit.name,
-                )
-            if isinstance(qubit, Coupler):
-                pulse = CouplerFluxPulse(
-                    start=0,
-                    duration=sequence.duration + sequence.start,
-                    amplitude=1,
-                    shape="Rectangular",
-                    channel=qubit.flux.name,
-                    qubit=qubit.name,
-                )
-
+    # pylint: disable=R0903
+    def __init__(self, sweeper, pulse=None):
+        if pulse:
             self.pulse = pulse
-
-            self.zhpulse = lo.pulse_library.const(
-                uid=(f"{pulse.type.name.lower()}_{pulse.qubit}_"),
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-            )
-
-        elif sweeper.parameter is Parameter.start:
-            if pulse:
-                self.pulse = pulse
-                self.zhpulse = ZhPulse(pulse).zhpulse
+            self.zhpulse = ZhPulse(pulse).zhpulse
 
         # Need something better to store multiple sweeps on the same pulse
-        self.zhsweepers = [self.select_sweeper(sweeper)]
-
-    @staticmethod  # pylint: disable=R0903
-    def select_sweeper(sweeper):
-        """Sweeper translation."""
-        if sweeper.parameter is Parameter.bias:
-            return lo.SweepParameter(
-                uid=sweeper.parameter.name,
-                values=sweeper.values,
-            )
-        if sweeper.parameter is Parameter.start:
-            return lo.SweepParameter(
-                uid=sweeper.parameter.name,
-                values=sweeper.values * NANO_TO_SECONDS,
-            )
+        self.zhsweepers = [select_sweeper(sweeper)]
 
 
 class Zurich(Controller):
@@ -654,14 +614,11 @@ def sequence_zh(self, sequence, qubits, couplers):
                     aux_list = zhsequence[pulse.channel]
                 for element in aux_list:
                     if pulse == element.pulse:
-                        if isinstance(aux_list[aux_list.index(element)], ZhPulse):
-                            aux_list.insert(
-                                aux_list.index(element),
-                                ZhSweeperLine(
-                                    sweeper, qubits[pulse.qubit], sequence, pulse
-                                ),
-                            )
-                            break
+                        aux_list.insert(
+                            aux_list.index(element),
+                            ZhSweeperLine(sweeper, pulse),
+                        )
+                        break
 
         self.sequence = zhsequence
         self.nt_sweeps, self.sweepers = classify_sweepers(self.sweepers)
@@ -787,14 +744,7 @@ def play_sweep(self, exp, qubit, pulse, section):
             channel_name = measure_channel_name(qubit)
         else:
             channel_name = getattr(qubit, section).name
-        if isinstance(pulse, ZhSweeperLine):
-            if pulse.zhsweepers[0].uid == "bias":
-                exp.play(
-                    signal=channel_name,
-                    pulse=pulse.zhpulse,
-                    amplitude=pulse.zhsweepers[0],
-                )
-        else:
+        if not isinstance(pulse, ZhSweeperLine):
             parameters = []
             for partial_sweep in pulse.zhsweepers:
                 parameters.append(partial_sweep.uid)
@@ -829,11 +779,9 @@ def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
-                        if isinstance(pulse, ZhSweeperLine):
-                            self.play_sweep(exp, coupler, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse) and pulse.zhsweepers:
+                        if pulse.zhsweepers:
                             self.play_sweep(exp, coupler, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse):
+                        else:
                             exp.play(signal=channel_name, pulse=pulse.zhpulse)
 
                 previous_section = section_uid
@@ -864,11 +812,9 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             time = round(
                                 pulse.pulse.duration * NANO_TO_SECONDS, 9
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                        if isinstance(pulse, ZhSweeperLine):
-                            self.play_sweep(exp, qubit, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse) and pulse.zhsweepers:
+                        if pulse.zhsweepers:
                             self.play_sweep(exp, qubit, pulse, section="flux")
-                        elif isinstance(pulse, ZhPulse):
+                        else:
                             exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
                 previous_section = section_uid
 
@@ -1149,8 +1095,8 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
         if sweeper.parameter is Parameter.frequency:
             for pulse in sweeper.pulses:
                 line = "drive" if pulse.type is PulseType.DRIVE else "readout"
-                zhsweeper = ZhPulse.select_sweeper(
-                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
+                zhsweeper = select_sweeper(
+                    sweeper, pulse.type, qubits[sweeper.pulses[0].qubit]
                 )
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
                 if line == "readout":
@@ -1171,20 +1117,18 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                parameter = ZhPulse.select_sweeper(
-                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
-                )
+                parameter = select_sweeper(sweeper)
                 sweeper.values *= aux_max
 
         if sweeper.parameter is Parameter.bias:
-            parameter = ZhSweeperLine.select_sweeper(sweeper)
+            parameter = select_sweeper(sweeper)
 
         elif sweeper.parameter is Parameter.start:
-            parameter = ZhSweeperLine.select_sweeper(sweeper)
+            parameter = select_sweeper(sweeper)
 
         elif parameter is None:
-            parameter = ZhPulse.select_sweeper(
-                sweeper.pulses[0].type, sweeper, qubits[sweeper.pulses[0].qubit]
+            parameter = select_sweeper(
+                sweeper, sweeper.pulses[0].type, qubits[sweeper.pulses[0].qubit]
             )
 
         with exp.sweep(
@@ -1254,7 +1198,7 @@ def sweep_recursion_nt(
         if sweeper.parameter is Parameter.bias:
             if sweeper.qubits:
                 for qubit in sweeper.qubits:
-                    zhsweeper = ZhSweeperLine.select_sweeper(sweeper)
+                    zhsweeper = select_sweeper(sweeper)
                     zhsweeper.uid = "bias"
                     path = self.find_instrument_address(qubit, "bias")
 
@@ -1270,9 +1214,7 @@ def sweep_recursion_nt(
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                zhsweeper = ZhPulse.select_sweeper(
-                    pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                )
+                zhsweeper = select_sweeper(sweeper)
                 sweeper.values *= aux_max
 
                 zhsweeper.uid = "amplitude"
@@ -1285,8 +1227,8 @@ def sweep_recursion_nt(
                 )
 
         elif parameter is None:
-            parameter = ZhPulse.select_sweeper(
-                sweeper.pulses[0], sweeper, qubits[sweeper.pulses[0].qubit]
+            parameter = select_sweeper(
+                sweeper, sweeper.pulses[0].type, qubits[sweeper.pulses[0].qubit]
             )
             device_path = f"/{path}/qachannels/*/oscs/0/gain"  # Hardcoded SHFQA device
 

From 258457121d1952e7778724c283191a841789b7be Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:49:25 +0400
Subject: [PATCH 20/45] untangle sequence_zh

---
 src/qibolab/instruments/zhinst.py | 72 ++++++++++++-------------------
 1 file changed, 27 insertions(+), 45 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 4e9c6315d..0e40f2214 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -17,7 +17,7 @@
 from qibolab.couplers import Coupler
 
 from qibolab.instruments.unrolling import batch_max_sequences
-from qibolab.pulses import CouplerFluxPulse, PulseSequence, PulseType
+from qibolab.pulses import PulseSequence, PulseType
 from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
@@ -574,53 +574,35 @@ def sequence_zh(self, sequence, qubits, couplers):
                 ch = pulse.channel
             zhsequence[ch].append(ZhPulse(pulse))
 
-        for sweeper in self.sweepers:
-            if sweeper.parameter.name in SWEEPER_SET:
-                for pulse in sweeper.pulses:
-                    if pulse.type == PulseType.READOUT:
-                        aux_list = zhsequence[measure_channel_name(qubits[pulse.qubit])]
-                    else:
-                        aux_list = zhsequence[pulse.channel]
-                    if (
-                        sweeper.parameter is Parameter.frequency
-                        and pulse.type is PulseType.READOUT
-                    ):
-                        self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
-                    if (
-                        sweeper.parameter is Parameter.amplitude
-                        and pulse.type is PulseType.READOUT
-                    ):
-                        self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
-
-                    for element in aux_list:
-                        if pulse == element.pulse:
-                            if isinstance(aux_list[aux_list.index(element)], ZhPulse):
-                                if isinstance(pulse, CouplerFluxPulse):
-                                    aux_list[aux_list.index(element)].add_sweeper(
-                                        sweeper, couplers[pulse.qubit]
-                                    )
-                                else:
-                                    aux_list[aux_list.index(element)].add_sweeper(
-                                        sweeper, qubits[pulse.qubit]
-                                    )
-
-            # This may not place the Zhsweeper when the start occurs among different sections or lines
-            if sweeper.parameter.name in SWEEPER_START:
-                pulse = sweeper.pulses[0]
-
-                if pulse.type == PulseType.READOUT:
-                    aux_list = zhsequence[measure_channel_name(qubits[pulse.qubit])]
-                else:
-                    aux_list = zhsequence[pulse.channel]
-                for element in aux_list:
-                    if pulse == element.pulse:
-                        aux_list.insert(
-                            aux_list.index(element),
-                            ZhSweeperLine(sweeper, pulse),
-                        )
+        def get_sweeps(pulse: ZhPulse):
+            sweepers_for_pulse = []
+            for sweeper in self.sweepers:
+                for p in sweeper.pulses:
+                    if p == pulse:
+                        sweepers_for_pulse.append(sweeper)
                         break
+            return sweepers_for_pulse
+
+        for ch, zhpulses in zhsequence.items():
+            for i, zhpulse in enumerate(zhpulses):
+                for s in get_sweeps(zhpulse.pulse):
+                    if s.parameter in SWEEPER_SET:
+                        zhpulse.add_sweeper(s, qubits[zhpulse.pulse.qubit])
+                    if s.parameter in SWEEPER_START:
+                        zhpulses.insert(i, ZhSweeperLine(s, zhpulse.pulse))
 
         self.sequence = zhsequence
+
+        def set_acquisition_type():
+            for sweeper in self.sweepers:
+                if sweeper.parameter.name in {Parameter.frequency, Parameter.amplitude}:
+                    for pulse in sweeper.pulses:
+                        if pulse.type is PulseType.READOUT:
+                            # FIXME: case of multiple pulses
+                            self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
+
+        set_acquisition_type()
+
         self.nt_sweeps, self.sweepers = classify_sweepers(self.sweepers)
 
     def create_exp(self, qubits, couplers, options):

From aa779e20e856e9e09edec810bbb15f1f8f48cc7a Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 20 Feb 2024 12:26:23 +0400
Subject: [PATCH 21/45] use channel names instead of qubits and couplers to
 construct exp

---
 src/qibolab/instruments/zhinst.py | 104 +++++++++---------------------
 1 file changed, 32 insertions(+), 72 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 0e40f2214..71ee8c236 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -652,9 +652,11 @@ def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
 
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
-        self.drive(exp, qubits)
-        self.flux(exp, qubits)
-        self.couplerflux(exp, couplers)
+        self.drive(
+            exp, [(q.drive.name, measure_channel_name(q)) for q in qubits.values()]
+        )
+        self.flux(exp, [q.flux.name for q in qubits.values()])
+        self.flux(exp, [c.flux.name for c in couplers.values()])
         self.measure_relax(
             exp,
             qubits,
@@ -664,9 +666,7 @@ def select_exp(self, exp, qubits, couplers, exp_options):
         )
 
     @staticmethod
-    def play_sweep_select_single(
-        exp, qubit, pulse, channel_name, parameters, partial_sweep
-    ):
+    def play_sweep_select_single(exp, pulse, channel_name, parameters, partial_sweep):
         """Play Zurich pulse when a single sweeper is involved."""
         if any("amplitude" in param for param in parameters):
             pulse.zhpulse.amplitude *= max(pulse.zhsweepers[0].values)
@@ -699,7 +699,7 @@ def play_sweep_select_single(
 
     # Hardcoded for the flux pulse for 2q gates
     @staticmethod
-    def play_sweep_select_dual(exp, qubit, pulse, channel_name, parameters):
+    def play_sweep_select_dual(exp, pulse, channel_name, parameters):
         """Play Zurich pulse when two sweepers are involved on the same
         pulse."""
         if "amplitude" in parameters and "duration" in parameters:
@@ -719,75 +719,39 @@ def play_sweep_select_dual(exp, qubit, pulse, channel_name, parameters):
                 length=pulse.zhsweepers[sweeper_dur_index],
             )
 
-    def play_sweep(self, exp, qubit, pulse, section):
+    def play_sweep(self, exp, channel_name, pulse):
         """Takes care of playing the sweepers and involved pulses for different
         options."""
-        if section == "readout":
-            channel_name = measure_channel_name(qubit)
-        else:
-            channel_name = getattr(qubit, section).name
         if not isinstance(pulse, ZhSweeperLine):
             parameters = []
             for partial_sweep in pulse.zhsweepers:
                 parameters.append(partial_sweep.uid)
             # Recheck partial sweeps
             if len(parameters) == 2:
-                self.play_sweep_select_dual(exp, qubit, pulse, channel_name, parameters)
+                self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
             else:
                 self.play_sweep_select_single(
-                    exp, qubit, pulse, channel_name, parameters, partial_sweep
+                    exp, pulse, channel_name, parameters, partial_sweep
                 )
 
-    def couplerflux(self, exp: lo.Experiment, couplers: Dict[str, Coupler]):
-        """Coupler flux for bias sweep or pulses.
-
-        Args:
-            exp (lo.Experiment): laboneq experiment on which register sequences.
-            couplers (dict[str, Coupler]): coupler on which pulses are played.
-        """
-        for coupler in couplers.values():
-            channel_name = coupler.flux.name
-            time = 0
-            previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[channel_name]):
-                section_uid = f"sequence_{channel_name}_{i}"
-                with exp.section(uid=section_uid, play_after=previous_section):
-                    for j, pulse in enumerate(sequence):
-                        pulse.zhpulse.uid += f"{i}_{j}"
-                        exp.delay(
-                            signal=channel_name,
-                            time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
-                        )
-                        time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
-                            pulse.pulse.start * NANO_TO_SECONDS, 9
-                        )
-                        if pulse.zhsweepers:
-                            self.play_sweep(exp, coupler, pulse, section="flux")
-                        else:
-                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
-
-                previous_section = section_uid
-
-    def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
+    def flux(self, exp: lo.Experiment, channels: list[str]):
         """Qubit flux for bias sweep or pulses.
 
         Args:
             exp (lo.Experiment): laboneq experiment on which register sequences.
-            qubits (dict[str, Qubit]): qubits on which pulses are played.
+            channels: list of qubit flux channel names.
         """
-        for qubit in qubits.values():
-            q = qubit.name  # pylint: disable=C0103
-            channel_name = qubit.flux.name
+        for channel_name in channels:
             time = 0
             previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[qubit.flux.name]):
+            for i, sequence in enumerate(self.sub_sequences[channel_name]):
                 section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
                         if not isinstance(pulse, ZhSweeperLine):
                             pulse.zhpulse.uid += f"{i}_{j}"
                             exp.delay(
-                                signal=qubit.flux.name,
+                                signal=channel_name,
                                 time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                                 - time,
                             )
@@ -795,31 +759,31 @@ def flux(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                                 pulse.pulse.duration * NANO_TO_SECONDS, 9
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                         if pulse.zhsweepers:
-                            self.play_sweep(exp, qubit, pulse, section="flux")
+                            self.play_sweep(exp, channel_name, pulse)
                         else:
-                            exp.play(signal=qubit.flux.name, pulse=pulse.zhpulse)
+                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
+
                 previous_section = section_uid
 
-    def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
+    def drive(self, exp: lo.Experiment, channels: list[tuple[str, str]]):
         """Qubit driving pulses.
 
         Args:
             exp (lo.Experiment): laboneq experiment on which register sequences.
-            qubits (dict[str, Qubit]): qubits on which pulses are played.
+            channels: list of (drive channel name, measure channel name) pairs.  # FIXME: remove measure channel
         """
-        for qubit in qubits.values():
-            channel_name = qubit.drive.name
+        for channel_name, channel_name_measure in channels:
             time = 0
             previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[qubit.drive.name]):
+            for i, sequence in enumerate(self.sub_sequences[channel_name]):
                 section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
                         if isinstance(pulse, ZhSweeperLine):
-                            exp.delay(signal=qubit.drive.name, time=pulse.zhsweepers[0])
+                            exp.delay(signal=channel_name, time=pulse.zhsweepers[0])
                         else:
                             exp.delay(
-                                signal=qubit.drive.name,
+                                signal=channel_name,
                                 time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                                 - time,
                             )
@@ -828,27 +792,23 @@ def drive(self, exp: lo.Experiment, qubits: Dict[str, Qubit]):
                             ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
                             pulse.zhpulse.uid += f"{i}_{j}"
                             if pulse.zhsweepers:
-                                self.play_sweep(exp, qubit, pulse, section="drive")
+                                self.play_sweep(exp, channel_name, pulse)
                             else:
                                 exp.play(
-                                    signal=qubit.drive.name,
+                                    signal=channel_name,
                                     pulse=pulse.zhpulse,
                                     phase=pulse.pulse.relative_phase,
                                 )
 
-                    if len(
-                        self.sequence[measure_channel_name(qubit)]
-                    ) > 0 and isinstance(
-                        self.sequence[measure_channel_name(qubit)][0], ZhSweeperLine
+                    if len(self.sequence[channel_name_measure]) > 0 and isinstance(
+                        self.sequence[channel_name_measure][0], ZhSweeperLine
                     ):
                         exp.delay(
-                            signal=qubit.drive.name,
-                            time=self.sequence[measure_channel_name(qubit)][
-                                0
-                            ].zhsweepers[0],
+                            signal=channel_name,
+                            time=self.sequence[channel_name_measure][0].zhsweepers[0],
                         )
-                        self.sequence[measure_channel_name(qubit)].remove(
-                            self.sequence[measure_channel_name(qubit)][0]
+                        self.sequence[channel_name_measure].remove(
+                            self.sequence[channel_name_measure][0]
                         )
 
                 previous_section = section_uid

From 41d4df2987be6bc48311fe17085ddbf7939eba54 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 21 Feb 2024 09:35:39 +0400
Subject: [PATCH 22/45] get rid of ZhSweeperLine

---
 src/qibolab/instruments/zhinst.py | 100 +++++++++++++-----------------
 1 file changed, 43 insertions(+), 57 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 71ee8c236..fcffd031f 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -201,24 +201,15 @@ def __init__(self, pulse):
         """Zurich pulse."""
         self.zhsweepers = []
 
+        self.delay_sweeper = None
+
     # pylint: disable=R0903
     def add_sweeper(self, sweeper, qubit):
         """Add sweeper to list of sweepers associated with this pulse."""
         self.zhsweepers.append(select_sweeper(sweeper, self.pulse.type, qubit))
 
-
-class ZhSweeperLine:
-    """Zurich sweeper from qibolab sweeper for non pulse parameters Bias, Delay
-    (, power_range, local_oscillator frequency, offset ???)"""
-
-    # pylint: disable=R0903
-    def __init__(self, sweeper, pulse=None):
-        if pulse:
-            self.pulse = pulse
-            self.zhpulse = ZhPulse(pulse).zhpulse
-
-        # Need something better to store multiple sweeps on the same pulse
-        self.zhsweepers = [select_sweeper(sweeper)]
+    def add_delay_sweeper(self, sweeper):
+        self.delay_sweeper = select_sweeper(sweeper)
 
 
 class Zurich(Controller):
@@ -589,7 +580,7 @@ def get_sweeps(pulse: ZhPulse):
                     if s.parameter in SWEEPER_SET:
                         zhpulse.add_sweeper(s, qubits[zhpulse.pulse.qubit])
                     if s.parameter in SWEEPER_START:
-                        zhpulses.insert(i, ZhSweeperLine(s, zhpulse.pulse))
+                        zhpulse.add_delay_sweeper(s)
 
         self.sequence = zhsequence
 
@@ -722,17 +713,16 @@ def play_sweep_select_dual(exp, pulse, channel_name, parameters):
     def play_sweep(self, exp, channel_name, pulse):
         """Takes care of playing the sweepers and involved pulses for different
         options."""
-        if not isinstance(pulse, ZhSweeperLine):
-            parameters = []
-            for partial_sweep in pulse.zhsweepers:
-                parameters.append(partial_sweep.uid)
-            # Recheck partial sweeps
-            if len(parameters) == 2:
-                self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
-            else:
-                self.play_sweep_select_single(
-                    exp, pulse, channel_name, parameters, partial_sweep
-                )
+        parameters = []
+        for partial_sweep in pulse.zhsweepers:
+            parameters.append(partial_sweep.uid)
+        # Recheck partial sweeps
+        if len(parameters) == 2:
+            self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
+        else:
+            self.play_sweep_select_single(
+                exp, pulse, channel_name, parameters, partial_sweep
+            )
 
     def flux(self, exp: lo.Experiment, channels: list[str]):
         """Qubit flux for bias sweep or pulses.
@@ -748,16 +738,16 @@ def flux(self, exp: lo.Experiment, channels: list[str]):
                 section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
-                        if not isinstance(pulse, ZhSweeperLine):
-                            pulse.zhpulse.uid += f"{i}_{j}"
-                            exp.delay(
-                                signal=channel_name,
-                                time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                                - time,
-                            )
-                            time = round(
-                                pulse.pulse.duration * NANO_TO_SECONDS, 9
-                            ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
+                        if pulse.delay_sweeper:
+                            exp.delay(signal=channel_name, time=pulse.delay_sweeper)
+                        pulse.zhpulse.uid += f"{i}_{j}"
+                        exp.delay(
+                            signal=channel_name,
+                            time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
+                        )
+                        time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
+                            pulse.pulse.start * NANO_TO_SECONDS, 9
+                        )
                         if pulse.zhsweepers:
                             self.play_sweep(exp, channel_name, pulse)
                         else:
@@ -779,37 +769,33 @@ def drive(self, exp: lo.Experiment, channels: list[tuple[str, str]]):
                 section_uid = f"sequence_{channel_name}_{i}"
                 with exp.section(uid=section_uid, play_after=previous_section):
                     for j, pulse in enumerate(sequence):
-                        if isinstance(pulse, ZhSweeperLine):
-                            exp.delay(signal=channel_name, time=pulse.zhsweepers[0])
+                        if pulse.delay_sweeper:
+                            exp.delay(signal=channel_name, time=pulse.delay_sweeper)
+                        exp.delay(
+                            signal=channel_name,
+                            time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
+                        )
+                        time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
+                            pulse.pulse.start * NANO_TO_SECONDS, 9
+                        )
+                        pulse.zhpulse.uid += f"{i}_{j}"
+                        if pulse.zhsweepers:
+                            self.play_sweep(exp, channel_name, pulse)
                         else:
-                            exp.delay(
+                            exp.play(
                                 signal=channel_name,
-                                time=round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                                - time,
+                                pulse=pulse.zhpulse,
+                                phase=pulse.pulse.relative_phase,
                             )
-                            time = round(
-                                pulse.pulse.duration * NANO_TO_SECONDS, 9
-                            ) + round(pulse.pulse.start * NANO_TO_SECONDS, 9)
-                            pulse.zhpulse.uid += f"{i}_{j}"
-                            if pulse.zhsweepers:
-                                self.play_sweep(exp, channel_name, pulse)
-                            else:
-                                exp.play(
-                                    signal=channel_name,
-                                    pulse=pulse.zhpulse,
-                                    phase=pulse.pulse.relative_phase,
-                                )
 
-                    if len(self.sequence[channel_name_measure]) > 0 and isinstance(
-                        self.sequence[channel_name_measure][0], ZhSweeperLine
+                    if (
+                        len(self.sequence[channel_name_measure]) > 0
+                        and self.sequence[channel_name_measure][0].delay_sweeper
                     ):
                         exp.delay(
                             signal=channel_name,
                             time=self.sequence[channel_name_measure][0].zhsweepers[0],
                         )
-                        self.sequence[channel_name_measure].remove(
-                            self.sequence[channel_name_measure][0]
-                        )
 
                 previous_section = section_uid
 

From 52fe7072c767a4fb44631b4479f440d618302c26 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 21 Feb 2024 10:04:36 +0400
Subject: [PATCH 23/45] new section implementation

---
 src/qibolab/instruments/zhinst.py | 412 +++++++++++-------------------
 1 file changed, 146 insertions(+), 266 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index fcffd031f..e6dc5c55c 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -17,7 +17,7 @@
 from qibolab.couplers import Coupler
 
 from qibolab.instruments.unrolling import batch_max_sequences
-from qibolab.pulses import PulseSequence, PulseType
+from qibolab.pulses import Pulse, PulseSequence, PulseType
 from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
@@ -212,6 +212,12 @@ def add_delay_sweeper(self, sweeper):
         self.delay_sweeper = select_sweeper(sweeper)
 
 
+@dataclass
+class SubSequence:
+    measurements: list[str, Pulse]
+    control_sequence: dict[str, list[Pulse]]
+
+
 class Zurich(Controller):
     """Zurich driver main class."""
 
@@ -448,59 +454,45 @@ def frequency_from_pulses(qubits, sequence):
             if pulse.type is PulseType.DRIVE:
                 qubit.drive_frequency = pulse.frequency
 
-    def create_sub_sequence(
-        self,
-        line_name: str,
-        quantum_elements: Union[Dict[str, Qubit], Dict[str, Coupler]],
-    ):
-        """Create a list of sequences for each measurement.
-
-        Args:
-            line_name (str): Name of the line from which extract the sequence.
-            quantum_elements (dict[str, Qubit]|dict[str, Coupler]): qubits or couplers for the platform.
-        """
-        for quantum_element in quantum_elements.values():
-            if isinstance(quantum_element, Coupler):
-                # Find arbitrary measurement from sequence and split coupler sequence based on that
-                # FIXME: neither the previous nor this way of doing this is not right - division into subsequences should happen
-                # for the entire sequence as a whole and not channel by channel (which is the source of prblem here, that we need to
-                # decide which qubit's measurement to take as baseline in splitting)
-                ch = next(
-                    iter(
-                        ch
-                        for ch, pulses in self.sequence.items()
-                        if any(p.pulse.type == PulseType.READOUT for p in pulses)
-                    )
-                )
-                measurements = self.sequence[ch]
-            else:
-                measurements = self.sequence[measure_channel_name(quantum_element)]
-
-            channel_name = getattr(quantum_element, line_name).name
-            pulses = self.sequence[channel_name]
-            pulse_sequences = [[] for _ in measurements]
-            pulse_sequences.append([])
-            measurement_index = 0
-            for pulse in pulses:
-                if measurement_index < len(measurements):
-                    if pulse.pulse.finish > measurements[measurement_index].pulse.start:
-                        measurement_index += 1
-                pulse_sequences[measurement_index].append(pulse)
-            self.sub_sequences[channel_name] = pulse_sequences
-
-    def create_sub_sequences(
-        self, qubits: Dict[str, Qubit], couplers: Dict[str, Coupler]
-    ):
+    def create_sub_sequences(self, qubits: list[Qubit]) -> list[SubSequence]:
         """Create subsequences for different lines (drive, flux, coupler flux).
 
         Args:
             qubits (dict[str, Qubit]): qubits for the platform.
             couplers (dict[str, Coupler]): couplers for the platform.
         """
-        self.sub_sequences = {}
-        self.create_sub_sequence("drive", qubits)
-        self.create_sub_sequence("flux", qubits)
-        self.create_sub_sequence("flux", couplers)
+        # collect and group all measurements. In one sequence, qubits may be measured multiple times (e.g. when the sequence was unrolled
+        # at platform level), however it is assumed that 1. all qubits are measured the same number of times, 2. each time all qubits are
+        # measured simultaneously.
+        measure_channels = {measure_channel_name(qb) for qb in qubits}
+        other_channels = set(self.sequence.keys()) - measure_channels
+
+        measurement_groups = defaultdict(list)
+        for ch in measure_channels:
+            for i, pulse in enumerate(self.sequence[ch]):
+                measurement_groups[i].append((ch, pulse))
+
+        measurement_starts = {}
+        for i, group in measurement_groups.items():
+            starts = np.array([meas.pulse.start for _, meas in group])
+            # TODO: validate that all start times are equal
+            measurement_starts[i] = max(starts)
+
+        # split all non-measurement channels according to the locations of the measurements
+        sub_sequences = defaultdict(lambda: defaultdict(list))
+        for ch in other_channels:
+            measurement_index = 0
+            for pulse in self.sequence[ch]:
+                if pulse.pulse.finish > measurement_starts[measurement_index]:
+                    measurement_index += 1
+                sub_sequences[measurement_index][ch].append(pulse)
+        if len(sub_sequences) > len(measurement_groups):
+            log.warning("There are control pulses after the last measurement start.")
+
+        return [
+            SubSequence(measurement_groups[i], sub_sequences[i])
+            for i in range(len(measurement_groups))
+        ]
 
     def experiment_flow(
         self,
@@ -520,7 +512,7 @@ def experiment_flow(
             sequence (PulseSequence): sequence of pulses to be played in the experiment.
         """
         self.sequence_zh(sequence, qubits, couplers)
-        self.create_sub_sequences(qubits, couplers)
+        self.sub_sequences = self.create_sub_sequences(list(qubits.values()))
         self.calibration_step(qubits, couplers, options)
         self.create_exp(qubits, couplers, options)
 
@@ -565,10 +557,10 @@ def sequence_zh(self, sequence, qubits, couplers):
                 ch = pulse.channel
             zhsequence[ch].append(ZhPulse(pulse))
 
-        def get_sweeps(pulse: ZhPulse):
+        def get_sweeps(pulse: Pulse):
             sweepers_for_pulse = []
             for sweeper in self.sweepers:
-                for p in sweeper.pulses:
+                for p in sweeper.pulses or []:
                     if p == pulse:
                         sweepers_for_pulse.append(sweeper)
                         break
@@ -577,9 +569,9 @@ def get_sweeps(pulse: ZhPulse):
         for ch, zhpulses in zhsequence.items():
             for i, zhpulse in enumerate(zhpulses):
                 for s in get_sweeps(zhpulse.pulse):
-                    if s.parameter in SWEEPER_SET:
+                    if s.parameter.name in SWEEPER_SET:
                         zhpulse.add_sweeper(s, qubits[zhpulse.pulse.qubit])
-                    if s.parameter in SWEEPER_START:
+                    if s.parameter.name in SWEEPER_START:
                         zhpulse.add_delay_sweeper(s)
 
         self.sequence = zhsequence
@@ -643,233 +635,48 @@ def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
 
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
-        self.drive(
-            exp, [(q.drive.name, measure_channel_name(q)) for q in qubits.values()]
-        )
-        self.flux(exp, [q.flux.name for q in qubits.values()])
-        self.flux(exp, [c.flux.name for c in couplers.values()])
-        self.measure_relax(
-            exp,
-            qubits,
-            couplers,
-            exp_options.relaxation_time,
-            exp_options.acquisition_type,
-        )
-
-    @staticmethod
-    def play_sweep_select_single(exp, pulse, channel_name, parameters, partial_sweep):
-        """Play Zurich pulse when a single sweeper is involved."""
-        if any("amplitude" in param for param in parameters):
-            pulse.zhpulse.amplitude *= max(pulse.zhsweepers[0].values)
-            pulse.zhsweepers[0].values /= max(pulse.zhsweepers[0].values)
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                amplitude=pulse.zhsweepers[0],
-                phase=pulse.pulse.relative_phase,
-            )
-        elif any("duration" in param for param in parameters):
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                length=pulse.zhsweepers[0],
-                phase=pulse.pulse.relative_phase,
-            )
-        elif any("relative_phase" in param for param in parameters):
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                phase=pulse.zhsweepers[0],
-            )
-        elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                phase=pulse.pulse.relative_phase,
-            )
-
-    # Hardcoded for the flux pulse for 2q gates
-    @staticmethod
-    def play_sweep_select_dual(exp, pulse, channel_name, parameters):
-        """Play Zurich pulse when two sweepers are involved on the same
-        pulse."""
-        if "amplitude" in parameters and "duration" in parameters:
-            for sweeper in pulse.zhsweepers:
-                if sweeper.uid == "amplitude":
-                    sweeper_amp_index = pulse.zhsweepers.index(sweeper)
-                    sweeper.values = sweeper.values.copy()
-                    pulse.zhpulse.amplitude *= max(abs(sweeper.values))
-                    sweeper.values /= max(abs(sweeper.values))
-                else:
-                    sweeper_dur_index = pulse.zhsweepers.index(sweeper)
-
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                amplitude=pulse.zhsweepers[sweeper_amp_index],
-                length=pulse.zhsweepers[sweeper_dur_index],
-            )
-
-    def play_sweep(self, exp, channel_name, pulse):
-        """Takes care of playing the sweepers and involved pulses for different
-        options."""
-        parameters = []
-        for partial_sweep in pulse.zhsweepers:
-            parameters.append(partial_sweep.uid)
-        # Recheck partial sweeps
-        if len(parameters) == 2:
-            self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
-        else:
-            self.play_sweep_select_single(
-                exp, pulse, channel_name, parameters, partial_sweep
-            )
-
-    def flux(self, exp: lo.Experiment, channels: list[str]):
-        """Qubit flux for bias sweep or pulses.
-
-        Args:
-            exp (lo.Experiment): laboneq experiment on which register sequences.
-            channels: list of qubit flux channel names.
-        """
-        for channel_name in channels:
-            time = 0
-            previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[channel_name]):
-                section_uid = f"sequence_{channel_name}_{i}"
-                with exp.section(uid=section_uid, play_after=previous_section):
-                    for j, pulse in enumerate(sequence):
+        weights = {}
+        previous_section = None
+        for i, seq in enumerate(self.sub_sequences):
+            section_uid = f"control_{i}"
+            with exp.section(uid=section_uid, play_after=previous_section):
+                for ch, pulses in seq.control_sequence.items():
+                    time = 0
+                    for j, pulse in enumerate(pulses):
                         if pulse.delay_sweeper:
-                            exp.delay(signal=channel_name, time=pulse.delay_sweeper)
-                        pulse.zhpulse.uid += f"{i}_{j}"
+                            exp.delay(signal=ch, time=pulse.delay_sweeper)
                         exp.delay(
-                            signal=channel_name,
+                            signal=ch,
                             time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
                         )
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
-                        if pulse.zhsweepers:
-                            self.play_sweep(exp, channel_name, pulse)
-                        else:
-                            exp.play(signal=channel_name, pulse=pulse.zhpulse)
-
-                previous_section = section_uid
-
-    def drive(self, exp: lo.Experiment, channels: list[tuple[str, str]]):
-        """Qubit driving pulses.
-
-        Args:
-            exp (lo.Experiment): laboneq experiment on which register sequences.
-            channels: list of (drive channel name, measure channel name) pairs.  # FIXME: remove measure channel
-        """
-        for channel_name, channel_name_measure in channels:
-            time = 0
-            previous_section = None
-            for i, sequence in enumerate(self.sub_sequences[channel_name]):
-                section_uid = f"sequence_{channel_name}_{i}"
-                with exp.section(uid=section_uid, play_after=previous_section):
-                    for j, pulse in enumerate(sequence):
-                        if pulse.delay_sweeper:
-                            exp.delay(signal=channel_name, time=pulse.delay_sweeper)
-                        exp.delay(
-                            signal=channel_name,
-                            time=round(pulse.pulse.start * NANO_TO_SECONDS, 9) - time,
-                        )
-                        time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
-                            pulse.pulse.start * NANO_TO_SECONDS, 9
+                        pulse.zhpulse.uid += (
+                            f"{i}_{j}"  # FIXME: is this actually needed?
                         )
-                        pulse.zhpulse.uid += f"{i}_{j}"
                         if pulse.zhsweepers:
-                            self.play_sweep(exp, channel_name, pulse)
+                            self.play_sweep(exp, ch, pulse)
                         else:
                             exp.play(
-                                signal=channel_name,
+                                signal=ch,
                                 pulse=pulse.zhpulse,
                                 phase=pulse.pulse.relative_phase,
                             )
+            previous_section = section_uid
 
-                    if (
-                        len(self.sequence[channel_name_measure]) > 0
-                        and self.sequence[channel_name_measure][0].delay_sweeper
-                    ):
-                        exp.delay(
-                            signal=channel_name,
-                            time=self.sequence[channel_name_measure][0].zhsweepers[0],
-                        )
-
+            if any(m.delay_sweeper is not None for _, m in seq.measurements):
+                section_uid = f"measurement_delay_{i}"
+                with exp.section(uid=section_uid, play_after=previous_section):
+                    for ch, m in seq.measurements:
+                        if m.delay_sweeper:
+                            exp.delay(signal=ch, time=m.delay_sweeper)
                 previous_section = section_uid
 
-    def find_subsequence_finish(
-        self,
-        measurement_number: int,
-        line: str,
-        quantum_elements: Union[Qubit, Coupler],
-    ) -> Tuple[int, str]:
-        """Find the finishing time and qubit for a given sequence.
-
-        Args:
-            measurement_number (int): number of the measure pulse.
-            line (str): line from which measure the finishing time.
-                e.g.: "drive", "flux"
-            quantum_elements: qubits or couplers from which to measure the finishing time.
-
-        Returns:
-            time_finish (int): Finish time of the last pulse of the subsequence
-                before the measurement.
-            sequence_finish (str): Name of the last subsequence before measurement. If
-                there are no sequences after the previous measurement, use "None".
-        """
-        time_finish = 0
-        sequence_finish = "None"
-        for quantum_element in quantum_elements:
-            channel_name = getattr(quantum_element, line).name
-            if len(self.sub_sequences[channel_name]) <= measurement_number:
-                continue
-            for pulse in self.sub_sequences[channel_name][measurement_number]:
-                if pulse.pulse.finish > time_finish:
-                    time_finish = pulse.pulse.finish
-                    sequence_finish = channel_name
-        return time_finish, sequence_finish
-
-    # For pulsed spectroscopy, set integration_length and either measure_pulse or measure_pulse_length.
-    # For CW spectroscopy, set only integration_length and do not specify the measure signal.
-    # For all other measurements, set either length or pulse for both the measure pulse and integration kernel.
-    def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type):
-        """Qubit readout pulse, data acquisition and qubit relaxation."""
-        readout_schedule = defaultdict(list)
-        qubit_readout_schedule = defaultdict(list)
-        for qubit in qubits.values():
-            channel_name = measure_channel_name(qubit)
-            for i, pulse in enumerate(self.sequence[channel_name]):
-                readout_schedule[i].append(pulse)
-                qubit_readout_schedule[i].append(qubit)
-
-        weights = {}
-        for i, (pulses, qubits_readout) in enumerate(
-            zip(
-                readout_schedule.values(),
-                qubit_readout_schedule.values(),
-            )
-        ):
-            qd_finish = self.find_subsequence_finish(i, "drive", qubits_readout)
-            qf_finish = self.find_subsequence_finish(i, "flux", qubits_readout)
-            cf_finish = self.find_subsequence_finish(i, "flux", couplers.values())
-            finish_times = np.array(
-                [
-                    qd_finish,
-                    qf_finish,
-                    cf_finish,
-                ],
-                dtype=[("finish", "i4"), ("line", "U15")],
-            )
-            latest_sequence = finish_times[finish_times["finish"].argmax()]
-            if latest_sequence["line"] == "None":
-                play_after = None
-            else:
-                play_after = f"sequence_{latest_sequence['line']}_{i}"
-            # Section on the outside loop allows for multiplex
-            with exp.section(uid=f"sequence_measure_{i}", play_after=play_after):
-                for pulse, qubit in zip(pulses, qubits_readout):
+            section_uid = f"measure_{i}"
+            with exp.section(uid=section_uid, play_after=previous_section):
+                for ch, pulse in seq.measurements:
+                    qubit = qubits[pulse.pulse.qubit]
                     q = qubit.name
                     pulse.zhpulse.uid += str(i)
 
@@ -880,7 +687,8 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
 
                     if (
                         qubit.kernel is not None
-                        and acquisition_type == lo.AcquisitionType.DISCRIMINATION
+                        and exp_options.acquisition_type
+                        == lo.AcquisitionType.DISCRIMINATION
                     ):
                         weight = lo.pulse_library.sampled_pulse_complex(
                             uid="weight" + str(q),
@@ -889,7 +697,10 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
 
                     else:
                         if i == 0:
-                            if acquisition_type == lo.AcquisitionType.DISCRIMINATION:
+                            if (
+                                exp_options.acquisition_type
+                                == lo.AcquisitionType.DISCRIMINATION
+                            ):
                                 weight = lo.pulse_library.sampled_pulse_complex(
                                     samples=np.ones(
                                         [
@@ -921,7 +732,7 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                     measure_pulse_parameters = {"phase": 0}
 
                     if i == len(self.sequence[measure_channel_name(qubit)]) - 1:
-                        reset_delay = relaxation_time * NANO_TO_SECONDS
+                        reset_delay = exp_options.relaxation_time * NANO_TO_SECONDS
                     else:
                         reset_delay = 0
 
@@ -941,6 +752,75 @@ def measure_relax(self, exp, qubits, couplers, relaxation_time, acquisition_type
                         acquire_delay=self.time_of_flight * NANO_TO_SECONDS,
                         reset_delay=reset_delay,
                     )
+            previous_section = section_uid
+
+    @staticmethod
+    def play_sweep_select_single(exp, pulse, channel_name, parameters, partial_sweep):
+        """Play Zurich pulse when a single sweeper is involved."""
+        if any("amplitude" in param for param in parameters):
+            pulse.zhpulse.amplitude *= max(pulse.zhsweepers[0].values)
+            pulse.zhsweepers[0].values /= max(pulse.zhsweepers[0].values)
+            exp.play(
+                signal=channel_name,
+                pulse=pulse.zhpulse,
+                amplitude=pulse.zhsweepers[0],
+                phase=pulse.pulse.relative_phase,
+            )
+        elif any("duration" in param for param in parameters):
+            exp.play(
+                signal=channel_name,
+                pulse=pulse.zhpulse,
+                length=pulse.zhsweepers[0],
+                phase=pulse.pulse.relative_phase,
+            )
+        elif any("relative_phase" in param for param in parameters):
+            exp.play(
+                signal=channel_name,
+                pulse=pulse.zhpulse,
+                phase=pulse.zhsweepers[0],
+            )
+        elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
+            exp.play(
+                signal=channel_name,
+                pulse=pulse.zhpulse,
+                phase=pulse.pulse.relative_phase,
+            )
+
+    # Hardcoded for the flux pulse for 2q gates
+    @staticmethod
+    def play_sweep_select_dual(exp, pulse, channel_name, parameters):
+        """Play Zurich pulse when two sweepers are involved on the same
+        pulse."""
+        if "amplitude" in parameters and "duration" in parameters:
+            for sweeper in pulse.zhsweepers:
+                if sweeper.uid == "amplitude":
+                    sweeper_amp_index = pulse.zhsweepers.index(sweeper)
+                    sweeper.values = sweeper.values.copy()
+                    pulse.zhpulse.amplitude *= max(abs(sweeper.values))
+                    sweeper.values /= max(abs(sweeper.values))
+                else:
+                    sweeper_dur_index = pulse.zhsweepers.index(sweeper)
+
+            exp.play(
+                signal=channel_name,
+                pulse=pulse.zhpulse,
+                amplitude=pulse.zhsweepers[sweeper_amp_index],
+                length=pulse.zhsweepers[sweeper_dur_index],
+            )
+
+    def play_sweep(self, exp, channel_name, pulse):
+        """Takes care of playing the sweepers and involved pulses for different
+        options."""
+        parameters = []
+        for partial_sweep in pulse.zhsweepers:
+            parameters.append(partial_sweep.uid)
+        # Recheck partial sweeps
+        if len(parameters) == 2:
+            self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
+        else:
+            self.play_sweep_select_single(
+                exp, pulse, channel_name, parameters, partial_sweep
+            )
 
     @staticmethod
     def rearrange_sweepers(sweepers: List[Sweeper]) -> Tuple[np.ndarray, List[Sweeper]]:

From 01d07c5dccb0086a2069f04c15b30a2069f0031a Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 22 Feb 2024 10:44:32 +0400
Subject: [PATCH 24/45] use nt/rt _sweeps wherever needed

---
 src/qibolab/instruments/zhinst.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index e6dc5c55c..9433ab95a 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -261,6 +261,7 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
 
         self.sweepers = []
         self.nt_sweeps = []
+        self.rt_sweeps = []
         "Storing sweepers"
         # Improve the storing of multiple sweeps
 
@@ -586,8 +587,6 @@ def set_acquisition_type():
 
         set_acquisition_type()
 
-        self.nt_sweeps, self.sweepers = classify_sweepers(self.sweepers)
-
     def create_exp(self, qubits, couplers, options):
         """Zurich experiment initialization using their Experiment class."""
 
@@ -625,7 +624,7 @@ def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
             averaging_mode=exp_options.averaging_mode,
         ):
             # Recursion loop for sweepers or just play a sequence
-            if len(self.sweepers) > 0:
+            if len(self.rt_sweeps) > 0:
                 self.sweep_recursion(qubits, couplers, exp, exp_calib, exp_options)
             else:
                 self.select_exp(exp, qubits, couplers, exp_options)
@@ -856,9 +855,8 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
-        self.nt_sweeps = []
-        sweepers = list(sweepers)
-        rearranging_axes, sweepers = self.rearrange_sweepers(sweepers)
+        self.nt_sweeps, self.rt_sweeps = classify_sweepers(self.sweepers)
+        rearranging_axes, sweepers = self.rearrange_sweepers(list(sweepers))
         self.sweepers = sweepers
         # if using singleshot, the first axis contains shots,
         # i.e.: (nshots, sweeper_1, sweeper_2)
@@ -894,10 +892,10 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
     def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
         """Sweepers recursion for multiple nested Real Time sweepers."""
 
-        sweeper = self.sweepers[0]
+        sweeper = self.rt_sweeps[0]
 
-        i = len(self.sweepers) - 1
-        self.sweepers.remove(sweeper)
+        i = len(self.rt_sweeps) - 1
+        self.rt_sweeps.remove(sweeper)
         parameter = None
 
         if sweeper.parameter is Parameter.frequency:
@@ -906,7 +904,6 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 zhsweeper = select_sweeper(
                     sweeper, pulse.type, qubits[sweeper.pulses[0].qubit]
                 )
-                zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
                 if line == "readout":
                     channel_name = measure_channel_name(qubits[pulse.qubit])
                 else:
@@ -944,7 +941,7 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
             parameter=parameter,
             reset_oscillator_phase=True,
         ):
-            if len(self.sweepers) > 0:
+            if len(self.rt_sweeps) > 0:
                 self.sweep_recursion(qubits, couplers, exp, exp_calib, exp_options)
             else:
                 self.select_exp(exp, qubits, couplers, exp_options)

From dde6bea679830c254c7bd77e228830134fb9e643 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 22 Feb 2024 11:25:50 +0400
Subject: [PATCH 25/45] get rid of play_sweep_select_single/dual

---
 src/qibolab/instruments/zhinst.py | 80 ++++++-------------------------
 1 file changed, 15 insertions(+), 65 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 9433ab95a..4b53be8dc 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -754,72 +754,22 @@ def select_exp(self, exp, qubits, couplers, exp_options):
             previous_section = section_uid
 
     @staticmethod
-    def play_sweep_select_single(exp, pulse, channel_name, parameters, partial_sweep):
+    def play_sweep(exp, channel_name, pulse):
         """Play Zurich pulse when a single sweeper is involved."""
-        if any("amplitude" in param for param in parameters):
-            pulse.zhpulse.amplitude *= max(pulse.zhsweepers[0].values)
-            pulse.zhsweepers[0].values /= max(pulse.zhsweepers[0].values)
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                amplitude=pulse.zhsweepers[0],
-                phase=pulse.pulse.relative_phase,
-            )
-        elif any("duration" in param for param in parameters):
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                length=pulse.zhsweepers[0],
-                phase=pulse.pulse.relative_phase,
-            )
-        elif any("relative_phase" in param for param in parameters):
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                phase=pulse.zhsweepers[0],
-            )
-        elif "frequency" in partial_sweep.uid or partial_sweep.uid == "start":
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                phase=pulse.pulse.relative_phase,
-            )
-
-    # Hardcoded for the flux pulse for 2q gates
-    @staticmethod
-    def play_sweep_select_dual(exp, pulse, channel_name, parameters):
-        """Play Zurich pulse when two sweepers are involved on the same
-        pulse."""
-        if "amplitude" in parameters and "duration" in parameters:
-            for sweeper in pulse.zhsweepers:
-                if sweeper.uid == "amplitude":
-                    sweeper_amp_index = pulse.zhsweepers.index(sweeper)
-                    sweeper.values = sweeper.values.copy()
-                    pulse.zhpulse.amplitude *= max(abs(sweeper.values))
-                    sweeper.values /= max(abs(sweeper.values))
-                else:
-                    sweeper_dur_index = pulse.zhsweepers.index(sweeper)
-
-            exp.play(
-                signal=channel_name,
-                pulse=pulse.zhpulse,
-                amplitude=pulse.zhsweepers[sweeper_amp_index],
-                length=pulse.zhsweepers[sweeper_dur_index],
-            )
-
-    def play_sweep(self, exp, channel_name, pulse):
-        """Takes care of playing the sweepers and involved pulses for different
-        options."""
-        parameters = []
-        for partial_sweep in pulse.zhsweepers:
-            parameters.append(partial_sweep.uid)
-        # Recheck partial sweeps
-        if len(parameters) == 2:
-            self.play_sweep_select_dual(exp, pulse, channel_name, parameters)
-        else:
-            self.play_sweep_select_single(
-                exp, pulse, channel_name, parameters, partial_sweep
-            )
+        play_parameters = {}
+        for zhs in pulse.zhsweepers:
+            if zhs.uid == "amplitude":
+                pulse.zhpulse.amplitude *= max(zhs.values)
+                zhs.values /= max(zhs.values)
+                play_parameters["amplitude"] = zhs
+            if zhs.uid == "duration":
+                play_parameters["length"] = zhs
+            if zhs.uid == "relative_phase":
+                play_parameters["phase"] = zhs
+        if "phase" not in play_parameters:
+            play_parameters["phase"] = pulse.pulse.relative_phase
+
+        exp.play(signal=channel_name, pulse=pulse.zhpulse, **play_parameters)
 
     @staticmethod
     def rearrange_sweepers(sweepers: List[Sweeper]) -> Tuple[np.ndarray, List[Sweeper]]:

From 732d2f679b183061cd7b8abce2673dfefa1cafa6 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 22 Feb 2024 12:25:30 +0400
Subject: [PATCH 26/45] rearrange only rt sweeps

---
 src/qibolab/instruments/zhinst.py | 32 +++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 4b53be8dc..7ff16bbb0 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -788,31 +788,31 @@ def rearrange_sweepers(sweepers: List[Sweeper]) -> Tuple[np.ndarray, List[Sweepe
             sweepers (list[Sweeper]): updated list of sweepers used in the experiment. If
                 sweepers must be swapped, the list is updated accordingly.
         """
-        rearranging_axes = np.zeros(2, dtype=int)
-        if len(sweepers) == 2:
-            if sweepers[1].parameter is Parameter.frequency:
-                if not sweepers[0].pulses is None:
-                    if (sweepers[0].parameter is Parameter.bias) or (
-                        not sweepers[0].parameter is Parameter.amplitude
-                        and sweepers[0].pulses[0].type is not PulseType.READOUT
-                    ):
-                        rearranging_axes[:] = [1, 0]
-                        sweepers = sweepers[::-1]
-                        log.warning("Sweepers were reordered")
-        return rearranging_axes, sweepers
+        swapped_axis_pair = np.zeros(2, dtype=int)
+        freq_sweeper = next(
+            iter(s for s in sweepers if s.parameter is Parameter.frequency), None
+        )
+        if freq_sweeper:
+            freq_sweeper_idx = sweepers.index(freq_sweeper)
+            sweepers[freq_sweeper_idx] = sweepers[0]
+            sweepers[0] = freq_sweeper
+            swapped_axis_pair = np.array([0, freq_sweeper_idx])
+            log.warning("Sweepers were reordered")
+        return swapped_axis_pair, sweepers
 
     def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
         self.nt_sweeps, self.rt_sweeps = classify_sweepers(self.sweepers)
-        rearranging_axes, sweepers = self.rearrange_sweepers(list(sweepers))
-        self.sweepers = sweepers
+        swapped_axis_pair, self.rt_sweeps = self.rearrange_sweepers(self.rt_sweeps)
+        self.sweepers = list(sweepers)
+        swapped_axis_pair += len(self.nt_sweeps)
         # if using singleshot, the first axis contains shots,
         # i.e.: (nshots, sweeper_1, sweeper_2)
         # if using integration: (sweeper_1, sweeper_2)
         if options.averaging_mode is AveragingMode.SINGLESHOT:
-            rearranging_axes += 1
+            swapped_axis_pair += 1
 
         self.frequency_from_pulses(qubits, sequence)
 
@@ -827,7 +827,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
                 exp_res = self.results.get_data(f"sequence{q}_{i}")
 
                 # Reorder dimensions
-                data = np.moveaxis(exp_res, rearranging_axes[0], rearranging_axes[1])
+                data = np.moveaxis(exp_res, swapped_axis_pair[0], swapped_axis_pair[1])
                 if options.acquisition_type is AcquisitionType.DISCRIMINATION:
                     data = (
                         np.ones(data.shape) - data.real

From 0a9905712903921ce5fa6fb6a343aaf47ba38e5d Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 22 Feb 2024 16:40:44 +0400
Subject: [PATCH 27/45] don't use sweep parameter uid for decision making

---
 src/qibolab/instruments/zhinst.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 7ff16bbb0..e6dab3083 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -206,7 +206,9 @@ def __init__(self, pulse):
     # pylint: disable=R0903
     def add_sweeper(self, sweeper, qubit):
         """Add sweeper to list of sweepers associated with this pulse."""
-        self.zhsweepers.append(select_sweeper(sweeper, self.pulse.type, qubit))
+        self.zhsweepers.append(
+            (sweeper.parameter, select_sweeper(sweeper, self.pulse.type, qubit))
+        )
 
     def add_delay_sweeper(self, sweeper):
         self.delay_sweeper = select_sweeper(sweeper)
@@ -757,14 +759,14 @@ def select_exp(self, exp, qubits, couplers, exp_options):
     def play_sweep(exp, channel_name, pulse):
         """Play Zurich pulse when a single sweeper is involved."""
         play_parameters = {}
-        for zhs in pulse.zhsweepers:
-            if zhs.uid == "amplitude":
+        for p, zhs in pulse.zhsweepers:
+            if p is Parameter.amplitude:
                 pulse.zhpulse.amplitude *= max(zhs.values)
                 zhs.values /= max(zhs.values)
                 play_parameters["amplitude"] = zhs
-            if zhs.uid == "duration":
+            if p is Parameter.duration:
                 play_parameters["length"] = zhs
-            if zhs.uid == "relative_phase":
+            if p is Parameter.relative_phase:
                 play_parameters["phase"] = zhs
         if "phase" not in play_parameters:
             play_parameters["phase"] = pulse.pulse.relative_phase

From e84738e979d02fab3d516558a18db682723b0e94 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Fri, 23 Feb 2024 10:12:29 +0400
Subject: [PATCH 28/45] fix classified sweepers being empty

---
 src/qibolab/instruments/zhinst.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index e6dab3083..af4e22a90 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -806,9 +806,9 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
+        self.sweepers = list(sweepers)
         self.nt_sweeps, self.rt_sweeps = classify_sweepers(self.sweepers)
         swapped_axis_pair, self.rt_sweeps = self.rearrange_sweepers(self.rt_sweeps)
-        self.sweepers = list(sweepers)
         swapped_axis_pair += len(self.nt_sweeps)
         # if using singleshot, the first axis contains shots,
         # i.e.: (nshots, sweeper_1, sweeper_2)

From 7564dc0b33fc5b2dcd47f33942d22eeaa6857a2f Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Fri, 23 Feb 2024 10:22:05 +0400
Subject: [PATCH 29/45] fix acquisition type not being set to spectroscopy

---
 src/qibolab/instruments/zhinst.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index af4e22a90..d010bff22 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -581,7 +581,7 @@ def get_sweeps(pulse: Pulse):
 
         def set_acquisition_type():
             for sweeper in self.sweepers:
-                if sweeper.parameter.name in {Parameter.frequency, Parameter.amplitude}:
+                if sweeper.parameter in {Parameter.frequency, Parameter.amplitude}:
                     for pulse in sweeper.pulses:
                         if pulse.type is PulseType.READOUT:
                             # FIXME: case of multiple pulses

From a15ba0a51d23dbff19479e95ac544aa0dd503ff4 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Fri, 23 Feb 2024 15:01:03 +0400
Subject: [PATCH 30/45] decoupler sweep processing

---
 src/qibolab/instruments/zhinst.py | 350 +++++++++++++-----------------
 1 file changed, 152 insertions(+), 198 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index d010bff22..1e3bd8ea2 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -1,9 +1,9 @@
 """Instrument for using the Zurich Instruments (Zhinst) devices."""
 
-import copy
 from collections import defaultdict
+from copy import copy
 from dataclasses import dataclass, replace
-from typing import Dict, List, Tuple, Union
+from typing import Dict, Iterable, List, Tuple
 
 import laboneq.simple as lo
 import numpy as np
@@ -132,41 +132,7 @@ def select_pulse(pulse, pulse_type):
         )
 
 
-def select_sweeper(sweeper, ptype=None, qubit=None):
-    """Sweeper translation."""
-    if sweeper.parameter in (Parameter.duration, Parameter.start):
-        return lo.SweepParameter(
-            uid=sweeper.parameter.name,
-            values=sweeper.values * NANO_TO_SECONDS,
-        )
-    if sweeper.parameter is Parameter.frequency:
-        if ptype is None or qubit is None:
-            raise ValueError(
-                "For frequency sweep ptype and qubit arguments are mandatory"
-            )  # FIXME
-        if ptype is PulseType.READOUT:
-            intermediate_frequency = (
-                qubit.readout_frequency - qubit.readout.local_oscillator.frequency
-            )
-        elif ptype is PulseType.DRIVE:
-            intermediate_frequency = (
-                qubit.drive_frequency - qubit.drive.local_oscillator.frequency
-            )
-        else:
-            raise ValueError(
-                f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
-            )
-        return lo.LinearSweepParameter(
-            uid=sweeper.parameter.name,
-            start=sweeper.values[0] + intermediate_frequency,
-            stop=sweeper.values[-1] + intermediate_frequency,
-            count=len(sweeper.values),
-        )
-
-    return lo.SweepParameter(uid=sweeper.parameter.name, values=sweeper.values)
-
-
-def classify_sweepers(sweepers):
+def classify_sweepers(sweepers: Iterable[Sweeper]):
     """"""
     nt_sweepers = []
     rt_sweepers = []
@@ -204,14 +170,95 @@ def __init__(self, pulse):
         self.delay_sweeper = None
 
     # pylint: disable=R0903
-    def add_sweeper(self, sweeper, qubit):
+    def add_sweeper(self, param, sweeper):
         """Add sweeper to list of sweepers associated with this pulse."""
-        self.zhsweepers.append(
-            (sweeper.parameter, select_sweeper(sweeper, self.pulse.type, qubit))
-        )
+        if param.name in SWEEPER_SET:
+            self.zhsweepers.append((param, sweeper))
+        if param.name in SWEEPER_START:
+            # FIXME: add exception when delay sweeper already exists
+            self.delay_sweeper = sweeper
+
+
+class ProcessedSweeps:
+
+    def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
+        pulse_sweeps = []
+        channel_sweeps = []
+        parallel_sweeps = []
+        for sweeper in sweepers:
+            for pulse in sweeper.pulses or []:
+                if sweeper.parameter in (Parameter.duration, Parameter.start):
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values * NANO_TO_SECONDS
+                    )
+                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
+                elif sweeper.parameter is Parameter.frequency:
+                    ptype, qubit = pulse.type, qubits[pulse.qubit]
+                    if ptype is PulseType.READOUT:
+                        ch = measure_channel_name(qubit)
+                        intermediate_frequency = (
+                            qubit.readout_frequency
+                            - qubit.readout.local_oscillator.frequency
+                        )
+                    elif ptype is PulseType.DRIVE:
+                        ch = qubit.drive.name
+                        intermediate_frequency = (
+                            qubit.drive_frequency
+                            - qubit.drive.local_oscillator.frequency
+                        )
+                    else:
+                        raise ValueError(
+                            f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
+                        )
+                    # sweep_param = lo.LinearSweepParameter(start=sweeper.values[0] + intermediate_frequency, stop=sweeper.values[-1] + intermediate_frequency, count=len(sweeper.values))
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values + intermediate_frequency
+                    )
+                    channel_sweeps.append((ch, sweeper.parameter, sweep_param))
+                else:
+                    sweep_param = lo.SweepParameter(values=copy(sweeper.values))
+                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
+                parallel_sweeps.append((sweeper, sweep_param))
+
+            for qubit in sweeper.qubits or []:
+                if sweeper.parameter is not Parameter.bias:
+                    raise ValueError(
+                        f"Sweeping {sweeper.parameter.name} for {qubit} is not supported"
+                    )
+                sweep_param = lo.SweepParameter(
+                    values=sweeper.values + qubit.flux.offset
+                )
+                channel_sweeps.append((qubit.flux.name, sweeper.parameter, sweep_param))
+                parallel_sweeps.append((sweeper, sweep_param))
+
+            for coupler in sweeper.couplers or []:
+                if sweeper.parameter is not Parameter.bias:
+                    raise ValueError(
+                        f"Sweeping {sweeper.parameter.name} for {coupler} is not supported"
+                    )
+                sweep_param = lo.SweepParameter(
+                    values=sweeper.values + coupler.flux.offset
+                )
+                channel_sweeps.append(
+                    (coupler.flux.name, sweeper.parameter, sweep_param)
+                )
+                parallel_sweeps.append((sweeper, sweep_param))
+
+        self._pulse_sweeps = pulse_sweeps
+        self._channel_sweeps = channel_sweeps
+        self._parallel_sweeps = parallel_sweeps
 
-    def add_delay_sweeper(self, sweeper):
-        self.delay_sweeper = select_sweeper(sweeper)
+    def sweeps_for_pulse(self, pulse: Pulse):
+        return [item[1:] for item in self._pulse_sweeps if item[0] == pulse]
+
+    def sweeps_for_channel(self, ch: str):
+        return [item[1:] for item in self._channel_sweeps if item[0] == ch]
+
+    def sweeps_for_sweeper(self, sweeper: Sweeper):
+        return [item[1] for item in self._parallel_sweeps if item[0] == sweeper]
+
+    def channels_with_sweeps(self):
+        return {ch for ch, _, _ in self._channel_sweeps}
 
 
 @dataclass
@@ -261,7 +308,7 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
         self.sub_sequences = {}
         "Sub sequences between each measurement"
 
-        self.sweepers = []
+        self.processed_sweeps: ProcessedSweeps = None  # FIXME
         self.nt_sweeps = []
         self.rt_sweeps = []
         "Storing sweepers"
@@ -560,35 +607,16 @@ def sequence_zh(self, sequence, qubits, couplers):
                 ch = pulse.channel
             zhsequence[ch].append(ZhPulse(pulse))
 
-        def get_sweeps(pulse: Pulse):
-            sweepers_for_pulse = []
-            for sweeper in self.sweepers:
-                for p in sweeper.pulses or []:
-                    if p == pulse:
-                        sweepers_for_pulse.append(sweeper)
-                        break
-            return sweepers_for_pulse
-
-        for ch, zhpulses in zhsequence.items():
-            for i, zhpulse in enumerate(zhpulses):
-                for s in get_sweeps(zhpulse.pulse):
-                    if s.parameter.name in SWEEPER_SET:
-                        zhpulse.add_sweeper(s, qubits[zhpulse.pulse.qubit])
-                    if s.parameter.name in SWEEPER_START:
-                        zhpulse.add_delay_sweeper(s)
+        if self.processed_sweeps:
+            for ch, zhpulses in zhsequence.items():
+                for zhpulse in zhpulses:
+                    for param, sweep in self.processed_sweeps.sweeps_for_pulse(
+                        zhpulse.pulse
+                    ):
+                        zhpulse.add_sweeper(param, sweep)
 
         self.sequence = zhsequence
 
-        def set_acquisition_type():
-            for sweeper in self.sweepers:
-                if sweeper.parameter in {Parameter.frequency, Parameter.amplitude}:
-                    for pulse in sweeper.pulses:
-                        if pulse.type is PulseType.READOUT:
-                            # FIXME: case of multiple pulses
-                            self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
-
-        set_acquisition_type()
-
     def create_exp(self, qubits, couplers, options):
         """Zurich experiment initialization using their Experiment class."""
 
@@ -610,14 +638,13 @@ def create_exp(self, qubits, couplers, options):
             options, acquisition_type=acquisition_type, averaging_mode=averaging_mode
         )
 
-        exp_calib = lo.Calibration()
         # Near Time recursion loop or directly to Real Time recursion loop
         if self.nt_sweeps:
-            self.sweep_recursion_nt(qubits, couplers, exp_options, exp, exp_calib)
+            self.sweep_recursion_nt(qubits, couplers, exp_options, exp)
         else:
-            self.define_exp(qubits, couplers, exp_options, exp, exp_calib)
+            self.define_exp(qubits, couplers, exp_options, exp)
 
-    def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
+    def define_exp(self, qubits, couplers, exp_options, exp):
         """Real time definition."""
         with exp.acquire_loop_rt(
             uid="shots",
@@ -627,13 +654,31 @@ def define_exp(self, qubits, couplers, exp_options, exp, exp_calib):
         ):
             # Recursion loop for sweepers or just play a sequence
             if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, couplers, exp, exp_calib, exp_options)
+                self.sweep_recursion(qubits, couplers, exp, exp_options)
             else:
                 self.select_exp(exp, qubits, couplers, exp_options)
-            exp.set_calibration(exp_calib)
+            self.set_calibration(exp)
             exp.set_signal_map(self.signal_map)
             self.experiment = exp
 
+    def set_calibration(self, exp):
+        calib = lo.Calibration()
+        if self.processed_sweeps:
+            for ch in (
+                set(self.sequence.keys()) | self.processed_sweeps.channels_with_sweeps()
+            ):
+                for param, sweep_param in self.processed_sweeps.sweeps_for_channel(ch):
+                    if param is Parameter.frequency:
+                        calib[ch] = lo.SignalCalibration(
+                            oscillator=lo.Oscillator(
+                                frequency=sweep_param,
+                                modulation_type=lo.ModulationType.HARDWARE,
+                            )
+                        )
+                    if param is Parameter.bias:
+                        calib[ch] = lo.SignalCalibration(voltage_offset=sweep_param)
+                exp.set_calibration(calib)
+
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
         weights = {}
@@ -731,6 +776,15 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                             weight = weights[q]
 
                     measure_pulse_parameters = {"phase": 0}
+                    amplitude = None
+                    if self.processed_sweeps:
+                        for param, sweep in self.processed_sweeps.sweeps_for_pulse(
+                            pulse.pulse
+                        ):
+                            if param is Parameter.amplitude:
+                                pulse.zhpulse.amplitude *= max(sweep.values)
+                                sweep.values /= max(sweep.values)
+                                amplitude = sweep
 
                     if i == len(self.sequence[measure_channel_name(qubit)]) - 1:
                         reset_delay = exp_options.relaxation_time * NANO_TO_SECONDS
@@ -749,7 +803,7 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                             pulse.pulse.duration * NANO_TO_SECONDS, 9
                         ),
                         measure_pulse_parameters=measure_pulse_parameters,
-                        measure_pulse_amplitude=None,
+                        measure_pulse_amplitude=amplitude,
                         acquire_delay=self.time_of_flight * NANO_TO_SECONDS,
                         reset_delay=reset_delay,
                     )
@@ -806,8 +860,8 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
-        self.sweepers = list(sweepers)
-        self.nt_sweeps, self.rt_sweeps = classify_sweepers(self.sweepers)
+        self.processed_sweeps = ProcessedSweeps(sweepers, qubits)
+        self.nt_sweeps, self.rt_sweeps = classify_sweepers(sweepers)
         swapped_axis_pair, self.rt_sweeps = self.rearrange_sweepers(self.rt_sweeps)
         swapped_axis_pair += len(self.nt_sweeps)
         # if using singleshot, the first axis contains shots,
@@ -818,6 +872,13 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
 
         self.frequency_from_pulses(qubits, sequence)
 
+        for sweeper in sweepers:
+            if sweeper.parameter in {Parameter.frequency, Parameter.amplitude}:
+                for pulse in sweeper.pulses:
+                    if pulse.type is PulseType.READOUT:
+                        # FIXME: case of multiple pulses
+                        self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
+
         self.experiment_flow(qubits, couplers, sequence, options)
         self.run_exp()
 
@@ -841,99 +902,33 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
 
         return results
 
-    def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
+    def sweep_recursion(self, qubits, couplers, exp, exp_options):
         """Sweepers recursion for multiple nested Real Time sweepers."""
 
         sweeper = self.rt_sweeps[0]
 
         i = len(self.rt_sweeps) - 1
         self.rt_sweeps.remove(sweeper)
-        parameter = None
-
-        if sweeper.parameter is Parameter.frequency:
-            for pulse in sweeper.pulses:
-                line = "drive" if pulse.type is PulseType.DRIVE else "readout"
-                zhsweeper = select_sweeper(
-                    sweeper, pulse.type, qubits[sweeper.pulses[0].qubit]
-                )
-                if line == "readout":
-                    channel_name = measure_channel_name(qubits[pulse.qubit])
-                else:
-                    channel_name = getattr(qubits[pulse.qubit], line).name
-                exp_calib[channel_name] = lo.SignalCalibration(
-                    oscillator=lo.Oscillator(
-                        frequency=zhsweeper,
-                        modulation_type=lo.ModulationType.HARDWARE,
-                    )
-                )
-        if sweeper.parameter is Parameter.amplitude:
-            for pulse in sweeper.pulses:
-                pulse = pulse.copy()
-                pulse.amplitude *= max(abs(sweeper.values))
-
-                aux_max = max(abs(sweeper.values))
-
-                sweeper.values /= aux_max
-                parameter = select_sweeper(sweeper)
-                sweeper.values *= aux_max
-
-        if sweeper.parameter is Parameter.bias:
-            parameter = select_sweeper(sweeper)
-
-        elif sweeper.parameter is Parameter.start:
-            parameter = select_sweeper(sweeper)
-
-        elif parameter is None:
-            parameter = select_sweeper(
-                sweeper, sweeper.pulses[0].type, qubits[sweeper.pulses[0].qubit]
-            )
 
         with exp.sweep(
             uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
-            parameter=parameter,
+            parameter=[
+                sweep_param
+                for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
+            ],
             reset_oscillator_phase=True,
         ):
             if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, couplers, exp, exp_calib, exp_options)
+                self.sweep_recursion(qubits, couplers, exp, exp_options)
             else:
                 self.select_exp(exp, qubits, couplers, exp_options)
 
-    def find_instrument_address(
-        self, quantum_element: Union[Qubit, Coupler], parameter: str
-    ) -> str:
-        """Find path of the instrument connected to a specified line and
-        qubit/coupler.
-
-        Args:
-            quantum_element (Qubit | Coupler): qubits or couplers on which perform the near time sweep.
-            parameter (str): parameter on which perform the near time sweep.
-        """
-        line_names = {
-            "bias": "flux",
-            "amplitude": "drive",
-        }
-        line_name = line_names[parameter]
-        channel_uid = (
-            self.device_setup.logical_signal_groups[f"q{quantum_element.name}"]
-            .logical_signals[f"{line_name}_line"]
-            .physical_channel.uid
-        )
-        channel_name = channel_uid.split("/")[0]
-        instruments = self.device_setup.instruments
-        for instrument in instruments:
-            if instrument.uid == channel_name:
-                return instrument.address
-        raise RuntimeError(
-            f"Could not find instrument for {quantum_element} {line_name}"
-        )
-
     def sweep_recursion_nt(
         self,
         qubits: Dict[str, Qubit],
         couplers: Dict[str, Coupler],
         options: ExecutionParameters,
         exp: lo.Experiment,
-        exp_calib: lo.Calibration,
     ):
         """Sweepers recursion for Near Time sweepers. Faster than regular
         software sweepers as they are executed on the actual device by
@@ -950,58 +945,17 @@ def sweep_recursion_nt(
         i = len(self.nt_sweeps) - 1
         self.nt_sweeps.remove(sweeper)
 
-        parameter = None
-
-        if sweeper.parameter is Parameter.bias:
-            if sweeper.qubits:
-                for qubit in sweeper.qubits:
-                    zhsweeper = select_sweeper(sweeper)
-                    zhsweeper.uid = "bias"
-                    path = self.find_instrument_address(qubit, "bias")
-
-                    parameter = copy.deepcopy(zhsweeper)
-                    parameter.values += qubit.flux.offset
-                    device_path = f"{path}/sigouts/0/offset"
-
-        elif sweeper.parameter is Parameter.amplitude:
-            for pulse in sweeper.pulses:
-                pulse = pulse.copy()
-                pulse.amplitude *= max(abs(sweeper.values))
-
-                aux_max = max(abs(sweeper.values))
-
-                sweeper.values /= aux_max
-                zhsweeper = select_sweeper(sweeper)
-                sweeper.values *= aux_max
-
-                zhsweeper.uid = "amplitude"
-                path = self.find_instrument_address(
-                    qubits[sweeper.pulses[0].qubit], "amplitude"
-                )
-                parameter = zhsweeper
-                device_path = (
-                    f"/{path}/qachannels/*/oscs/0/gain"  # Hardcoded SHFQA device
-                )
-
-        elif parameter is None:
-            parameter = select_sweeper(
-                sweeper, sweeper.pulses[0].type, qubits[sweeper.pulses[0].qubit]
-            )
-            device_path = f"/{path}/qachannels/*/oscs/0/gain"  # Hardcoded SHFQA device
-
         with exp.sweep(
             uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
-            parameter=parameter,
+            parameter=[
+                sweep_param
+                for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
+            ],
         ):
-            exp.set_node(
-                path=device_path,
-                value=parameter,
-            )
-
             if len(self.nt_sweeps) > 0:
-                self.sweep_recursion_nt(qubits, couplers, options, exp, exp_calib)
+                self.sweep_recursion_nt(qubits, couplers, options, exp)
             else:
-                self.define_exp(qubits, couplers, options, exp, exp_calib)
+                self.define_exp(qubits, couplers, options, exp)
 
     def split_batches(self, sequences):
         return batch_max_sequences(sequences, MAX_SEQUENCES)

From 02448912efeb40dbae46473d6a1c1723ff4e664d Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 27 Feb 2024 10:45:15 +0400
Subject: [PATCH 31/45] polish handling of channel parameter sweeps

---
 src/qibolab/instruments/zhinst.py | 64 +++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 1e3bd8ea2..8c4abff28 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -1,5 +1,6 @@
 """Instrument for using the Zurich Instruments (Zhinst) devices."""
 
+import re
 from collections import defaultdict
 from copy import copy
 from dataclasses import dataclass, replace
@@ -215,6 +216,20 @@ def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
                         values=sweeper.values + intermediate_frequency
                     )
                     channel_sweeps.append((ch, sweeper.parameter, sweep_param))
+                elif (
+                    pulse.type is PulseType.READOUT
+                    and sweeper.parameter is Parameter.amplitude
+                ):
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values / max(sweeper.values)
+                    )
+                    channel_sweeps.append(
+                        (
+                            measure_channel_name(qubits[pulse.qubit]),
+                            sweeper.parameter,
+                            sweep_param,
+                        )
+                    )
                 else:
                     sweep_param = lo.SweepParameter(values=copy(sweeper.values))
                     pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
@@ -662,22 +677,46 @@ def define_exp(self, qubits, couplers, exp_options, exp):
             self.experiment = exp
 
     def set_calibration(self, exp):
-        calib = lo.Calibration()
         if self.processed_sweeps:
             for ch in (
                 set(self.sequence.keys()) | self.processed_sweeps.channels_with_sweeps()
             ):
                 for param, sweep_param in self.processed_sweeps.sweeps_for_channel(ch):
                     if param is Parameter.frequency:
-                        calib[ch] = lo.SignalCalibration(
-                            oscillator=lo.Oscillator(
-                                frequency=sweep_param,
-                                modulation_type=lo.ModulationType.HARDWARE,
+                        exp.set_calibration(
+                            lo.Calibration(
+                                {
+                                    ch: lo.SignalCalibration(
+                                        oscillator=lo.Oscillator(
+                                            frequency=sweep_param,
+                                            modulation_type=lo.ModulationType.HARDWARE,
+                                        )
+                                    )
+                                }
                             )
                         )
+
+                    # for the remaining cases we can only achieve our goal by manipulating the instrument node directly
+                    channel_node_path = None
+                    logical_signal = self.signal_map[ch]
+                    for instrument in self.device_setup.instruments:
+                        for conn in instrument.connections:
+                            if conn.remote_path == logical_signal.path:
+                                channel_node_path = (
+                                    f"{instrument.address}/{conn.local_port}"
+                                )
+                                break
+                    if channel_node_path is None:
+                        raise RuntimeError(
+                            f"Could not find instrument node corresponding to {param.name} of channel {ch}"
+                        )
                     if param is Parameter.bias:
-                        calib[ch] = lo.SignalCalibration(voltage_offset=sweep_param)
-                exp.set_calibration(calib)
+                        offset_node_path = f"{channel_node_path}/offset"
+                        exp.set_node(path=offset_node_path, value=sweep_param)
+                    if param is Parameter.amplitude:
+                        a, b = re.match(r"(.*)/(\d)/.*", channel_node_path).groups()
+                        gain_node_path = f"{a}/{b}/oscs/{b}/gain"
+                        exp.set_node(path=gain_node_path, value=sweep_param)
 
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
@@ -776,15 +815,6 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                             weight = weights[q]
 
                     measure_pulse_parameters = {"phase": 0}
-                    amplitude = None
-                    if self.processed_sweeps:
-                        for param, sweep in self.processed_sweeps.sweeps_for_pulse(
-                            pulse.pulse
-                        ):
-                            if param is Parameter.amplitude:
-                                pulse.zhpulse.amplitude *= max(sweep.values)
-                                sweep.values /= max(sweep.values)
-                                amplitude = sweep
 
                     if i == len(self.sequence[measure_channel_name(qubit)]) - 1:
                         reset_delay = exp_options.relaxation_time * NANO_TO_SECONDS
@@ -803,7 +833,7 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                             pulse.pulse.duration * NANO_TO_SECONDS, 9
                         ),
                         measure_pulse_parameters=measure_pulse_parameters,
-                        measure_pulse_amplitude=amplitude,
+                        measure_pulse_amplitude=None,
                         acquire_delay=self.time_of_flight * NANO_TO_SECONDS,
                         reset_delay=reset_delay,
                     )

From a9f63ec06a61350901302e20091fb8afa467deed Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 27 Feb 2024 19:29:38 +0400
Subject: [PATCH 32/45] fix the place where instrument node should be set

---
 src/qibolab/instruments/zhinst.py | 74 +++++++++++++++++--------------
 1 file changed, 41 insertions(+), 33 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 8c4abff28..03a5e23a0 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -272,6 +272,13 @@ def sweeps_for_channel(self, ch: str):
     def sweeps_for_sweeper(self, sweeper: Sweeper):
         return [item[1] for item in self._parallel_sweeps if item[0] == sweeper]
 
+    def channel_sweeps_for_sweeper(self, sweeper: Sweeper):
+        return [
+            item
+            for item in self._channel_sweeps
+            if item[2] in self.sweeps_for_sweeper(sweeper)
+        ]
+
     def channels_with_sweeps(self):
         return {ch for ch, _, _ in self._channel_sweeps}
 
@@ -672,51 +679,51 @@ def define_exp(self, qubits, couplers, exp_options, exp):
                 self.sweep_recursion(qubits, couplers, exp, exp_options)
             else:
                 self.select_exp(exp, qubits, couplers, exp_options)
-            self.set_calibration(exp)
+            self.set_calibration_for_rt_sweep(exp)
             exp.set_signal_map(self.signal_map)
             self.experiment = exp
 
-    def set_calibration(self, exp):
+    def set_calibration_for_rt_sweep(self, exp):
         if self.processed_sweeps:
+            calib = lo.Calibration()
             for ch in (
                 set(self.sequence.keys()) | self.processed_sweeps.channels_with_sweeps()
             ):
                 for param, sweep_param in self.processed_sweeps.sweeps_for_channel(ch):
                     if param is Parameter.frequency:
-                        exp.set_calibration(
-                            lo.Calibration(
-                                {
-                                    ch: lo.SignalCalibration(
-                                        oscillator=lo.Oscillator(
-                                            frequency=sweep_param,
-                                            modulation_type=lo.ModulationType.HARDWARE,
-                                        )
-                                    )
-                                }
+                        calib[ch] = lo.SignalCalibration(
+                            oscillator=lo.Oscillator(
+                                frequency=sweep_param,
+                                modulation_type=lo.ModulationType.HARDWARE,
                             )
                         )
+            exp.set_calibration(calib)
 
-                    # for the remaining cases we can only achieve our goal by manipulating the instrument node directly
-                    channel_node_path = None
-                    logical_signal = self.signal_map[ch]
-                    for instrument in self.device_setup.instruments:
-                        for conn in instrument.connections:
-                            if conn.remote_path == logical_signal.path:
-                                channel_node_path = (
-                                    f"{instrument.address}/{conn.local_port}"
-                                )
-                                break
-                    if channel_node_path is None:
-                        raise RuntimeError(
-                            f"Could not find instrument node corresponding to {param.name} of channel {ch}"
-                        )
-                    if param is Parameter.bias:
-                        offset_node_path = f"{channel_node_path}/offset"
-                        exp.set_node(path=offset_node_path, value=sweep_param)
-                    if param is Parameter.amplitude:
-                        a, b = re.match(r"(.*)/(\d)/.*", channel_node_path).groups()
-                        gain_node_path = f"{a}/{b}/oscs/{b}/gain"
-                        exp.set_node(path=gain_node_path, value=sweep_param)
+    def set_instrument_nodes_for_nt_sweep(
+        self, exp: lo.Experiment, sweeper: Sweeper
+    ) -> None:
+        for ch, param, sweep_param in self.processed_sweeps.channel_sweeps_for_sweeper(
+            sweeper
+        ):
+            # for the remaining cases we can only achieve our goal by manipulating the instrument node directly
+            channel_node_path = None
+            logical_signal = self.signal_map[ch]
+            for instrument in self.device_setup.instruments:
+                for conn in instrument.connections:
+                    if conn.remote_path == logical_signal.path:
+                        channel_node_path = f"{instrument.address}/{conn.local_port}"
+                        break
+            if channel_node_path is None:
+                raise RuntimeError(
+                    f"Could not find instrument node corresponding to {param.name} of channel {ch}"
+                )
+            if param is Parameter.bias:
+                offset_node_path = f"{channel_node_path}/offset"
+                exp.set_node(path=offset_node_path, value=sweep_param)
+            if param is Parameter.amplitude:
+                a, b = re.match(r"(.*)/(\d)/.*", channel_node_path).groups()
+                gain_node_path = f"{a}/{b}/oscs/{b}/gain"
+                exp.set_node(path=gain_node_path, value=sweep_param)
 
     def select_exp(self, exp, qubits, couplers, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
@@ -982,6 +989,7 @@ def sweep_recursion_nt(
                 for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
             ],
         ):
+            self.set_instrument_nodes_for_nt_sweep(exp, sweeper)
             if len(self.nt_sweeps) > 0:
                 self.sweep_recursion_nt(qubits, couplers, options, exp)
             else:

From 58a734edd0362111607c626335c3b0a82273021a Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 27 Feb 2024 22:01:06 +0400
Subject: [PATCH 33/45] polish docstrings and type annotations

---
 src/qibolab/instruments/zhinst.py | 250 ++++++++++++++++--------------
 1 file changed, 134 insertions(+), 116 deletions(-)

diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index 03a5e23a0..be4e825a0 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -2,9 +2,10 @@
 
 import re
 from collections import defaultdict
+from collections.abc import Iterable
 from copy import copy
 from dataclasses import dataclass, replace
-from typing import Dict, Iterable, List, Tuple
+from typing import Optional
 
 import laboneq.simple as lo
 import numpy as np
@@ -69,14 +70,12 @@ def acquire_channel_name(qubit: Qubit) -> str:
     return f"acquire{qubit.name}"
 
 
-def select_pulse(pulse, pulse_type):
-    """Pulse translation."""
-
+def select_pulse(pulse: Pulse):
+    """Return laboneq pulse object corresponding to the given qibolab pulse."""
     if "IIR" not in str(pulse.shape):
         if str(pulse.shape) == "Rectangular()":
             can_compress = pulse.type is not PulseType.READOUT
             return lo.pulse_library.const(
-                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 can_compress=can_compress,
@@ -84,7 +83,6 @@ def select_pulse(pulse, pulse_type):
         if "Gaussian" in str(pulse.shape):
             sigma = pulse.shape.rel_sigma
             return lo.pulse_library.gaussian(
-                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 sigma=2 / sigma,
@@ -96,7 +94,6 @@ def select_pulse(pulse, pulse_type):
             width = pulse.shape.width
             can_compress = pulse.type is not PulseType.READOUT
             return lo.pulse_library.gaussian_square(
-                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
                 amplitude=pulse.amplitude,
@@ -109,7 +106,6 @@ def select_pulse(pulse, pulse_type):
             sigma = pulse.shape.rel_sigma
             beta = pulse.shape.beta
             return lo.pulse_library.drag(
-                uid=f"{pulse_type}_{pulse.qubit}_",
                 length=round(pulse.duration * NANO_TO_SECONDS, 9),
                 amplitude=pulse.amplitude,
                 sigma=2 / sigma,
@@ -119,24 +115,23 @@ def select_pulse(pulse, pulse_type):
 
     if np.all(pulse.envelope_waveform_q(SAMPLING_RATE).data == 0):
         return sampled_pulse_real(
-            uid=f"{pulse_type}_{pulse.qubit}_",
             samples=pulse.envelope_waveform_i(SAMPLING_RATE).data,
             can_compress=True,
         )
     else:
-        # Test this when we have pulses that use it
         return sampled_pulse_complex(
-            uid=f"{pulse_type}_{pulse.qubit}_",
             samples=pulse.envelope_waveform_i(SAMPLING_RATE).data
             + (1j * pulse.envelope_waveform_q(SAMPLING_RATE).data),
             can_compress=True,
         )
 
 
-def classify_sweepers(sweepers: Iterable[Sweeper]):
-    """"""
-    nt_sweepers = []
-    rt_sweepers = []
+def classify_sweepers(
+    sweepers: Iterable[Sweeper],
+) -> tuple[list[Sweeper], list[Sweeper]]:
+    """Divide sweepers into two lists: 1. sweeps that can be done in the laboneq near-time sweep loop, 2. sweeps that
+    can be done in real-time (i.e. on hardware)"""
+    nt_sweepers, rt_sweepers = [], []
     for sweeper in sweepers:
         if (
             sweeper.parameter is Parameter.amplitude
@@ -152,35 +147,43 @@ def classify_sweepers(sweepers: Iterable[Sweeper]):
 
 @dataclass
 class ZhPort(Port):
-    name: Tuple[str, str]
+    name: tuple[str, str]
     offset: float = 0.0
     power_range: int = 0
 
 
 class ZhPulse:
-    """Zurich pulse from qibolab pulse translation."""
+    """Wrapper data type that holds a qibolab pulse, the corresponding laboneq
+    pulse object, and any sweeps associated with this pulse."""
 
     def __init__(self, pulse):
-        """Zurich pulse from qibolab pulse."""
-        self.pulse = pulse
+        self.pulse: Pulse = pulse
         """Qibolab pulse."""
-        self.zhpulse = select_pulse(pulse, pulse.type.name.lower())
-        """Zurich pulse."""
-        self.zhsweepers = []
-
-        self.delay_sweeper = None
+        self.zhpulse = select_pulse(pulse)
+        """Laboneq pulse."""
+        self.zhsweepers: list[tuple[Parameter, lo.SweepParameter]] = []
+        """Parameters to be swept, along with their laboneq sweep parameter
+        definitions."""
+        self.delay_sweeper: Optional[lo.SweepParameter] = None
+        """Laboneq sweep parameter if the delay of the pulse should be
+        swept."""
 
     # pylint: disable=R0903
-    def add_sweeper(self, param, sweeper):
+    def add_sweeper(self, param: Parameter, sweeper: lo.SweepParameter):
         """Add sweeper to list of sweepers associated with this pulse."""
         if param.name in SWEEPER_SET:
             self.zhsweepers.append((param, sweeper))
         if param.name in SWEEPER_START:
-            # FIXME: add exception when delay sweeper already exists
+            if self.delay_sweeper:
+                raise ValueError(
+                    "Cannot have multiple delay sweepers for a single pulse"
+                )
             self.delay_sweeper = sweeper
 
 
 class ProcessedSweeps:
+    """Data type that centralizes and allows extracting information about given
+    sweeps."""
 
     def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
         pulse_sweeps = []
@@ -211,7 +214,6 @@ def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
                         raise ValueError(
                             f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
                         )
-                    # sweep_param = lo.LinearSweepParameter(start=sweeper.values[0] + intermediate_frequency, stop=sweeper.values[-1] + intermediate_frequency, count=len(sweeper.values))
                     sweep_param = lo.SweepParameter(
                         values=sweeper.values + intermediate_frequency
                     )
@@ -263,34 +265,55 @@ def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
         self._channel_sweeps = channel_sweeps
         self._parallel_sweeps = parallel_sweeps
 
-    def sweeps_for_pulse(self, pulse: Pulse):
+    def sweeps_for_pulse(
+        self, pulse: Pulse
+    ) -> list[tuple[Parameter, lo.SweepParameter]]:
         return [item[1:] for item in self._pulse_sweeps if item[0] == pulse]
 
-    def sweeps_for_channel(self, ch: str):
+    def sweeps_for_channel(self, ch: str) -> list[tuple[Parameter, lo.SweepParameter]]:
         return [item[1:] for item in self._channel_sweeps if item[0] == ch]
 
-    def sweeps_for_sweeper(self, sweeper: Sweeper):
+    def sweeps_for_sweeper(self, sweeper: Sweeper) -> list[lo.SweepParameter]:
         return [item[1] for item in self._parallel_sweeps if item[0] == sweeper]
 
-    def channel_sweeps_for_sweeper(self, sweeper: Sweeper):
+    def channel_sweeps_for_sweeper(
+        self, sweeper: Sweeper
+    ) -> list[tuple[str, Parameter, lo.SweepParameter]]:
         return [
             item
             for item in self._channel_sweeps
             if item[2] in self.sweeps_for_sweeper(sweeper)
         ]
 
-    def channels_with_sweeps(self):
+    def channels_with_sweeps(self) -> set[str]:
         return {ch for ch, _, _ in self._channel_sweeps}
 
 
 @dataclass
 class SubSequence:
-    measurements: list[str, Pulse]
-    control_sequence: dict[str, list[Pulse]]
+    """A subsequence is a slice (in time) of a sequence that contains at most
+    one measurement per qubit.
+
+    When the driver is asked to execute a sequence, it will first split
+    it into sub-sequences. This is needed so that we can create a
+    separate laboneq section for each measurement (multiple measurements
+    per section are not allowed). When splitting a sequence, it is
+    assumed that 1. a measurement operation can be parallel (in time) to
+    another measurement operation (i.e. measuring multiple qubits
+    simultaneously), but other channels (e.g. drive) do not contain any
+    pulses parallel to measurements, 2. ith measurement on some channel
+    is in the same subsequence as the ith measurement (if any) on
+    another measurement channel, 3. all measurements in one subsequence
+    happen at the same time.
+    """
+
+    measurements: list[tuple[str, ZhPulse]]
+    control_sequence: dict[str, list[ZhPulse]]
 
 
 class Zurich(Controller):
-    """Zurich driver main class."""
+    """Driver for a collection of ZI instruments that are automatically
+    synchronized via ZSync protocol."""
 
     PortType = ZhPort
 
@@ -310,7 +333,6 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
         self.smearing = smearing
         "Parameters read from the runcard not part of ExecutionParameters"
 
-        self.exp = None
         self.experiment = None
         self.results = None
         "Zurich experiment definitions"
@@ -326,15 +348,12 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
 
         self.sequence = defaultdict(list)
         "Zurich pulse sequence"
-        self.sequence_qibo = None
-        self.sub_sequences = {}
+        self.sub_sequences: list[SubSequence] = []
         "Sub sequences between each measurement"
 
-        self.processed_sweeps: ProcessedSweeps = None  # FIXME
-        self.nt_sweeps = []
-        self.rt_sweeps = []
-        "Storing sweepers"
-        # Improve the storing of multiple sweeps
+        self.processed_sweeps: Optional[ProcessedSweeps] = None
+        self.nt_sweeps: list[Sweeper] = []
+        self.rt_sweeps: list[Sweeper] = []
 
     @property
     def sampling_rate(self):
@@ -413,7 +432,6 @@ def register_readout_line(self, qubit, intermediate_frequency, options):
                     modulation_type=lo.ModulationType.SOFTWARE,
                 ),
                 local_oscillator=lo.Oscillator(
-                    uid="lo_shfqa_m" + str(q),
                     frequency=int(qubit.readout.local_oscillator.frequency),
                 ),
                 range=qubit.readout.power_range,
@@ -464,7 +482,6 @@ def register_drive_line(self, qubit, intermediate_frequency):
                     modulation_type=lo.ModulationType.HARDWARE,
                 ),
                 local_oscillator=lo.Oscillator(
-                    uid="lo_shfqc" + str(q),
                     frequency=int(qubit.drive.local_oscillator.frequency),
                 ),
                 range=qubit.drive.power_range,
@@ -511,10 +528,10 @@ def run_exp(self):
         - Save a experiment compiled experiment ():
         self.exp.save("saved_exp")  # saving compiled experiment
         """
-        self.exp = self.session.compile(
+        compiled_experiment = self.session.compile(
             self.experiment, compiler_settings=COMPILER_SETTINGS
         )
-        self.results = self.session.run(self.exp)
+        self.results = self.session.run(compiled_experiment)
 
     @staticmethod
     def frequency_from_pulses(qubits, sequence):
@@ -527,15 +544,7 @@ def frequency_from_pulses(qubits, sequence):
                 qubit.drive_frequency = pulse.frequency
 
     def create_sub_sequences(self, qubits: list[Qubit]) -> list[SubSequence]:
-        """Create subsequences for different lines (drive, flux, coupler flux).
-
-        Args:
-            qubits (dict[str, Qubit]): qubits for the platform.
-            couplers (dict[str, Coupler]): couplers for the platform.
-        """
-        # collect and group all measurements. In one sequence, qubits may be measured multiple times (e.g. when the sequence was unrolled
-        # at platform level), however it is assumed that 1. all qubits are measured the same number of times, 2. each time all qubits are
-        # measured simultaneously.
+        """Create subsequences based on locations of measurements."""
         measure_channels = {measure_channel_name(qb) for qb in qubits}
         other_channels = set(self.sequence.keys()) - measure_channels
 
@@ -547,7 +556,6 @@ def create_sub_sequences(self, qubits: list[Qubit]) -> list[SubSequence]:
         measurement_starts = {}
         for i, group in measurement_groups.items():
             starts = np.array([meas.pulse.start for _, meas in group])
-            # TODO: validate that all start times are equal
             measurement_starts[i] = max(starts)
 
         # split all non-measurement channels according to the locations of the measurements
@@ -568,8 +576,8 @@ def create_sub_sequences(self, qubits: list[Qubit]) -> list[SubSequence]:
 
     def experiment_flow(
         self,
-        qubits: Dict[str, Qubit],
-        couplers: Dict[str, Coupler],
+        qubits: dict[str, Qubit],
+        couplers: dict[str, Coupler],
         sequence: PulseSequence,
         options: ExecutionParameters,
     ):
@@ -583,10 +591,10 @@ def experiment_flow(
             couplers (dict[str, Coupler]): couplers for the platform.
             sequence (PulseSequence): sequence of pulses to be played in the experiment.
         """
-        self.sequence_zh(sequence, qubits, couplers)
+        self.sequence = self.sequence_zh(sequence, qubits)
         self.sub_sequences = self.create_sub_sequences(list(qubits.values()))
         self.calibration_step(qubits, couplers, options)
-        self.create_exp(qubits, couplers, options)
+        self.create_exp(qubits, options)
 
     # pylint: disable=W0221
     def play(self, qubits, couplers, sequence, options):
@@ -615,11 +623,17 @@ def play(self, qubits, couplers, sequence, options):
 
         return results
 
-    def sequence_zh(self, sequence, qubits, couplers):
-        """Qibo sequence to Zurich sequence."""
+    def sequence_zh(
+        self, sequence: PulseSequence, qubits: dict[str, Qubit]
+    ) -> dict[str, list[ZhPulse]]:
+        """Convert Qibo sequence to a sequence where all pulses are replaced
+        with ZhPulse instances.
+
+        The resulting object is a dictionary mapping from channel name
+        to corresponding sequence of ZhPulse instances
+        """
         # Define and assign the sequence
         zhsequence = defaultdict(list)
-        self.sequence_qibo = sequence
 
         # Fill the sequences with pulses according to their lines in temporal order
         for pulse in sequence:
@@ -637,9 +651,9 @@ def sequence_zh(self, sequence, qubits, couplers):
                     ):
                         zhpulse.add_sweeper(param, sweep)
 
-        self.sequence = zhsequence
+        return zhsequence
 
-    def create_exp(self, qubits, couplers, options):
+    def create_exp(self, qubits, options):
         """Zurich experiment initialization using their Experiment class."""
 
         # Setting experiment signal lines
@@ -662,11 +676,11 @@ def create_exp(self, qubits, couplers, options):
 
         # Near Time recursion loop or directly to Real Time recursion loop
         if self.nt_sweeps:
-            self.sweep_recursion_nt(qubits, couplers, exp_options, exp)
+            self.sweep_recursion_nt(qubits, exp_options, exp)
         else:
-            self.define_exp(qubits, couplers, exp_options, exp)
+            self.define_exp(qubits, exp_options, exp)
 
-    def define_exp(self, qubits, couplers, exp_options, exp):
+    def define_exp(self, qubits, exp_options, exp):
         """Real time definition."""
         with exp.acquire_loop_rt(
             uid="shots",
@@ -676,14 +690,16 @@ def define_exp(self, qubits, couplers, exp_options, exp):
         ):
             # Recursion loop for sweepers or just play a sequence
             if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, couplers, exp, exp_options)
+                self.sweep_recursion(qubits, exp, exp_options)
             else:
-                self.select_exp(exp, qubits, couplers, exp_options)
+                self.select_exp(exp, qubits, exp_options)
             self.set_calibration_for_rt_sweep(exp)
             exp.set_signal_map(self.signal_map)
             self.experiment = exp
 
-    def set_calibration_for_rt_sweep(self, exp):
+    def set_calibration_for_rt_sweep(self, exp: lo.Experiment) -> None:
+        """Set laboneq calibration of parameters that are to be swept in real-
+        time."""
         if self.processed_sweeps:
             calib = lo.Calibration()
             for ch in (
@@ -702,30 +718,38 @@ def set_calibration_for_rt_sweep(self, exp):
     def set_instrument_nodes_for_nt_sweep(
         self, exp: lo.Experiment, sweeper: Sweeper
     ) -> None:
+        """In some cases there is no straightforward way to sweep a parameter.
+
+        In these cases we achieve sweeping by directly manipulating the
+        instrument nodes
+        """
         for ch, param, sweep_param in self.processed_sweeps.channel_sweeps_for_sweeper(
             sweeper
         ):
-            # for the remaining cases we can only achieve our goal by manipulating the instrument node directly
-            channel_node_path = None
-            logical_signal = self.signal_map[ch]
-            for instrument in self.device_setup.instruments:
-                for conn in instrument.connections:
-                    if conn.remote_path == logical_signal.path:
-                        channel_node_path = f"{instrument.address}/{conn.local_port}"
-                        break
-            if channel_node_path is None:
-                raise RuntimeError(
-                    f"Could not find instrument node corresponding to {param.name} of channel {ch}"
-                )
+            channel_node_path = self.get_channel_node_path(ch)
             if param is Parameter.bias:
                 offset_node_path = f"{channel_node_path}/offset"
                 exp.set_node(path=offset_node_path, value=sweep_param)
+
+            # This is supposed to happen only for measurement, but we do not validate it here.
             if param is Parameter.amplitude:
                 a, b = re.match(r"(.*)/(\d)/.*", channel_node_path).groups()
                 gain_node_path = f"{a}/{b}/oscs/{b}/gain"
                 exp.set_node(path=gain_node_path, value=sweep_param)
 
-    def select_exp(self, exp, qubits, couplers, exp_options):
+    def get_channel_node_path(self, channel_name: str) -> str:
+        """Return the path of the instrument node corresponding to the given
+        channel."""
+        logical_signal = self.signal_map[channel_name]
+        for instrument in self.device_setup.instruments:
+            for conn in instrument.connections:
+                if conn.remote_path == logical_signal.path:
+                    return f"{instrument.address}/{conn.local_port}"
+        raise RuntimeError(
+            f"Could not find instrument node corresponding to channel {channel_name}"
+        )
+
+    def select_exp(self, exp, qubits, exp_options):
         """Build Zurich Experiment selecting the relevant sections."""
         weights = {}
         previous_section = None
@@ -734,7 +758,7 @@ def select_exp(self, exp, qubits, couplers, exp_options):
             with exp.section(uid=section_uid, play_after=previous_section):
                 for ch, pulses in seq.control_sequence.items():
                     time = 0
-                    for j, pulse in enumerate(pulses):
+                    for pulse in pulses:
                         if pulse.delay_sweeper:
                             exp.delay(signal=ch, time=pulse.delay_sweeper)
                         exp.delay(
@@ -744,9 +768,6 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                         time = round(pulse.pulse.duration * NANO_TO_SECONDS, 9) + round(
                             pulse.pulse.start * NANO_TO_SECONDS, 9
                         )
-                        pulse.zhpulse.uid += (
-                            f"{i}_{j}"  # FIXME: is this actually needed?
-                        )
                         if pulse.zhsweepers:
                             self.play_sweep(exp, ch, pulse)
                         else:
@@ -770,7 +791,6 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                 for ch, pulse in seq.measurements:
                     qubit = qubits[pulse.pulse.qubit]
                     q = qubit.name
-                    pulse.zhpulse.uid += str(i)
 
                     exp.delay(
                         signal=acquire_channel_name(qubit),
@@ -783,7 +803,6 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                         == lo.AcquisitionType.DISCRIMINATION
                     ):
                         weight = lo.pulse_library.sampled_pulse_complex(
-                            uid="weight" + str(q),
                             samples=qubit.kernel * np.exp(1j * qubit.iq_angle),
                         )
 
@@ -803,13 +822,10 @@ def select_exp(self, exp, qubits, couplers, exp_options):
                                         ]
                                     )
                                     * np.exp(1j * qubit.iq_angle),
-                                    uid="weights" + str(q),
                                 )
                                 weights[q] = weight
                             else:
-                                # TODO: Patch for multiple readouts: Remove different uids
                                 weight = lo.pulse_library.const(
-                                    uid="weight" + str(q),
                                     length=round(
                                         pulse.pulse.duration * NANO_TO_SECONDS, 9
                                     )
@@ -865,33 +881,35 @@ def play_sweep(exp, channel_name, pulse):
         exp.play(signal=channel_name, pulse=pulse.zhpulse, **play_parameters)
 
     @staticmethod
-    def rearrange_sweepers(sweepers: List[Sweeper]) -> Tuple[np.ndarray, List[Sweeper]]:
+    def rearrange_rt_sweepers(
+        sweepers: list[Sweeper],
+    ) -> tuple[np.ndarray, list[Sweeper]]:
         """Rearranges sweepers from qibocal based on device hardware
         limitations.
 
-        Frequency sweepers must be applied before (on the outer loop) bias or amplitude sweepers.
+        The only known limitation currently is that frequency sweepers must be applied before (on the outer loop) other
+        (e.g. amplitude) sweepers. Consequently, the only thing done here is to swap the frequency sweeper with the
+        first sweeper in the list.
 
         Args:
-            sweepers (list[Sweeper]): list of sweepers used in the experiment.
+            sweepers (list[Sweeper]): Sweepers to rearrange.
 
         Returns:
-            rearranging_axes (np.ndarray): array of shape (2,) and dtype=int containing
-                the indexes of the sweepers to be swapped. Defaults to np.array([0, 0])
-                if no swap is needed.
-            sweepers (list[Sweeper]): updated list of sweepers used in the experiment. If
-                sweepers must be swapped, the list is updated accordingly.
+            swapped_axis_pair (np.ndarray): array of shape (2,) containing the indices of the two swapped axes.
+            sweepers (list[Sweeper]): rearranged list of sweepers.
         """
         swapped_axis_pair = np.zeros(2, dtype=int)
+        sweepers_copy = sweepers.copy()
         freq_sweeper = next(
-            iter(s for s in sweepers if s.parameter is Parameter.frequency), None
+            iter(s for s in sweepers_copy if s.parameter is Parameter.frequency), None
         )
         if freq_sweeper:
-            freq_sweeper_idx = sweepers.index(freq_sweeper)
-            sweepers[freq_sweeper_idx] = sweepers[0]
-            sweepers[0] = freq_sweeper
+            freq_sweeper_idx = sweepers_copy.index(freq_sweeper)
+            sweepers_copy[freq_sweeper_idx] = sweepers_copy[0]
+            sweepers_copy[0] = freq_sweeper
             swapped_axis_pair = np.array([0, freq_sweeper_idx])
             log.warning("Sweepers were reordered")
-        return swapped_axis_pair, sweepers
+        return swapped_axis_pair, sweepers_copy
 
     def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
@@ -899,7 +917,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         self.signal_map = {}
         self.processed_sweeps = ProcessedSweeps(sweepers, qubits)
         self.nt_sweeps, self.rt_sweeps = classify_sweepers(sweepers)
-        swapped_axis_pair, self.rt_sweeps = self.rearrange_sweepers(self.rt_sweeps)
+        swapped_axis_pair, self.rt_sweeps = self.rearrange_rt_sweepers(self.rt_sweeps)
         swapped_axis_pair += len(self.nt_sweeps)
         # if using singleshot, the first axis contains shots,
         # i.e.: (nshots, sweeper_1, sweeper_2)
@@ -913,7 +931,6 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
             if sweeper.parameter in {Parameter.frequency, Parameter.amplitude}:
                 for pulse in sweeper.pulses:
                     if pulse.type is PulseType.READOUT:
-                        # FIXME: case of multiple pulses
                         self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
 
         self.experiment_flow(qubits, couplers, sequence, options)
@@ -939,7 +956,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
 
         return results
 
-    def sweep_recursion(self, qubits, couplers, exp, exp_options):
+    def sweep_recursion(self, qubits, exp, exp_options):
         """Sweepers recursion for multiple nested Real Time sweepers."""
 
         sweeper = self.rt_sweeps[0]
@@ -956,14 +973,13 @@ def sweep_recursion(self, qubits, couplers, exp, exp_options):
             reset_oscillator_phase=True,
         ):
             if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, couplers, exp, exp_options)
+                self.sweep_recursion(qubits, exp, exp_options)
             else:
-                self.select_exp(exp, qubits, couplers, exp_options)
+                self.select_exp(exp, qubits, exp_options)
 
     def sweep_recursion_nt(
         self,
-        qubits: Dict[str, Qubit],
-        couplers: Dict[str, Coupler],
+        qubits: dict[str, Qubit],
         options: ExecutionParameters,
         exp: lo.Experiment,
     ):
@@ -989,11 +1005,13 @@ def sweep_recursion_nt(
                 for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
             ],
         ):
+            # This has to be called exactly here, otherwise laboneq will not identify the sweepable node
             self.set_instrument_nodes_for_nt_sweep(exp, sweeper)
+
             if len(self.nt_sweeps) > 0:
-                self.sweep_recursion_nt(qubits, couplers, options, exp)
+                self.sweep_recursion_nt(qubits, options, exp)
             else:
-                self.define_exp(qubits, couplers, options, exp)
+                self.define_exp(qubits, options, exp)
 
     def split_batches(self, sequences):
         return batch_max_sequences(sequences, MAX_SEQUENCES)

From 1869147a75b38a0b1f3edb6848a82d80b5f34645 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 29 Feb 2024 09:55:52 +0400
Subject: [PATCH 34/45] update tests

---
 tests/dummy_qrc/zurich/platform.py |   1 -
 tests/test_instruments_zhinst.py   | 599 ++++++++++++++++-------------
 2 files changed, 325 insertions(+), 275 deletions(-)

diff --git a/tests/dummy_qrc/zurich/platform.py b/tests/dummy_qrc/zurich/platform.py
index e74c11672..3014bbc22 100644
--- a/tests/dummy_qrc/zurich/platform.py
+++ b/tests/dummy_qrc/zurich/platform.py
@@ -86,7 +86,6 @@ def create(path: pathlib.Path = FOLDER):
     controller = Zurich(
         "EL_ZURO",
         device_setup=device_setup,
-        use_emulation=False,
         time_of_flight=75,
         smearing=50,
     )
diff --git a/tests/test_instruments_zhinst.py b/tests/test_instruments_zhinst.py
index c94fe8135..5cb175580 100644
--- a/tests/test_instruments_zhinst.py
+++ b/tests/test_instruments_zhinst.py
@@ -1,11 +1,20 @@
 import math
-import re
+from collections import defaultdict
 
+import laboneq.dsl.experiment.pulse as laboneq_pulse
+import laboneq.simple as lo
 import numpy as np
 import pytest
 
 from qibolab import AcquisitionType, AveragingMode, ExecutionParameters, create_platform
-from qibolab.instruments.zhinst import ZhPulse, ZhSweeperLine, Zurich
+from qibolab.instruments.zhinst import (
+    ProcessedSweeps,
+    ZhPulse,
+    Zurich,
+    acquire_channel_name,
+    classify_sweepers,
+    measure_channel_name,
+)
 from qibolab.pulses import (
     IIR,
     SNZ,
@@ -15,31 +24,25 @@
     Gaussian,
     Pulse,
     PulseSequence,
+    PulseType,
     ReadoutPulse,
     Rectangular,
 )
-from qibolab.sweeper import Parameter, Sweeper, SweeperType
+from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import batch
 
 from .conftest import get_instrument
 
 
 @pytest.mark.parametrize(
-    "shape", ["Rectangular", "Gaussian", "GaussianSquare", "Drag", "SNZ", "IIR"]
-)
-def test_zhpulse(shape):
-    if shape == "Rectangular":
-        pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
-    if shape == "Gaussian":
-        pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch0", qubit=0)
-    if shape == "GaussianSquare":
-        pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch0", qubit=0)
-    if shape == "Drag":
-        pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Drag(5, 0.4), "ch0", qubit=0)
-    if shape == "SNZ":
-        pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, SNZ(10, 0.01), "ch0", qubit=0)
-    if shape == "IIR":
-        pulse = Pulse(
+    "pulse",
+    [
+        Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0),
+        Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch0", qubit=0),
+        Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch0", qubit=0),
+        Pulse(0, 40, 0.05, int(3e9), 0.0, Drag(5, 0.4), "ch0", qubit=0),
+        Pulse(0, 40, 0.05, int(3e9), 0.0, SNZ(10, 0.01), "ch0", qubit=0),
+        Pulse(
             0,
             40,
             0.05,
@@ -48,39 +51,198 @@ def test_zhpulse(shape):
             IIR([10, 1], [0.4, 1], target=Gaussian(5)),
             "ch0",
             qubit=0,
-        )
+        ),
+    ],
+)
+def test_zhpulse_pulse_conversion(pulse):
+    shape = pulse.shape
+    zhpulse = ZhPulse(pulse).zhpulse
+    assert isinstance(zhpulse, laboneq_pulse.Pulse)
+    if isinstance(shape, (SNZ, IIR)):
+        assert len(zhpulse.samples) == 80
+    else:
+        assert zhpulse.length == 40e-9
+
 
+def test_zhpulse_add_sweeper():
+    pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch", qubit=0)
     zhpulse = ZhPulse(pulse)
-    assert zhpulse.pulse.serial == pulse.serial
-    if shape == "SNZ" or shape == "IIR":
-        assert len(zhpulse.zhpulse.samples) == 80
-    else:
-        assert zhpulse.zhpulse.length == 40e-9
+    assert zhpulse.zhsweepers == []
+    assert zhpulse.delay_sweeper is None
 
+    zhpulse.add_sweeper(
+        Parameter.duration, lo.SweepParameter(values=np.array([1, 2, 3]))
+    )
+    assert len(zhpulse.zhsweepers) == 1
+    assert zhpulse.delay_sweeper is None
 
-@pytest.mark.parametrize("parameter", [Parameter.bias, Parameter.start])
-def test_select_sweeper(dummy_qrc, parameter):
-    swept_points = 5
+    zhpulse.add_sweeper(
+        Parameter.start, lo.SweepParameter(values=np.array([4, 5, 6, 7]))
+    )
+    assert len(zhpulse.zhsweepers) == 1
+    assert zhpulse.delay_sweeper is not None
+
+    zhpulse.add_sweeper(
+        Parameter.amplitude, lo.SweepParameter(values=np.array([3, 2, 1, 0]))
+    )
+    assert len(zhpulse.zhsweepers) == 2
+    assert zhpulse.delay_sweeper is not None
+
+
+def test_measure_channel_name(dummy_qrc):
     platform = create_platform("zurich")
-    qubits = {0: platform.qubits[0]}
-    sequence = PulseSequence()
-    ro_pulses = {}
-    qd_pulses = {}
-    for qubit in qubits.values():
-        q = qubit.name
-        qd_pulses[q] = platform.create_RX_pulse(q, start=0)
-        sequence.add(qd_pulses[q])
-        ro_pulses[q] = platform.create_qubit_readout_pulse(q, start=qd_pulses[q].finish)
-        sequence.add(ro_pulses[q])
+    qubits = platform.qubits.values()
+    meas_ch_names = {measure_channel_name(q) for q in qubits}
+    assert len(qubits) > 0
+    assert len(meas_ch_names) == len(qubits)
+
+
+def test_acquire_channel_name(dummy_qrc):
+    platform = create_platform("zurich")
+    qubits = platform.qubits.values()
+    acq_ch_names = {acquire_channel_name(q) for q in qubits}
+    assert len(qubits) > 0
+    assert len(acq_ch_names) == len(qubits)
+
+
+def test_classify_sweepers(dummy_qrc):
+    platform = create_platform("zurich")
+    qubit_id, qubit = 0, platform.qubits[0]
+    pulse_1 = Pulse(0, 40, 0.05, int(3e9), 0.0, Gaussian(5), "ch0", qubit=qubit_id)
+    pulse_2 = Pulse(
+        0,
+        40,
+        0.05,
+        int(3e9),
+        0.0,
+        Rectangular(),
+        "ch7",
+        PulseType.READOUT,
+        qubit=qubit_id,
+    )
+    amplitude_sweeper = Sweeper(Parameter.amplitude, np.array([1, 2, 3]), [pulse_1])
+    readout_amplitude_sweeper = Sweeper(
+        Parameter.amplitude, np.array([1, 2, 3, 4, 5]), [pulse_2]
+    )
+    freq_sweeper = Sweeper(Parameter.frequency, np.array([4, 5, 6, 7]), [pulse_1])
+    bias_sweeper = Sweeper(Parameter.bias, np.array([3, 2, 1]), qubits=[qubit])
+    nt_sweeps, rt_sweeps = classify_sweepers(
+        [amplitude_sweeper, readout_amplitude_sweeper, bias_sweeper, freq_sweeper]
+    )
 
-        parameter_range = np.random.randint(swept_points, size=swept_points)
-        if parameter is Parameter.start:
-            sweeper = Sweeper(parameter, parameter_range, pulses=[qd_pulses[q]])
-        if parameter is Parameter.bias:
-            sweeper = Sweeper(parameter, parameter_range, qubits=q)
+    assert amplitude_sweeper in rt_sweeps
+    assert freq_sweeper in rt_sweeps
+    assert bias_sweeper in nt_sweeps
+    assert readout_amplitude_sweeper in nt_sweeps
 
-        ZhSweeper = ZhSweeperLine(sweeper, qubit, sequence)
-        assert ZhSweeper.sweeper == sweeper
+
+def test_processed_sweeps_pulse_properties(dummy_qrc):
+    platform = create_platform("zurich")
+    qubit_id_1, qubit_1 = 0, platform.qubits[0]
+    qubit_id_2, qubit_2 = 3, platform.qubits[3]
+    pulse_1 = Pulse(
+        0, 40, 0.05, int(3e9), 0.0, Gaussian(5), qubit_1.drive.name, qubit=qubit_id_1
+    )
+    pulse_2 = Pulse(
+        0, 40, 0.05, int(3e9), 0.0, Gaussian(5), qubit_2.drive.name, qubit=qubit_id_2
+    )
+    sweeper_amplitude = Sweeper(
+        Parameter.amplitude, np.array([1, 2, 3]), [pulse_1, pulse_2]
+    )
+    sweeper_duration = Sweeper(Parameter.duration, np.array([1, 2, 3, 4]), [pulse_2])
+    processed_sweeps = ProcessedSweeps(
+        [sweeper_duration, sweeper_amplitude], qubits=platform.qubits
+    )
+
+    assert len(processed_sweeps.sweeps_for_pulse(pulse_1)) == 1
+    assert processed_sweeps.sweeps_for_pulse(pulse_1)[0][0] == Parameter.amplitude
+    assert isinstance(
+        processed_sweeps.sweeps_for_pulse(pulse_1)[0][1], lo.SweepParameter
+    )
+    assert len(processed_sweeps.sweeps_for_pulse(pulse_2)) == 2
+
+    assert len(processed_sweeps.sweeps_for_sweeper(sweeper_amplitude)) == 2
+    parallel_sweep_ids = {
+        s.uid for s in processed_sweeps.sweeps_for_sweeper(sweeper_amplitude)
+    }
+    assert len(parallel_sweep_ids) == 2
+    assert processed_sweeps.sweeps_for_pulse(pulse_1)[0][1].uid in parallel_sweep_ids
+    assert any(
+        s.uid in parallel_sweep_ids
+        for _, s in processed_sweeps.sweeps_for_pulse(pulse_2)
+    )
+
+    assert len(processed_sweeps.sweeps_for_sweeper(sweeper_duration)) == 1
+    pulse_2_sweep_ids = {s.uid for _, s in processed_sweeps.sweeps_for_pulse(pulse_2)}
+    assert len(pulse_2_sweep_ids) == 2
+    assert (
+        processed_sweeps.sweeps_for_sweeper(sweeper_duration)[0].uid
+        in pulse_2_sweep_ids
+    )
+
+    assert processed_sweeps.channels_with_sweeps() == set()
+
+
+def test_processed_sweeps_frequency(dummy_qrc):
+    platform = create_platform("zurich")
+    qubit_id, qubit = 1, platform.qubits[1]
+    pulse = Pulse(
+        0, 40, 0.05, int(3e9), 0.0, Gaussian(5), qubit.drive.name, qubit=qubit_id
+    )
+    freq_sweeper = Sweeper(Parameter.frequency, np.array([1, 2, 3]), [pulse])
+    processed_sweeps = ProcessedSweeps([freq_sweeper], platform.qubits)
+
+    # Frequency sweepers should result into channel property sweeps
+    assert len(processed_sweeps.sweeps_for_pulse(pulse)) == 0
+    assert processed_sweeps.channels_with_sweeps() == {qubit.drive.name}
+    assert len(processed_sweeps.sweeps_for_channel(qubit.drive.name)) == 1
+
+    with pytest.raises(ValueError):
+        flux_pulse = Pulse(
+            0,
+            40,
+            0.05,
+            int(3e9),
+            0.0,
+            Gaussian(5),
+            qubit.flux.name,
+            PulseType.FLUX,
+            qubit=qubit_id,
+        )
+        freq_sweeper = Sweeper(
+            Parameter.frequency, np.array([1, 3, 5, 7]), [flux_pulse]
+        )
+        ProcessedSweeps([freq_sweeper], platform.qubits)
+
+
+def test_processed_sweeps_readout_amplitude(dummy_qrc):
+    platform = create_platform("zurich")
+    qubit_id, qubit = 0, platform.qubits[0]
+    readout_ch = measure_channel_name(qubit)
+    pulse_readout = Pulse(
+        0,
+        40,
+        0.05,
+        int(3e9),
+        0.0,
+        Rectangular(),
+        readout_ch,
+        PulseType.READOUT,
+        qubit_id,
+    )
+    readout_amplitude_sweeper = Sweeper(
+        Parameter.amplitude, np.array([1, 2, 3, 4]), [pulse_readout]
+    )
+    processed_sweeps = ProcessedSweeps(
+        [readout_amplitude_sweeper], qubits=platform.qubits
+    )
+
+    # Readout amplitude should result into channel property (gain) sweep
+    assert len(processed_sweeps.sweeps_for_pulse(pulse_readout)) == 0
+    assert processed_sweeps.channels_with_sweeps() == {
+        readout_ch,
+    }
+    assert len(processed_sweeps.sweeps_for_channel(readout_ch)) == 1
 
 
 def test_zhinst_setup(dummy_qrc):
@@ -90,110 +252,84 @@ def test_zhinst_setup(dummy_qrc):
 
 
 def test_zhsequence(dummy_qrc):
-    qd_pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
-    ro_pulse = ReadoutPulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch1", qubit=0)
-    sequence = PulseSequence()
-    sequence.add(qd_pulse)
-    sequence.add(ro_pulse)
     IQM5q = create_platform("zurich")
     controller = IQM5q.instruments["EL_ZURO"]
 
-    controller.sequence_zh(sequence, IQM5q.qubits, IQM5q.couplers)
-    zhsequence = controller.sequence
-
-    with pytest.raises(AttributeError):
-        controller.sequence_zh("sequence", IQM5q.qubits, IQM5q.couplers)
-        zhsequence = controller.sequence
-
-    assert len(zhsequence) == 2
-    assert len(zhsequence["readout0"]) == 1
-
-
-def test_zhsequence_couplers(dummy_qrc):
-    qd_pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
-    ro_pulse = ReadoutPulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch1", qubit=0)
-    qc_pulse = CouplerFluxPulse(0, 40, 0.05, Rectangular(), "ch_c0", qubit=3)
+    drive_channel, readout_channel = IQM5q.qubits[0].drive.name, measure_channel_name(
+        IQM5q.qubits[0]
+    )
+    qd_pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), drive_channel, qubit=0)
+    ro_pulse = ReadoutPulse(
+        0, 40, 0.05, int(3e9), 0.0, Rectangular(), readout_channel, qubit=0
+    )
     sequence = PulseSequence()
     sequence.add(qd_pulse)
+    sequence.add(qd_pulse)
     sequence.add(ro_pulse)
-    sequence.add(qc_pulse)
-    IQM5q = create_platform("zurich")
-    controller = IQM5q.instruments["EL_ZURO"]
 
-    controller.sequence_zh(sequence, IQM5q.qubits, IQM5q.couplers)
-    zhsequence = controller.sequence
+    zhsequence = controller.sequence_zh(sequence, IQM5q.qubits)
 
-    with pytest.raises(AttributeError):
-        controller.sequence_zh("sequence", IQM5q.qubits, IQM5q.couplers)
-        zhsequence = controller.sequence
+    assert len(zhsequence) == 2
+    assert len(zhsequence[drive_channel]) == 2
+    assert len(zhsequence[readout_channel]) == 1
 
-    assert len(zhsequence) == 3
-    assert len(zhsequence["readout0"]) == 1
-    assert len(zhsequence["couplerflux3"]) == 1
+    with pytest.raises(AttributeError):
+        controller.sequence_zh("sequence", IQM5q.qubits)
 
 
-def test_zhsequence_couplers_sweeper(dummy_qrc):
-    ro_pulse = ReadoutPulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch1", qubit=0)
-    sequence = PulseSequence()
-    sequence.add(ro_pulse)
+def test_zhsequence_couplers(dummy_qrc):
     IQM5q = create_platform("zurich")
     controller = IQM5q.instruments["EL_ZURO"]
 
-    delta_bias_range = np.arange(-1, 1, 0.5)
-
-    sweeper = Sweeper(
-        Parameter.amplitude,
-        delta_bias_range,
-        pulses=[
-            CouplerFluxPulse(
-                start=0,
-                duration=sequence.duration + sequence.start,
-                amplitude=1,
-                shape="Rectangular",
-                qubit=IQM5q.couplers[0].name,
-            )
-        ],
-        type=SweeperType.ABSOLUTE,
+    drive_channel, readout_channel = IQM5q.qubits[0].drive.name, measure_channel_name(
+        IQM5q.qubits[0]
     )
+    couplerflux_channel = IQM5q.couplers[0].flux.name
+    qd_pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), drive_channel, qubit=0)
+    ro_pulse = ReadoutPulse(
+        0, 40, 0.05, int(3e9), 0.0, Rectangular(), readout_channel, qubit=0
+    )
+    qc_pulse = CouplerFluxPulse(
+        0, 40, 0.05, Rectangular(), couplerflux_channel, qubit=3
+    )
+    sequence = PulseSequence()
+    sequence.add(qd_pulse)
+    sequence.add(ro_pulse)
+    sequence.add(qc_pulse)
 
-    controller.sweepers = [sweeper]
-    controller.sequence_zh(sequence, IQM5q.qubits, IQM5q.couplers)
-    zhsequence = controller.sequence
-
-    with pytest.raises(AttributeError):
-        controller.sequence_zh("sequence", IQM5q.qubits, IQM5q.couplers)
-        zhsequence = controller.sequence
+    zhsequence = controller.sequence_zh(sequence, IQM5q.qubits)
 
-    assert len(zhsequence) == 2
-    assert len(zhsequence["readout0"]) == 1
-    assert len(zhsequence["couplerflux0"]) == 0  # is it correct?
+    assert len(zhsequence) == 3
+    assert len(zhsequence[couplerflux_channel]) == 1
 
 
 def test_zhsequence_multiple_ro(dummy_qrc):
+    platform = create_platform("zurich")
+    readout_channel = measure_channel_name(platform.qubits[0])
     sequence = PulseSequence()
     qd_pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
     sequence.add(qd_pulse)
-    ro_pulse = ReadoutPulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch1", qubit=0)
+    ro_pulse = ReadoutPulse(
+        0, 40, 0.05, int(3e9), 0.0, Rectangular(), readout_channel, qubit=0
+    )
     sequence.add(ro_pulse)
-    ro_pulse = ReadoutPulse(0, 5000, 0.05, int(3e9), 0.0, Rectangular(), "ch1", qubit=0)
+    ro_pulse = ReadoutPulse(
+        0, 5000, 0.05, int(3e9), 0.0, Rectangular(), readout_channel, qubit=0
+    )
     sequence.add(ro_pulse)
     platform = create_platform("zurich")
 
     controller = platform.instruments["EL_ZURO"]
-    controller.sequence_zh(sequence, platform.qubits, platform.couplers)
-    zhsequence = controller.sequence
-
-    with pytest.raises(AttributeError):
-        controller.sequence_zh("sequence", platform.qubits, platform.couplers)
-        zhsequence = controller.sequence
+    zhsequence = controller.sequence_zh(sequence, platform.qubits)
 
     assert len(zhsequence) == 2
-    assert len(zhsequence["readout0"]) == 2
+    assert len(zhsequence[readout_channel]) == 2
 
 
 def test_zhinst_register_readout_line(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
+    qubit = platform.qubits[0]
 
     options = ExecutionParameters(
         relaxation_time=300e-6,
@@ -201,12 +337,10 @@ def test_zhinst_register_readout_line(dummy_qrc):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.register_readout_line(
-        platform.qubits[0], intermediate_frequency=int(1e6), options=options
-    )
+    IQM5q.register_readout_line(qubit, intermediate_frequency=int(1e6), options=options)
 
-    assert "measure0" in IQM5q.signal_map
-    assert "acquire0" in IQM5q.signal_map
+    assert measure_channel_name(qubit) in IQM5q.signal_map
+    assert acquire_channel_name(qubit) in IQM5q.signal_map
     assert (
         "/logical_signal_groups/q0/measure_line" in IQM5q.calibration.calibration_items
     )
@@ -215,22 +349,24 @@ def test_zhinst_register_readout_line(dummy_qrc):
 def test_zhinst_register_drive_line(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
-    IQM5q.register_drive_line(platform.qubits[0], intermediate_frequency=int(1e6))
+    qubit = platform.qubits[0]
+    IQM5q.register_drive_line(qubit, intermediate_frequency=int(1e6))
 
-    assert "drive0" in IQM5q.signal_map
+    assert qubit.drive.name in IQM5q.signal_map
     assert "/logical_signal_groups/q0/drive_line" in IQM5q.calibration.calibration_items
 
 
 def test_zhinst_register_flux_line(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
-    IQM5q.register_flux_line(platform.qubits[0])
+    qubit = platform.qubits[0]
+    IQM5q.register_flux_line(qubit)
 
-    assert "flux0" in IQM5q.signal_map
+    assert qubit.flux.name in IQM5q.signal_map
     assert "/logical_signal_groups/q0/flux_line" in IQM5q.calibration.calibration_items
 
 
-def test_experiment_execute_pulse_sequence(dummy_qrc):
+def test_experiment_flow(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
@@ -263,12 +399,12 @@ def test_experiment_execute_pulse_sequence(dummy_qrc):
 
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "flux0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    assert qubits[0].flux.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
-def test_experiment_execute_pulse_sequence_coupler(dummy_qrc):
+def test_experiment_flow_coupler(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
@@ -315,89 +451,55 @@ def test_experiment_execute_pulse_sequence_coupler(dummy_qrc):
 
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "flux0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    assert qubits[0].flux.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
-def test_experiment_fast_reset_readout(dummy_qrc):
+def test_sweep_and_play_sim(dummy_qrc):
+    """Test end-to-end experiment run using ZI emulated connection."""
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
     sequence = PulseSequence()
-    qubits = {0: platform.qubits[0]}
-    couplers = {}
-    platform.qubits = qubits
-
-    ro_pulses = {}
-    fr_pulses = {}
-    for qubit in qubits:
-        fr_pulses[qubit] = platform.create_RX_pulse(qubit, start=0)
-        ro_pulses[qubit] = platform.create_qubit_readout_pulse(qubit, start=0)
-        sequence.add(ro_pulses[qubit])
-
-    options = ExecutionParameters(
-        relaxation_time=300e-6,
-        fast_reset=fr_pulses,
-        acquisition_type=AcquisitionType.INTEGRATION,
-        averaging_mode=AveragingMode.CYCLIC,
-    )
-
-    IQM5q.experiment_flow(qubits, couplers, sequence, options)
-
-    assert "drive0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
-
-
-@pytest.mark.parametrize("fast_reset", [True, False])
-def test_experiment_execute_pulse_sequence(dummy_qrc, fast_reset):
-    platform = create_platform("zurich")
-    IQM5q = platform.instruments["EL_ZURO"]
-
-    sequence = PulseSequence()
-    qubits = {0: platform.qubits[0]}
+    qubits = {0: platform.qubits[0], 2: platform.qubits[2]}
     platform.qubits = qubits
+    couplers = {}
 
     ro_pulses = {}
-    qd_pulses = {}
     qf_pulses = {}
-    fr_pulses = {}
-    for qubit in qubits:
-        if fast_reset:
-            fr_pulses[qubit] = platform.create_RX_pulse(qubit, start=0)
-        qd_pulses[qubit] = platform.create_RX_pulse(qubit, start=0)
-        sequence.add(qd_pulses[qubit])
-        ro_pulses[qubit] = platform.create_qubit_readout_pulse(
-            qubit, start=qd_pulses[qubit].finish
-        )
-        sequence.add(ro_pulses[qubit])
-        qf_pulses[qubit] = FluxPulse(
+    for qubit in qubits.values():
+        q = qubit.name
+        qf_pulses[q] = FluxPulse(
             start=0,
-            duration=ro_pulses[qubit].start,
+            duration=500,
             amplitude=1,
             shape=Rectangular(),
-            channel=platform.qubits[qubit].flux.name,
-            qubit=qubit,
+            channel=platform.qubits[q].flux.name,
+            qubit=q,
         )
-        sequence.add(qf_pulses[qubit])
-
-    if fast_reset:
-        fast_reset = fr_pulses
+        sequence.add(qf_pulses[q])
+        ro_pulses[q] = platform.create_qubit_readout_pulse(q, start=qf_pulses[q].finish)
+        sequence.add(ro_pulses[q])
 
     options = ExecutionParameters(
         relaxation_time=300e-6,
-        fast_reset=fast_reset,
         acquisition_type=AcquisitionType.INTEGRATION,
         averaging_mode=AveragingMode.CYCLIC,
+        nshots=12,
     )
 
-    IQM5q.experiment_flow(qubits, sequence, options)
+    IQM5q.session = lo.Session(IQM5q.device_setup)
+    IQM5q.session.connect(do_emulation=True)
+    res = IQM5q.play(qubits, couplers, sequence, options)
+    assert res is not None
+    assert all(qubit in res for qubit in qubits)
 
-    assert "drive0" in IQM5q.experiment.signals
-    assert "flux0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    sweep_1 = Sweeper(Parameter.start, np.array([1, 2, 3, 4]), list(qf_pulses.values()))
+    sweep_2 = Sweeper(Parameter.bias, np.array([1, 2, 3]), qubits=[qubits[0]])
+    res = IQM5q.sweep(qubits, couplers, sequence, options, sweep_1, sweep_2)
+    assert res is not None
+    assert all(qubit in res for qubit in qubits)
 
 
 @pytest.mark.parametrize("parameter1", [Parameter.start, Parameter.duration])
@@ -405,7 +507,6 @@ def test_experiment_sweep_single(dummy_qrc, parameter1):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0]}
     couplers = {}
 
@@ -436,13 +537,11 @@ def test_experiment_sweep_single(dummy_qrc, parameter1):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.sweepers = sweepers
-
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "drive0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    assert qubits[0].drive.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
 @pytest.mark.parametrize("parameter1", [Parameter.start, Parameter.duration])
@@ -450,7 +549,6 @@ def test_experiment_sweep_single_coupler(dummy_qrc, parameter1):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0], 2: platform.qubits[2]}
     couplers = {0: platform.couplers[0]}
 
@@ -494,14 +592,12 @@ def test_experiment_sweep_single_coupler(dummy_qrc, parameter1):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.sweepers = sweepers
-
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "couplerflux0" in IQM5q.experiment.signals
-    assert "drive0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    assert couplers[0].flux.name in IQM5q.experiment.signals
+    assert qubits[0].drive.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
 SweeperParameter = {
@@ -519,7 +615,6 @@ def test_experiment_sweep_2d_general(dummy_qrc, parameter1, parameter2):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0]}
     couplers = {}
 
@@ -566,21 +661,17 @@ def test_experiment_sweep_2d_general(dummy_qrc, parameter1, parameter2):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.sweepers = sweepers
-    rearranging_axes, sweepers = IQM5q.rearrange_sweepers(sweepers)
-    IQM5q.sweepers = sweepers  # to be changed
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "drive0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
+    assert qubits[0].drive.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
 def test_experiment_sweep_2d_specific(dummy_qrc):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0]}
     couplers = {}
 
@@ -621,15 +712,11 @@ def test_experiment_sweep_2d_specific(dummy_qrc):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.sweepers = sweepers
-    rearranging_axes, sweepers = IQM5q.rearrange_sweepers(sweepers)
-    IQM5q.sweepers = sweepers  # to be changed
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "drive0" in IQM5q.experiment.signals
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
-    assert rearranging_axes != [[], []]
+    assert qubits[0].drive.name in IQM5q.experiment.signals
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
 @pytest.mark.parametrize(
@@ -639,7 +726,6 @@ def test_experiment_sweep_punchouts(dummy_qrc, parameter):
     platform = create_platform("zurich")
     IQM5q = platform.instruments["EL_ZURO"]
 
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0]}
     couplers = {}
 
@@ -687,41 +773,10 @@ def test_experiment_sweep_punchouts(dummy_qrc, parameter):
         averaging_mode=AveragingMode.CYCLIC,
     )
 
-    IQM5q.sweepers = sweepers
-    rearranging_axes, sweepers = IQM5q.rearrange_sweepers(sweepers)
-    IQM5q.sweepers = sweepers  # to be changed
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
 
-    assert "measure0" in IQM5q.experiment.signals
-    assert "acquire0" in IQM5q.experiment.signals
-
-
-# TODO: Fix this
-def test_sim(dummy_qrc):
-    platform = create_platform("zurich")
-    IQM5q = platform.instruments["EL_ZURO"]
-    sequence = PulseSequence()
-    qubits = {0: platform.qubits[0]}
-    platform.qubits = qubits
-    ro_pulses = {}
-    qd_pulses = {}
-    qf_pulses = {}
-    for qubit in qubits:
-        qd_pulses[qubit] = platform.create_RX_pulse(qubit, start=0)
-        sequence.add(qd_pulses[qubit])
-        ro_pulses[qubit] = platform.create_qubit_readout_pulse(
-            qubit, start=qd_pulses[qubit].finish
-        )
-        sequence.add(ro_pulses[qubit])
-        qf_pulses[qubit] = FluxPulse(
-            start=0,
-            duration=500,
-            amplitude=1,
-            shape=Rectangular(),
-            channel=platform.qubits[qubit].flux.name,
-            qubit=qubit,
-        )
-        sequence.add(qf_pulses[qubit])
+    assert measure_channel_name(qubits[0]) in IQM5q.experiment.signals
+    assert acquire_channel_name(qubits[0]) in IQM5q.experiment.signals
 
 
 def test_batching(dummy_qrc):
@@ -756,7 +811,7 @@ def test_connections(instrument):
 
 
 @pytest.mark.qpu
-def test_experiment_execute_pulse_sequence(connected_platform, instrument):
+def test_experiment_execute_pulse_sequence_qpu(connected_platform, instrument):
     platform = connected_platform
     sequence = PulseSequence()
     qubits = {0: platform.qubits[0], "c0": platform.qubits["c0"]}
@@ -795,9 +850,8 @@ def test_experiment_execute_pulse_sequence(connected_platform, instrument):
 
 
 @pytest.mark.qpu
-def test_experiment_sweep_2d_specific(connected_platform, instrument):
+def test_experiment_sweep_2d_specific_qpu(connected_platform, instrument):
     platform = connected_platform
-    sequence = PulseSequence()
     qubits = {0: platform.qubits[0]}
 
     swept_points = 5
@@ -849,25 +903,22 @@ def test_experiment_sweep_2d_specific(connected_platform, instrument):
 
 def get_previous_subsequence_finish(instrument, name):
     """Look recursively for sub_section finish times."""
-    signal_name = re.sub("sequence_", "", name)
-    signal_name = re.sub(r"_\d+$", "", signal_name)
-    signal_name = re.sub(r"flux", "bias", signal_name)
-    finish = 0
-    for section in instrument.experiment.sections[0].children:
-        if section.uid == name:
-            for pulse in section.children:
-                if pulse.signal == signal_name:
-                    try:
-                        finish += pulse.time
-                    except AttributeError:
-                        # not a laboneq Delay class object, skipping
-                        pass
-                    try:
-                        finish += pulse.pulse.length
-                    except AttributeError:
-                        # not a laboneq PlayPulse class object, skipping
-                        pass
-    return finish
+    section = next(
+        iter(ch for ch in instrument.experiment.sections[0].children if ch.uid == name)
+    )
+    finish = defaultdict(int)
+    for pulse in section.children:
+        try:
+            finish[pulse.signal] += pulse.time
+        except AttributeError:
+            # not a laboneq Delay class object, skipping
+            pass
+        try:
+            finish[pulse.signal] += pulse.pulse.length
+        except AttributeError:
+            # not a laboneq PlayPulse class object, skipping
+            pass
+    return max(finish.values())
 
 
 def test_experiment_measurement_sequence(dummy_qrc):
@@ -902,11 +953,11 @@ def test_experiment_measurement_sequence(dummy_qrc):
     IQM5q.experiment_flow(qubits, couplers, sequence, options)
     measure_start = 0
     for section in IQM5q.experiment.sections[0].children:
-        if section.uid == "sequence_measure_0":
+        if section.uid == "measure_0":
             measure_start += get_previous_subsequence_finish(IQM5q, section.play_after)
             for pulse in section.children:
                 try:
-                    if pulse.signal == "measure0":
+                    if pulse.signal == measure_channel_name(qubits[0]):
                         measure_start += pulse.time
                 except AttributeError:
                     # not a laboneq delay class object, skipping

From 624dec0c69007d3ef034435d61ca03e7365f4f73 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Mon, 4 Mar 2024 15:11:06 +0400
Subject: [PATCH 35/45] restructure zi driver

---
 src/qibolab/instruments/zhinst/__init__.py    |   4 +
 .../{zhinst.py => zhinst/executor.py}         | 255 +-----------------
 src/qibolab/instruments/zhinst/pulse.py       | 107 ++++++++
 src/qibolab/instruments/zhinst/sweep.py       | 139 ++++++++++
 src/qibolab/instruments/zhinst/util.py        |  24 ++
 5 files changed, 285 insertions(+), 244 deletions(-)
 create mode 100644 src/qibolab/instruments/zhinst/__init__.py
 rename src/qibolab/instruments/{zhinst.py => zhinst/executor.py} (75%)
 create mode 100644 src/qibolab/instruments/zhinst/pulse.py
 create mode 100644 src/qibolab/instruments/zhinst/sweep.py
 create mode 100644 src/qibolab/instruments/zhinst/util.py

diff --git a/src/qibolab/instruments/zhinst/__init__.py b/src/qibolab/instruments/zhinst/__init__.py
new file mode 100644
index 000000000..6eff487e6
--- /dev/null
+++ b/src/qibolab/instruments/zhinst/__init__.py
@@ -0,0 +1,4 @@
+from .executor import Zurich
+from .pulse import ZhPulse
+from .sweep import ProcessedSweeps, classify_sweepers
+from .util import acquire_channel_name, measure_channel_name
diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst/executor.py
similarity index 75%
rename from src/qibolab/instruments/zhinst.py
rename to src/qibolab/instruments/zhinst/executor.py
index be4e825a0..f52fb412c 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -1,25 +1,19 @@
-"""Instrument for using the Zurich Instruments (Zhinst) devices."""
+"""Executing pulse sequences on a Zurich Instruments devices."""
 
 import re
 from collections import defaultdict
-from collections.abc import Iterable
-from copy import copy
 from dataclasses import dataclass, replace
 from typing import Optional
 
 import laboneq.simple as lo
 import numpy as np
-from laboneq.dsl.experiment.pulse_library import (
-    sampled_pulse_complex,
-    sampled_pulse_real,
-)
 from qibo.config import log
 
 from qibolab import AcquisitionType, AveragingMode, ExecutionParameters
 from qibolab.couplers import Coupler
 
 from qibolab.instruments.unrolling import batch_max_sequences
-from qibolab.pulses import Pulse, PulseSequence, PulseType
+from qibolab.pulses import PulseSequence, PulseType
 from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
@@ -27,8 +21,15 @@
 from .abstract import Controller
 from .port import Port
 
-SAMPLING_RATE = 2
-NANO_TO_SECONDS = 1e-9
+from .pulse import ZhPulse
+from .sweep import ProcessedSweeps, classify_sweepers
+from .util import (
+    NANO_TO_SECONDS,
+    SAMPLING_RATE,
+    acquire_channel_name,
+    measure_channel_name,
+)
+
 COMPILER_SETTINGS = {
     "SHFSG_MIN_PLAYWAVE_HINT": 32,
     "SHFSG_MIN_PLAYZERO_HINT": 32,
@@ -47,103 +48,6 @@
     AveragingMode.SINGLESHOT: lo.AveragingMode.SINGLE_SHOT,
 }
 
-SWEEPER_SET = {"amplitude", "frequency", "duration", "relative_phase"}
-SWEEPER_BIAS = {"bias"}
-SWEEPER_START = {"start"}
-
-
-def measure_channel_name(qubit: Qubit) -> str:
-    """Construct and return a name for qubit's measure channel.
-
-    FIXME: We cannot use channel name directly, because currently channels are named after wires, and due to multiplexed readout
-    multiple qubits have the same channel name for their readout. Should be fixed once channels are refactored.
-    """
-    return f"{qubit.readout.name}_{qubit.name}"
-
-
-def acquire_channel_name(qubit: Qubit) -> str:
-    """Construct and return a name for qubit's acquire channel.
-
-    FIXME: We cannot use acquire channel name, because qibolab does not have a concept of acquire channel. This function shall be removed
-    once all channel refactoring is done.
-    """
-    return f"acquire{qubit.name}"
-
-
-def select_pulse(pulse: Pulse):
-    """Return laboneq pulse object corresponding to the given qibolab pulse."""
-    if "IIR" not in str(pulse.shape):
-        if str(pulse.shape) == "Rectangular()":
-            can_compress = pulse.type is not PulseType.READOUT
-            return lo.pulse_library.const(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                can_compress=can_compress,
-            )
-        if "Gaussian" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            return lo.pulse_library.gaussian(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                sigma=2 / sigma,
-                zero_boundaries=False,
-            )
-
-        if "GaussianSquare" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            width = pulse.shape.width
-            can_compress = pulse.type is not PulseType.READOUT
-            return lo.pulse_library.gaussian_square(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
-                amplitude=pulse.amplitude,
-                can_compress=can_compress,
-                sigma=2 / sigma,
-                zero_boundaries=False,
-            )
-
-        if "Drag" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            beta = pulse.shape.beta
-            return lo.pulse_library.drag(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                sigma=2 / sigma,
-                beta=beta,
-                zero_boundaries=False,
-            )
-
-    if np.all(pulse.envelope_waveform_q(SAMPLING_RATE).data == 0):
-        return sampled_pulse_real(
-            samples=pulse.envelope_waveform_i(SAMPLING_RATE).data,
-            can_compress=True,
-        )
-    else:
-        return sampled_pulse_complex(
-            samples=pulse.envelope_waveform_i(SAMPLING_RATE).data
-            + (1j * pulse.envelope_waveform_q(SAMPLING_RATE).data),
-            can_compress=True,
-        )
-
-
-def classify_sweepers(
-    sweepers: Iterable[Sweeper],
-) -> tuple[list[Sweeper], list[Sweeper]]:
-    """Divide sweepers into two lists: 1. sweeps that can be done in the laboneq near-time sweep loop, 2. sweeps that
-    can be done in real-time (i.e. on hardware)"""
-    nt_sweepers, rt_sweepers = [], []
-    for sweeper in sweepers:
-        if (
-            sweeper.parameter is Parameter.amplitude
-            and sweeper.pulses[0].type is PulseType.READOUT
-        ):
-            nt_sweepers.append(sweeper)
-        elif sweeper.parameter.name in SWEEPER_BIAS:
-            nt_sweepers.append(sweeper)
-        else:
-            rt_sweepers.append(sweeper)
-    return nt_sweepers, rt_sweepers
-
 
 @dataclass
 class ZhPort(Port):
@@ -152,143 +56,6 @@ class ZhPort(Port):
     power_range: int = 0
 
 
-class ZhPulse:
-    """Wrapper data type that holds a qibolab pulse, the corresponding laboneq
-    pulse object, and any sweeps associated with this pulse."""
-
-    def __init__(self, pulse):
-        self.pulse: Pulse = pulse
-        """Qibolab pulse."""
-        self.zhpulse = select_pulse(pulse)
-        """Laboneq pulse."""
-        self.zhsweepers: list[tuple[Parameter, lo.SweepParameter]] = []
-        """Parameters to be swept, along with their laboneq sweep parameter
-        definitions."""
-        self.delay_sweeper: Optional[lo.SweepParameter] = None
-        """Laboneq sweep parameter if the delay of the pulse should be
-        swept."""
-
-    # pylint: disable=R0903
-    def add_sweeper(self, param: Parameter, sweeper: lo.SweepParameter):
-        """Add sweeper to list of sweepers associated with this pulse."""
-        if param.name in SWEEPER_SET:
-            self.zhsweepers.append((param, sweeper))
-        if param.name in SWEEPER_START:
-            if self.delay_sweeper:
-                raise ValueError(
-                    "Cannot have multiple delay sweepers for a single pulse"
-                )
-            self.delay_sweeper = sweeper
-
-
-class ProcessedSweeps:
-    """Data type that centralizes and allows extracting information about given
-    sweeps."""
-
-    def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
-        pulse_sweeps = []
-        channel_sweeps = []
-        parallel_sweeps = []
-        for sweeper in sweepers:
-            for pulse in sweeper.pulses or []:
-                if sweeper.parameter in (Parameter.duration, Parameter.start):
-                    sweep_param = lo.SweepParameter(
-                        values=sweeper.values * NANO_TO_SECONDS
-                    )
-                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
-                elif sweeper.parameter is Parameter.frequency:
-                    ptype, qubit = pulse.type, qubits[pulse.qubit]
-                    if ptype is PulseType.READOUT:
-                        ch = measure_channel_name(qubit)
-                        intermediate_frequency = (
-                            qubit.readout_frequency
-                            - qubit.readout.local_oscillator.frequency
-                        )
-                    elif ptype is PulseType.DRIVE:
-                        ch = qubit.drive.name
-                        intermediate_frequency = (
-                            qubit.drive_frequency
-                            - qubit.drive.local_oscillator.frequency
-                        )
-                    else:
-                        raise ValueError(
-                            f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
-                        )
-                    sweep_param = lo.SweepParameter(
-                        values=sweeper.values + intermediate_frequency
-                    )
-                    channel_sweeps.append((ch, sweeper.parameter, sweep_param))
-                elif (
-                    pulse.type is PulseType.READOUT
-                    and sweeper.parameter is Parameter.amplitude
-                ):
-                    sweep_param = lo.SweepParameter(
-                        values=sweeper.values / max(sweeper.values)
-                    )
-                    channel_sweeps.append(
-                        (
-                            measure_channel_name(qubits[pulse.qubit]),
-                            sweeper.parameter,
-                            sweep_param,
-                        )
-                    )
-                else:
-                    sweep_param = lo.SweepParameter(values=copy(sweeper.values))
-                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
-                parallel_sweeps.append((sweeper, sweep_param))
-
-            for qubit in sweeper.qubits or []:
-                if sweeper.parameter is not Parameter.bias:
-                    raise ValueError(
-                        f"Sweeping {sweeper.parameter.name} for {qubit} is not supported"
-                    )
-                sweep_param = lo.SweepParameter(
-                    values=sweeper.values + qubit.flux.offset
-                )
-                channel_sweeps.append((qubit.flux.name, sweeper.parameter, sweep_param))
-                parallel_sweeps.append((sweeper, sweep_param))
-
-            for coupler in sweeper.couplers or []:
-                if sweeper.parameter is not Parameter.bias:
-                    raise ValueError(
-                        f"Sweeping {sweeper.parameter.name} for {coupler} is not supported"
-                    )
-                sweep_param = lo.SweepParameter(
-                    values=sweeper.values + coupler.flux.offset
-                )
-                channel_sweeps.append(
-                    (coupler.flux.name, sweeper.parameter, sweep_param)
-                )
-                parallel_sweeps.append((sweeper, sweep_param))
-
-        self._pulse_sweeps = pulse_sweeps
-        self._channel_sweeps = channel_sweeps
-        self._parallel_sweeps = parallel_sweeps
-
-    def sweeps_for_pulse(
-        self, pulse: Pulse
-    ) -> list[tuple[Parameter, lo.SweepParameter]]:
-        return [item[1:] for item in self._pulse_sweeps if item[0] == pulse]
-
-    def sweeps_for_channel(self, ch: str) -> list[tuple[Parameter, lo.SweepParameter]]:
-        return [item[1:] for item in self._channel_sweeps if item[0] == ch]
-
-    def sweeps_for_sweeper(self, sweeper: Sweeper) -> list[lo.SweepParameter]:
-        return [item[1] for item in self._parallel_sweeps if item[0] == sweeper]
-
-    def channel_sweeps_for_sweeper(
-        self, sweeper: Sweeper
-    ) -> list[tuple[str, Parameter, lo.SweepParameter]]:
-        return [
-            item
-            for item in self._channel_sweeps
-            if item[2] in self.sweeps_for_sweeper(sweeper)
-        ]
-
-    def channels_with_sweeps(self) -> set[str]:
-        return {ch for ch, _, _ in self._channel_sweeps}
-
-
 @dataclass
 class SubSequence:
     """A subsequence is a slice (in time) of a sequence that contains at most
diff --git a/src/qibolab/instruments/zhinst/pulse.py b/src/qibolab/instruments/zhinst/pulse.py
new file mode 100644
index 000000000..14c0b18e7
--- /dev/null
+++ b/src/qibolab/instruments/zhinst/pulse.py
@@ -0,0 +1,107 @@
+"""Wrapper for qibolab and laboneq pulses and sweeps."""
+
+from typing import Optional
+
+import laboneq.simple as lo
+import numpy as np
+from laboneq.dsl.experiment.pulse_library import (
+    sampled_pulse_complex,
+    sampled_pulse_real,
+)
+
+from qibolab.pulses import Pulse, PulseType
+from qibolab.sweeper import Parameter
+
+from .util import NANO_TO_SECONDS, SAMPLING_RATE
+
+
+def select_pulse(pulse: Pulse):
+    """Return laboneq pulse object corresponding to the given qibolab pulse."""
+    if "IIR" not in str(pulse.shape):
+        if str(pulse.shape) == "Rectangular()":
+            can_compress = pulse.type is not PulseType.READOUT
+            return lo.pulse_library.const(
+                length=round(pulse.duration * NANO_TO_SECONDS, 9),
+                amplitude=pulse.amplitude,
+                can_compress=can_compress,
+            )
+        if "Gaussian" in str(pulse.shape):
+            sigma = pulse.shape.rel_sigma
+            return lo.pulse_library.gaussian(
+                length=round(pulse.duration * NANO_TO_SECONDS, 9),
+                amplitude=pulse.amplitude,
+                sigma=2 / sigma,
+                zero_boundaries=False,
+            )
+
+        if "GaussianSquare" in str(pulse.shape):
+            sigma = pulse.shape.rel_sigma
+            width = pulse.shape.width
+            can_compress = pulse.type is not PulseType.READOUT
+            return lo.pulse_library.gaussian_square(
+                length=round(pulse.duration * NANO_TO_SECONDS, 9),
+                width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
+                amplitude=pulse.amplitude,
+                can_compress=can_compress,
+                sigma=2 / sigma,
+                zero_boundaries=False,
+            )
+
+        if "Drag" in str(pulse.shape):
+            sigma = pulse.shape.rel_sigma
+            beta = pulse.shape.beta
+            return lo.pulse_library.drag(
+                length=round(pulse.duration * NANO_TO_SECONDS, 9),
+                amplitude=pulse.amplitude,
+                sigma=2 / sigma,
+                beta=beta,
+                zero_boundaries=False,
+            )
+
+    if np.all(pulse.envelope_waveform_q(SAMPLING_RATE).data == 0):
+        return sampled_pulse_real(
+            samples=pulse.envelope_waveform_i(SAMPLING_RATE).data,
+            can_compress=True,
+        )
+    else:
+        return sampled_pulse_complex(
+            samples=pulse.envelope_waveform_i(SAMPLING_RATE).data
+            + (1j * pulse.envelope_waveform_q(SAMPLING_RATE).data),
+            can_compress=True,
+        )
+
+
+class ZhPulse:
+    """Wrapper data type that holds a qibolab pulse, the corresponding laboneq
+    pulse object, and any sweeps associated with this pulse."""
+
+    def __init__(self, pulse):
+        self.pulse: Pulse = pulse
+        """Qibolab pulse."""
+        self.zhpulse = select_pulse(pulse)
+        """Laboneq pulse."""
+        self.zhsweepers: list[tuple[Parameter, lo.SweepParameter]] = []
+        """Parameters to be swept, along with their laboneq sweep parameter
+        definitions."""
+        self.delay_sweeper: Optional[lo.SweepParameter] = None
+        """Laboneq sweep parameter if the delay of the pulse should be
+        swept."""
+
+    # pylint: disable=R0903
+    def add_sweeper(self, param: Parameter, sweeper: lo.SweepParameter):
+        """Add sweeper to list of sweepers associated with this pulse."""
+        if param in {
+            Parameter.amplitude,
+            Parameter.frequency,
+            Parameter.duration,
+            Parameter.relative_phase,
+        }:
+            self.zhsweepers.append((param, sweeper))
+        elif param is Parameter.start:
+            if self.delay_sweeper:
+                raise ValueError(
+                    "Cannot have multiple delay sweepers for a single pulse"
+                )
+            self.delay_sweeper = sweeper
+        else:
+            raise ValueError(f"Sweeping {param} is not supported")
diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
new file mode 100644
index 000000000..918c98d0c
--- /dev/null
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -0,0 +1,139 @@
+"""Pre-execution processing of sweeps."""
+
+from collections.abc import Iterable
+from copy import copy
+
+import laboneq.simple as lo
+
+from qibolab.pulses import Pulse, PulseType
+from qibolab.qubits import Qubit
+from qibolab.sweeper import Parameter, Sweeper
+
+from .util import NANO_TO_SECONDS, measure_channel_name
+
+
+def classify_sweepers(
+    sweepers: Iterable[Sweeper],
+) -> tuple[list[Sweeper], list[Sweeper]]:
+    """Divide sweepers into two lists: 1. sweeps that can be done in the laboneq near-time sweep loop, 2. sweeps that
+    can be done in real-time (i.e. on hardware)"""
+    nt_sweepers, rt_sweepers = [], []
+    for sweeper in sweepers:
+        if (
+            sweeper.parameter is Parameter.amplitude
+            and sweeper.pulses[0].type is PulseType.READOUT
+        ):
+            nt_sweepers.append(sweeper)
+        elif sweeper.parameter is Parameter.bias:
+            nt_sweepers.append(sweeper)
+        else:
+            rt_sweepers.append(sweeper)
+    return nt_sweepers, rt_sweepers
+
+
+class ProcessedSweeps:
+    """Data type that centralizes and allows extracting information about given
+    sweeps."""
+
+    def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
+        pulse_sweeps = []
+        channel_sweeps = []
+        parallel_sweeps = []
+        for sweeper in sweepers:
+            for pulse in sweeper.pulses or []:
+                if sweeper.parameter in (Parameter.duration, Parameter.start):
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values * NANO_TO_SECONDS
+                    )
+                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
+                elif sweeper.parameter is Parameter.frequency:
+                    ptype, qubit = pulse.type, qubits[pulse.qubit]
+                    if ptype is PulseType.READOUT:
+                        ch = measure_channel_name(qubit)
+                        intermediate_frequency = (
+                            qubit.readout_frequency
+                            - qubit.readout.local_oscillator.frequency
+                        )
+                    elif ptype is PulseType.DRIVE:
+                        ch = qubit.drive.name
+                        intermediate_frequency = (
+                            qubit.drive_frequency
+                            - qubit.drive.local_oscillator.frequency
+                        )
+                    else:
+                        raise ValueError(
+                            f"Cannot sweep frequency of pulse of type {ptype}, because it does not have associated frequency"
+                        )
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values + intermediate_frequency
+                    )
+                    channel_sweeps.append((ch, sweeper.parameter, sweep_param))
+                elif (
+                    pulse.type is PulseType.READOUT
+                    and sweeper.parameter is Parameter.amplitude
+                ):
+                    sweep_param = lo.SweepParameter(
+                        values=sweeper.values / max(sweeper.values)
+                    )
+                    channel_sweeps.append(
+                        (
+                            measure_channel_name(qubits[pulse.qubit]),
+                            sweeper.parameter,
+                            sweep_param,
+                        )
+                    )
+                else:
+                    sweep_param = lo.SweepParameter(values=copy(sweeper.values))
+                    pulse_sweeps.append((pulse, sweeper.parameter, sweep_param))
+                parallel_sweeps.append((sweeper, sweep_param))
+
+            for qubit in sweeper.qubits or []:
+                if sweeper.parameter is not Parameter.bias:
+                    raise ValueError(
+                        f"Sweeping {sweeper.parameter.name} for {qubit} is not supported"
+                    )
+                sweep_param = lo.SweepParameter(
+                    values=sweeper.values + qubit.flux.offset
+                )
+                channel_sweeps.append((qubit.flux.name, sweeper.parameter, sweep_param))
+                parallel_sweeps.append((sweeper, sweep_param))
+
+            for coupler in sweeper.couplers or []:
+                if sweeper.parameter is not Parameter.bias:
+                    raise ValueError(
+                        f"Sweeping {sweeper.parameter.name} for {coupler} is not supported"
+                    )
+                sweep_param = lo.SweepParameter(
+                    values=sweeper.values + coupler.flux.offset
+                )
+                channel_sweeps.append(
+                    (coupler.flux.name, sweeper.parameter, sweep_param)
+                )
+                parallel_sweeps.append((sweeper, sweep_param))
+
+        self._pulse_sweeps = pulse_sweeps
+        self._channel_sweeps = channel_sweeps
+        self._parallel_sweeps = parallel_sweeps
+
+    def sweeps_for_pulse(
+        self, pulse: Pulse
+    ) -> list[tuple[Parameter, lo.SweepParameter]]:
+        return [item[1:] for item in self._pulse_sweeps if item[0] == pulse]
+
+    def sweeps_for_channel(self, ch: str) -> list[tuple[Parameter, lo.SweepParameter]]:
+        return [item[1:] for item in self._channel_sweeps if item[0] == ch]
+
+    def sweeps_for_sweeper(self, sweeper: Sweeper) -> list[lo.SweepParameter]:
+        return [item[1] for item in self._parallel_sweeps if item[0] == sweeper]
+
+    def channel_sweeps_for_sweeper(
+        self, sweeper: Sweeper
+    ) -> list[tuple[str, Parameter, lo.SweepParameter]]:
+        return [
+            item
+            for item in self._channel_sweeps
+            if item[2] in self.sweeps_for_sweeper(sweeper)
+        ]
+
+    def channels_with_sweeps(self) -> set[str]:
+        return {ch for ch, _, _ in self._channel_sweeps}
diff --git a/src/qibolab/instruments/zhinst/util.py b/src/qibolab/instruments/zhinst/util.py
new file mode 100644
index 000000000..8ca884306
--- /dev/null
+++ b/src/qibolab/instruments/zhinst/util.py
@@ -0,0 +1,24 @@
+"""Utility methods."""
+
+from qibolab.qubits import Qubit
+
+SAMPLING_RATE = 2
+NANO_TO_SECONDS = 1e-9
+
+
+def measure_channel_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's measure channel.
+
+    FIXME: We cannot use channel name directly, because currently channels are named after wires, and due to multiplexed readout
+    multiple qubits have the same channel name for their readout. Should be fixed once channels are refactored.
+    """
+    return f"{qubit.readout.name}_{qubit.name}"
+
+
+def acquire_channel_name(qubit: Qubit) -> str:
+    """Construct and return a name for qubit's acquire channel.
+
+    FIXME: We cannot use acquire channel name, because qibolab does not have a concept of acquire channel. This function shall be removed
+    once all channel refactoring is done.
+    """
+    return f"acquire{qubit.name}"

From d0fb3e4546629dec9e46019ba06d1598f143b571 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 5 Mar 2024 10:01:51 +0400
Subject: [PATCH 36/45] Update src/qibolab/instruments/zhinst/sweep.py

Co-authored-by: Alessandro Candido <candido.ale@gmail.com>
---
 src/qibolab/instruments/zhinst/sweep.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
index 918c98d0c..8ad27b9eb 100644
--- a/src/qibolab/instruments/zhinst/sweep.py
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -20,12 +20,11 @@ def classify_sweepers(
     nt_sweepers, rt_sweepers = [], []
     for sweeper in sweepers:
         if (
-            sweeper.parameter is Parameter.amplitude
-            and sweeper.pulses[0].type is PulseType.READOUT
+            sweeper.parameter is Parameter.bias or
+            (sweeper.parameter is Parameter.amplitude
+            and sweeper.pulses[0].type is PulseType.READOUT)
         ):
             nt_sweepers.append(sweeper)
-        elif sweeper.parameter is Parameter.bias:
-            nt_sweepers.append(sweeper)
         else:
             rt_sweepers.append(sweeper)
     return nt_sweepers, rt_sweepers

From 3d20d8f65473419bbcf1a305d213b5642bcf31cb Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 5 Mar 2024 06:02:04 +0000
Subject: [PATCH 37/45] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/qibolab/instruments/zhinst/sweep.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
index 8ad27b9eb..ecc2dae41 100644
--- a/src/qibolab/instruments/zhinst/sweep.py
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -19,10 +19,9 @@ def classify_sweepers(
     can be done in real-time (i.e. on hardware)"""
     nt_sweepers, rt_sweepers = [], []
     for sweeper in sweepers:
-        if (
-            sweeper.parameter is Parameter.bias or
-            (sweeper.parameter is Parameter.amplitude
-            and sweeper.pulses[0].type is PulseType.READOUT)
+        if sweeper.parameter is Parameter.bias or (
+            sweeper.parameter is Parameter.amplitude
+            and sweeper.pulses[0].type is PulseType.READOUT
         ):
             nt_sweepers.append(sweeper)
         else:

From 5c3e1e7724b580a1b6a6deb3bd0265e0f8811612 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 5 Mar 2024 12:26:37 +0400
Subject: [PATCH 38/45] Expand the documentation of ProcessedSweeps

---
 src/qibolab/instruments/zhinst/sweep.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
index ecc2dae41..8c99d6b90 100644
--- a/src/qibolab/instruments/zhinst/sweep.py
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -31,7 +31,29 @@ def classify_sweepers(
 
 class ProcessedSweeps:
     """Data type that centralizes and allows extracting information about given
-    sweeps."""
+    sweeps.
+
+    In laboneq, sweeps are represented with the help of SweepParameter
+    instances. When adding pulses to a laboneq experiment, some
+    properties can be set to be an instance of SweepParameter instead of
+    a fixed numeric value. In case of channel property sweeps, either
+    the relevant calibration property or the instrument node directly
+    can be set ot a SweepParameter instance. Parts of the laboneq
+    experiment that define the sweep loops refer to SweepParameter
+    instances as well. These should be linkable to instances that are
+    either set to a pulse property, a channel calibration or instrument
+    node. To achieve this, we use the exact same SweepParameter instance
+    in both places. This class takes care of creating these
+    SweepParameter instances and giving access to them in a consistent
+    way (i.e. whenever they need to be the same instance they will be
+    the same instance). When constructing sweep loops you may ask from
+    this class to provide all the SweepParameter instances related to a
+    given qibolab Sweeper (parallel sweeps). Later, when adding pulses
+    or setting channel properties, you may ask from this class to
+    provide all SweepParameter instances related to a given pulse or
+    channel, and you will get parameters that are linkable to the ones
+    in the sweep loop definition
+    """
 
     def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
         pulse_sweeps = []

From 59829e894792257b75d52937fb3ac26393c178e3 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Thu, 7 Mar 2024 11:38:12 +0400
Subject: [PATCH 39/45] reuse sweep for play

---
 src/qibolab/instruments/zhinst/executor.py | 63 +++++++---------------
 tests/test_instruments_zhinst.py           |  7 +++
 2 files changed, 27 insertions(+), 43 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/executor.py b/src/qibolab/instruments/zhinst/executor.py
index f52fb412c..982a40ef2 100644
--- a/src/qibolab/instruments/zhinst/executor.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -366,29 +366,7 @@ def experiment_flow(
     # pylint: disable=W0221
     def play(self, qubits, couplers, sequence, options):
         """Play pulse sequence."""
-        self.signal_map = {}
-
-        self.frequency_from_pulses(qubits, sequence)
-
-        self.experiment_flow(qubits, couplers, sequence, options)
-
-        self.run_exp()
-
-        # Get the results back
-        results = {}
-        for qubit in qubits.values():
-            q = qubit.name  # pylint: disable=C0103
-            for i, ropulse in enumerate(self.sequence[measure_channel_name(qubit)]):
-                data = np.array(self.results.get_data(f"sequence{q}_{i}"))
-                if options.acquisition_type is AcquisitionType.DISCRIMINATION:
-                    data = (
-                        np.ones(data.shape) - data.real
-                    )  # Probability inversion patch
-                serial = ropulse.pulse.serial
-                qubit = ropulse.pulse.qubit
-                results[serial] = results[qubit] = options.results_type(data)
-
-        return results
+        return self.sweep(qubits, couplers, sequence, options)
 
     def sequence_zh(
         self, sequence: PulseSequence, qubits: dict[str, Qubit]
@@ -650,33 +628,31 @@ def play_sweep(exp, channel_name, pulse):
     @staticmethod
     def rearrange_rt_sweepers(
         sweepers: list[Sweeper],
-    ) -> tuple[np.ndarray, list[Sweeper]]:
-        """Rearranges sweepers from qibocal based on device hardware
-        limitations.
+    ) -> tuple[Optional[tuple[int, int]], list[Sweeper]]:
+        """Rearranges list of real-time sweepers based on hardware limitations.
 
         The only known limitation currently is that frequency sweepers must be applied before (on the outer loop) other
         (e.g. amplitude) sweepers. Consequently, the only thing done here is to swap the frequency sweeper with the
         first sweeper in the list.
 
         Args:
-            sweepers (list[Sweeper]): Sweepers to rearrange.
+            sweepers: Sweepers to rearrange.
 
         Returns:
-            swapped_axis_pair (np.ndarray): array of shape (2,) containing the indices of the two swapped axes.
-            sweepers (list[Sweeper]): rearranged list of sweepers.
+            swapped_axis_pair: tuple containing indices of the two swapped axes, or None if nothing to rearrange.
+            sweepers: rearranged (or original, if nothing to rearrange) list of sweepers.
         """
-        swapped_axis_pair = np.zeros(2, dtype=int)
-        sweepers_copy = sweepers.copy()
         freq_sweeper = next(
-            iter(s for s in sweepers_copy if s.parameter is Parameter.frequency), None
+            iter(s for s in sweepers if s.parameter is Parameter.frequency), None
         )
         if freq_sweeper:
+            sweepers_copy = sweepers.copy()
             freq_sweeper_idx = sweepers_copy.index(freq_sweeper)
             sweepers_copy[freq_sweeper_idx] = sweepers_copy[0]
             sweepers_copy[0] = freq_sweeper
-            swapped_axis_pair = np.array([0, freq_sweeper_idx])
             log.warning("Sweepers were reordered")
-        return swapped_axis_pair, sweepers_copy
+            return (0, freq_sweeper_idx), sweepers_copy
+        return None, sweepers
 
     def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
@@ -685,12 +661,13 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         self.processed_sweeps = ProcessedSweeps(sweepers, qubits)
         self.nt_sweeps, self.rt_sweeps = classify_sweepers(sweepers)
         swapped_axis_pair, self.rt_sweeps = self.rearrange_rt_sweepers(self.rt_sweeps)
-        swapped_axis_pair += len(self.nt_sweeps)
-        # if using singleshot, the first axis contains shots,
-        # i.e.: (nshots, sweeper_1, sweeper_2)
-        # if using integration: (sweeper_1, sweeper_2)
-        if options.averaging_mode is AveragingMode.SINGLESHOT:
-            swapped_axis_pair += 1
+        if swapped_axis_pair:
+            # 1. axes corresponding to NT sweeps appear before axes corresponding to RT sweeps
+            # 2. in singleshot mode, the first axis contains shots, i.e.: (nshots, sweeper_1, sweeper_2)
+            axis_offset = len(self.nt_sweeps) + int(
+                options.averaging_mode is AveragingMode.SINGLESHOT
+            )
+            swapped_axis_pair = tuple(ax + axis_offset for ax in swapped_axis_pair)
 
         self.frequency_from_pulses(qubits, sequence)
 
@@ -708,10 +685,10 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         for qubit in qubits.values():
             q = qubit.name  # pylint: disable=C0103
             for i, ropulse in enumerate(self.sequence[measure_channel_name(qubit)]):
-                exp_res = self.results.get_data(f"sequence{q}_{i}")
+                data = self.results.get_data(f"sequence{q}_{i}")
 
-                # Reorder dimensions
-                data = np.moveaxis(exp_res, swapped_axis_pair[0], swapped_axis_pair[1])
+                if swapped_axis_pair:
+                    data = np.moveaxis(data, swapped_axis_pair[0], swapped_axis_pair[1])
                 if options.acquisition_type is AcquisitionType.DISCRIMINATION:
                     data = (
                         np.ones(data.shape) - data.real
diff --git a/tests/test_instruments_zhinst.py b/tests/test_instruments_zhinst.py
index 5cb175580..bbea85ecb 100644
--- a/tests/test_instruments_zhinst.py
+++ b/tests/test_instruments_zhinst.py
@@ -489,12 +489,19 @@ def test_sweep_and_play_sim(dummy_qrc):
         nshots=12,
     )
 
+    # check play
     IQM5q.session = lo.Session(IQM5q.device_setup)
     IQM5q.session.connect(do_emulation=True)
     res = IQM5q.play(qubits, couplers, sequence, options)
     assert res is not None
     assert all(qubit in res for qubit in qubits)
 
+    # check sweep with empty list of sweeps
+    res = IQM5q.sweep(qubits, couplers, sequence, options)
+    assert res is not None
+    assert all(qubit in res for qubit in qubits)
+
+    # check sweep with sweeps
     sweep_1 = Sweeper(Parameter.start, np.array([1, 2, 3, 4]), list(qf_pulses.values()))
     sweep_2 = Sweeper(Parameter.bias, np.array([1, 2, 3]), qubits=[qubits[0]])
     res = IQM5q.sweep(qubits, couplers, sequence, options, sweep_1, sweep_2)

From 0a2ded8ec4a0791fb7412c64217ed6066ce128b2 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 12 Mar 2024 10:12:56 +0400
Subject: [PATCH 40/45] simplify the recursive functions

---
 src/qibolab/instruments/zhinst/executor.py | 153 +++++++++------------
 1 file changed, 68 insertions(+), 85 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/executor.py b/src/qibolab/instruments/zhinst/executor.py
index 982a40ef2..e8d3a1042 100644
--- a/src/qibolab/instruments/zhinst/executor.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -3,7 +3,7 @@
 import re
 from collections import defaultdict
 from dataclasses import dataclass, replace
-from typing import Optional
+from typing import Any, Optional
 
 import laboneq.simple as lo
 import numpy as np
@@ -400,18 +400,8 @@ def sequence_zh(
 
     def create_exp(self, qubits, options):
         """Zurich experiment initialization using their Experiment class."""
-
-        # Setting experiment signal lines
-        signals = [lo.ExperimentSignal(name) for name in self.signal_map.keys()]
-
-        exp = lo.Experiment(
-            uid="Sequence",
-            signals=signals,
-        )
-
         if self.acquisition_type:
             acquisition_type = self.acquisition_type
-            self.acquisition_type = None
         else:
             acquisition_type = ACQUISITION_TYPE[options.acquisition_type]
         averaging_mode = AVERAGING_MODE[options.averaging_mode]
@@ -419,28 +409,77 @@ def create_exp(self, qubits, options):
             options, acquisition_type=acquisition_type, averaging_mode=averaging_mode
         )
 
-        # Near Time recursion loop or directly to Real Time recursion loop
-        if self.nt_sweeps:
-            self.sweep_recursion_nt(qubits, exp_options, exp)
-        else:
-            self.define_exp(qubits, exp_options, exp)
+        signals = [lo.ExperimentSignal(name) for name in self.signal_map.keys()]
+        exp = lo.Experiment(
+            uid="Sequence",
+            signals=signals,
+        )
+
+        contexts = self._contexts(exp, exp_options)
+        self._populate_exp(qubits, exp, exp_options, contexts)
+        self.set_calibration_for_rt_sweep(exp)
+        exp.set_signal_map(self.signal_map)
+        self.experiment = exp
 
-    def define_exp(self, qubits, exp_options, exp):
-        """Real time definition."""
-        with exp.acquire_loop_rt(
+    def _contexts(
+        self, exp: lo.Experiment, exp_options: ExecutionParameters
+    ) -> list[tuple[Optional[Sweeper], Any]]:
+        """To construct a laboneq experiment, we need to first define a certain
+        sequence of nested contexts.
+
+        This method returns the corresponding sequence of context
+        managers.
+        """
+        sweep_contexts = []
+        for i, sweeper in enumerate(self.nt_sweeps):
+            ctx = exp.sweep(
+                uid=f"nt_sweep_{sweeper.parameter.name.lower()}_{i}",
+                parameter=[
+                    sweep_param
+                    for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
+                ],
+            )
+            sweep_contexts.append((sweeper, ctx))
+
+        shots_ctx = exp.acquire_loop_rt(
             uid="shots",
             count=exp_options.nshots,
             acquisition_type=exp_options.acquisition_type,
             averaging_mode=exp_options.averaging_mode,
-        ):
-            # Recursion loop for sweepers or just play a sequence
-            if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, exp, exp_options)
-            else:
-                self.select_exp(exp, qubits, exp_options)
-            self.set_calibration_for_rt_sweep(exp)
-            exp.set_signal_map(self.signal_map)
-            self.experiment = exp
+        )
+        sweep_contexts.append((None, shots_ctx))
+
+        for i, sweeper in enumerate(self.rt_sweeps):
+            ctx = exp.sweep(
+                uid=f"rt_sweep_{sweeper.parameter.name.lower()}_{i}",
+                parameter=[
+                    sweep_param
+                    for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
+                ],
+                reset_oscillator_phase=True,
+            )
+            sweep_contexts.append((sweeper, ctx))
+
+        return sweep_contexts
+
+    def _populate_exp(
+        self,
+        qubits: dict[str, Qubit],
+        exp: lo.Experiment,
+        exp_options: ExecutionParameters,
+        contexts,
+    ):
+        """Recursively activate the nested contexts, then define the main
+        experiment body inside the innermost context."""
+        if len(contexts) == 0:
+            self.select_exp(exp, qubits, exp_options)
+            return
+
+        sweeper, ctx = contexts[0]
+        with ctx:
+            if sweeper in self.nt_sweeps:
+                self.set_instrument_nodes_for_nt_sweep(exp, sweeper)
+            self._populate_exp(qubits, exp, exp_options, contexts[1:])
 
     def set_calibration_for_rt_sweep(self, exp: lo.Experiment) -> None:
         """Set laboneq calibration of parameters that are to be swept in real-
@@ -671,6 +710,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
 
         self.frequency_from_pulses(qubits, sequence)
 
+        self.acquisition_type = None
         for sweeper in sweepers:
             if sweeper.parameter in {Parameter.frequency, Parameter.amplitude}:
                 for pulse in sweeper.pulses:
@@ -700,62 +740,5 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
 
         return results
 
-    def sweep_recursion(self, qubits, exp, exp_options):
-        """Sweepers recursion for multiple nested Real Time sweepers."""
-
-        sweeper = self.rt_sweeps[0]
-
-        i = len(self.rt_sweeps) - 1
-        self.rt_sweeps.remove(sweeper)
-
-        with exp.sweep(
-            uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
-            parameter=[
-                sweep_param
-                for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
-            ],
-            reset_oscillator_phase=True,
-        ):
-            if len(self.rt_sweeps) > 0:
-                self.sweep_recursion(qubits, exp, exp_options)
-            else:
-                self.select_exp(exp, qubits, exp_options)
-
-    def sweep_recursion_nt(
-        self,
-        qubits: dict[str, Qubit],
-        options: ExecutionParameters,
-        exp: lo.Experiment,
-    ):
-        """Sweepers recursion for Near Time sweepers. Faster than regular
-        software sweepers as they are executed on the actual device by
-        (software ? or slower hardware ones)
-
-        You want to avoid them so for now they are implement for a
-        specific sweep.
-        """
-
-        log.info("nt Loop")
-
-        sweeper = self.nt_sweeps[0]
-
-        i = len(self.nt_sweeps) - 1
-        self.nt_sweeps.remove(sweeper)
-
-        with exp.sweep(
-            uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",
-            parameter=[
-                sweep_param
-                for sweep_param in self.processed_sweeps.sweeps_for_sweeper(sweeper)
-            ],
-        ):
-            # This has to be called exactly here, otherwise laboneq will not identify the sweepable node
-            self.set_instrument_nodes_for_nt_sweep(exp, sweeper)
-
-            if len(self.nt_sweeps) > 0:
-                self.sweep_recursion_nt(qubits, options, exp)
-            else:
-                self.define_exp(qubits, options, exp)
-
     def split_batches(self, sequences):
         return batch_max_sequences(sequences, MAX_SEQUENCES)

From afffcc64ad53601ce5ee58685fc6dfcaa8acab3d Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 19 Mar 2024 10:33:47 +0400
Subject: [PATCH 41/45] remove unnecessarry IIR check

---
 src/qibolab/instruments/zhinst/pulse.py | 77 ++++++++++++-------------
 1 file changed, 38 insertions(+), 39 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/pulse.py b/src/qibolab/instruments/zhinst/pulse.py
index 14c0b18e7..44c223ea6 100644
--- a/src/qibolab/instruments/zhinst/pulse.py
+++ b/src/qibolab/instruments/zhinst/pulse.py
@@ -9,7 +9,7 @@
     sampled_pulse_real,
 )
 
-from qibolab.pulses import Pulse, PulseType
+from qibolab.pulses import Drag, Gaussian, GaussianSquare, Pulse, PulseType, Rectangular
 from qibolab.sweeper import Parameter
 
 from .util import NANO_TO_SECONDS, SAMPLING_RATE
@@ -17,46 +17,45 @@
 
 def select_pulse(pulse: Pulse):
     """Return laboneq pulse object corresponding to the given qibolab pulse."""
-    if "IIR" not in str(pulse.shape):
-        if str(pulse.shape) == "Rectangular()":
-            can_compress = pulse.type is not PulseType.READOUT
-            return lo.pulse_library.const(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                can_compress=can_compress,
-            )
-        if "Gaussian" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            return lo.pulse_library.gaussian(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                sigma=2 / sigma,
-                zero_boundaries=False,
-            )
+    if isinstance(pulse.shape, Rectangular):
+        can_compress = pulse.type is not PulseType.READOUT
+        return lo.pulse_library.const(
+            length=round(pulse.duration * NANO_TO_SECONDS, 9),
+            amplitude=pulse.amplitude,
+            can_compress=can_compress,
+        )
+    if isinstance(pulse.shape, Gaussian):
+        sigma = pulse.shape.rel_sigma
+        return lo.pulse_library.gaussian(
+            length=round(pulse.duration * NANO_TO_SECONDS, 9),
+            amplitude=pulse.amplitude,
+            sigma=2 / sigma,
+            zero_boundaries=False,
+        )
 
-        if "GaussianSquare" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            width = pulse.shape.width
-            can_compress = pulse.type is not PulseType.READOUT
-            return lo.pulse_library.gaussian_square(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
-                amplitude=pulse.amplitude,
-                can_compress=can_compress,
-                sigma=2 / sigma,
-                zero_boundaries=False,
-            )
+    if isinstance(pulse.shape, GaussianSquare):
+        sigma = pulse.shape.rel_sigma
+        width = pulse.shape.width
+        can_compress = pulse.type is not PulseType.READOUT
+        return lo.pulse_library.gaussian_square(
+            length=round(pulse.duration * NANO_TO_SECONDS, 9),
+            width=round(pulse.duration * NANO_TO_SECONDS, 9) * width,
+            amplitude=pulse.amplitude,
+            can_compress=can_compress,
+            sigma=2 / sigma,
+            zero_boundaries=False,
+        )
 
-        if "Drag" in str(pulse.shape):
-            sigma = pulse.shape.rel_sigma
-            beta = pulse.shape.beta
-            return lo.pulse_library.drag(
-                length=round(pulse.duration * NANO_TO_SECONDS, 9),
-                amplitude=pulse.amplitude,
-                sigma=2 / sigma,
-                beta=beta,
-                zero_boundaries=False,
-            )
+    if isinstance(pulse.shape, Drag):
+        sigma = pulse.shape.rel_sigma
+        beta = pulse.shape.beta
+        return lo.pulse_library.drag(
+            length=round(pulse.duration * NANO_TO_SECONDS, 9),
+            amplitude=pulse.amplitude,
+            sigma=2 / sigma,
+            beta=beta,
+            zero_boundaries=False,
+        )
 
     if np.all(pulse.envelope_waveform_q(SAMPLING_RATE).data == 0):
         return sampled_pulse_real(

From ca57d390a0fb85e47645ef679c9a3d20df70d1fd Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Tue, 19 Mar 2024 10:54:07 +0400
Subject: [PATCH 42/45] temporary fix for readout pulse amplitude scaling

---
 src/qibolab/instruments/zhinst/sweep.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
index 8c99d6b90..b08c9e926 100644
--- a/src/qibolab/instruments/zhinst/sweep.py
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -95,6 +95,11 @@ def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
                     sweep_param = lo.SweepParameter(
                         values=sweeper.values / max(sweeper.values)
                     )
+                    # FIXME: this implicitly relies on the fact that pulse is the same python object as appears in the
+                    # sequence that is being executed, hence the mutation is propagated. This is bad programming and
+                    # should be fixed once things become simpler
+                    pulse.amplitude *= max(sweeper.values)
+
                     channel_sweeps.append(
                         (
                             measure_channel_name(qubits[pulse.qubit]),

From 27f33070b4a38c2b93d2568dcc9cd3c55f0249a4 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 20 Mar 2024 16:41:23 +0400
Subject: [PATCH 43/45] fix amplitude scaling if negative values are present

---
 src/qibolab/instruments/zhinst/executor.py | 5 +++--
 src/qibolab/instruments/zhinst/sweep.py    | 8 ++++----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/executor.py b/src/qibolab/instruments/zhinst/executor.py
index e8d3a1042..70088e85a 100644
--- a/src/qibolab/instruments/zhinst/executor.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -652,8 +652,9 @@ def play_sweep(exp, channel_name, pulse):
         play_parameters = {}
         for p, zhs in pulse.zhsweepers:
             if p is Parameter.amplitude:
-                pulse.zhpulse.amplitude *= max(zhs.values)
-                zhs.values /= max(zhs.values)
+                max_value = max(np.abs(zhs.values))
+                pulse.zhpulse.amplitude *= max_value
+                zhs.values /= max_value
                 play_parameters["amplitude"] = zhs
             if p is Parameter.duration:
                 play_parameters["length"] = zhs
diff --git a/src/qibolab/instruments/zhinst/sweep.py b/src/qibolab/instruments/zhinst/sweep.py
index b08c9e926..d2371c79e 100644
--- a/src/qibolab/instruments/zhinst/sweep.py
+++ b/src/qibolab/instruments/zhinst/sweep.py
@@ -4,6 +4,7 @@
 from copy import copy
 
 import laboneq.simple as lo
+import numpy as np
 
 from qibolab.pulses import Pulse, PulseType
 from qibolab.qubits import Qubit
@@ -92,13 +93,12 @@ def __init__(self, sweepers: Iterable[Sweeper], qubits: dict[str, Qubit]):
                     pulse.type is PulseType.READOUT
                     and sweeper.parameter is Parameter.amplitude
                 ):
-                    sweep_param = lo.SweepParameter(
-                        values=sweeper.values / max(sweeper.values)
-                    )
+                    max_value = max(np.abs(sweeper.values))
+                    sweep_param = lo.SweepParameter(values=sweeper.values / max_value)
                     # FIXME: this implicitly relies on the fact that pulse is the same python object as appears in the
                     # sequence that is being executed, hence the mutation is propagated. This is bad programming and
                     # should be fixed once things become simpler
-                    pulse.amplitude *= max(sweeper.values)
+                    pulse.amplitude *= max_value
 
                     channel_sweeps.append(
                         (

From 798906edfc35591f6e756d61c0d108cdbbf0d240 Mon Sep 17 00:00:00 2001
From: Hayk Sargsyan <52532457+hay-k@users.noreply.github.com>
Date: Wed, 20 Mar 2024 16:53:37 +0400
Subject: [PATCH 44/45] port changes from main

---
 src/qibolab/instruments/zhinst/executor.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/src/qibolab/instruments/zhinst/executor.py b/src/qibolab/instruments/zhinst/executor.py
index 70088e85a..738d314c3 100644
--- a/src/qibolab/instruments/zhinst/executor.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -11,15 +11,13 @@
 
 from qibolab import AcquisitionType, AveragingMode, ExecutionParameters
 from qibolab.couplers import Coupler
-
-from qibolab.instruments.unrolling import batch_max_sequences
+from qibolab.instruments.abstract import Controller
+from qibolab.instruments.port import Port
 from qibolab.pulses import PulseSequence, PulseType
 from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
 
-from .abstract import Controller
-from .port import Port
 
 from .pulse import ZhPulse
 from .sweep import ProcessedSweeps, classify_sweepers
@@ -740,6 +738,3 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
                 results[serial] = results[qubit] = options.results_type(data)
 
         return results
-
-    def split_batches(self, sequences):
-        return batch_max_sequences(sequences, MAX_SEQUENCES)

From 92f964028b3957b681f05f17189b55d885400ca2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 20 Mar 2024 13:12:53 +0000
Subject: [PATCH 45/45] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/qibolab/instruments/zhinst/executor.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/qibolab/instruments/zhinst/executor.py b/src/qibolab/instruments/zhinst/executor.py
index 738d314c3..ff73fde62 100644
--- a/src/qibolab/instruments/zhinst/executor.py
+++ b/src/qibolab/instruments/zhinst/executor.py
@@ -18,7 +18,6 @@
 from qibolab.sweeper import Parameter, Sweeper
 from qibolab.unrolling import Bounds
 
-
 from .pulse import ZhPulse
 from .sweep import ProcessedSweeps, classify_sweepers
 from .util import (