From 96656f50ba055d6077e6ec746105bfdda5d53564 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 8 Jan 2024 16:48:05 +0400
Subject: [PATCH 01/19] feat: Support sequence unrolling for QM

---
 src/qibolab/instruments/qm/acquisition.py | 74 +++++++++++++++--------
 src/qibolab/instruments/qm/config.py      | 23 +++----
 src/qibolab/instruments/qm/controller.py  | 27 ++++++---
 src/qibolab/instruments/qm/sequence.py    | 42 +++++++++----
 4 files changed, 111 insertions(+), 55 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 44f1cb61e..80249522b 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -27,9 +27,11 @@ class Acquisition(ABC):
     variables.
     """
 
-    serial: str
-    """Serial of the readout pulse that generates this acquisition."""
+    name: str
+    """Name of the acquisition used as identifier to download results from the
+    instruments."""
     average: bool
+    npulses: int
 
     @abstractmethod
     def assign_element(self, element):
@@ -91,19 +93,18 @@ def download(self, *dimensions):
         if self.average:
             i_stream = i_stream.average()
             q_stream = q_stream.average()
-        i_stream.save(f"{self.serial}_I")
-        q_stream.save(f"{self.serial}_Q")
+        i_stream.save(f"{self.name}_I")
+        q_stream.save(f"{self.name}_Q")
 
     def fetch(self, handles):
-        ires = handles.get(f"{self.serial}_I").fetch_all()
-        qres = handles.get(f"{self.serial}_Q").fetch_all()
+        ires = handles.get(f"{self.name}_I").fetch_all()
+        qres = handles.get(f"{self.name}_Q").fetch_all()
         # convert raw ADC signal to volts
         u = unit()
-        ires = u.raw2volts(ires)
-        qres = u.raw2volts(qres)
+        signal = u.raw2volts(ires) + 1j * u.raw2volts(qres)
         if self.average:
-            return AveragedRawWaveformResults(ires + 1j * qres)
-        return RawWaveformResults(ires + 1j * qres)
+            return [AveragedRawWaveformResults(signal)]
+        return [RawWaveformResults(signal)]
 
 
 @dataclass
@@ -136,22 +137,35 @@ def save(self):
     def download(self, *dimensions):
         Istream = self.I_stream
         Qstream = self.Q_stream
+        if self.npulses > 1:
+            Istream = Istream.buffer(self.npulses)
+            Qstream = Qstream.buffer(self.npulses)
         for dim in dimensions:
             Istream = Istream.buffer(dim)
             Qstream = Qstream.buffer(dim)
         if self.average:
             Istream = Istream.average()
             Qstream = Qstream.average()
-        Istream.save(f"{self.serial}_I")
-        Qstream.save(f"{self.serial}_Q")
+        Istream.save(f"{self.name}_I")
+        Qstream.save(f"{self.name}_Q")
 
     def fetch(self, handles):
-        ires = handles.get(f"{self.serial}_I").fetch_all()
-        qres = handles.get(f"{self.serial}_Q").fetch_all()
-        if self.average:
-            # TODO: calculate std
-            return AveragedIntegratedResults(ires + 1j * qres)
-        return IntegratedResults(ires + 1j * qres)
+        ires = handles.get(f"{self.name}_I").fetch_all()
+        qres = handles.get(f"{self.name}_Q").fetch_all()
+        signal = ires + 1j * qres
+        if self.npulses > 1:
+            if self.average:
+                # TODO: calculate std
+                return [
+                    AveragedIntegratedResults(signal[..., i])
+                    for i in range(self.npulses)
+                ]
+            return [IntegratedResults(signal[..., i]) for i in range(self.npulses)]
+        else:
+            if self.average:
+                # TODO: calculate std
+                return [AveragedIntegratedResults(signal)]
+            return [IntegratedResults(signal)]
 
 
 @dataclass
@@ -199,15 +213,27 @@ def save(self):
 
     def download(self, *dimensions):
         shots = self.shots
+        if self.npulses > 1:
+            shots = shots.buffer(self.npulses)
         for dim in dimensions:
             shots = shots.buffer(dim)
         if self.average:
             shots = shots.average()
-        shots.save(f"{self.serial}_shots")
+        shots.save(f"{self.name}_shots")
 
     def fetch(self, handles):
-        shots = handles.get(f"{self.serial}_shots").fetch_all()
-        if self.average:
-            # TODO: calculate std
-            return AveragedSampleResults(shots)
-        return SampleResults(shots.astype(int))
+        shots = handles.get(f"{self.name}_shots").fetch_all()
+        if len(self.npulses) > 1:
+            if self.average:
+                # TODO: calculate std
+                return [
+                    AveragedSampleResults(shots[..., i]) for i in range(self.npulses)
+                ]
+            return [
+                SampleResults(shots[..., i].astype(int)) for i in range(self.npulses)
+            ]
+        else:
+            if self.average:
+                # TODO: calculate std
+                return [AveragedSampleResults(shots)]
+            return [SampleResults(shots.astype(int))]
diff --git a/src/qibolab/instruments/qm/config.py b/src/qibolab/instruments/qm/config.py
index 016ce6205..5e58d64ef 100644
--- a/src/qibolab/instruments/qm/config.py
+++ b/src/qibolab/instruments/qm/config.py
@@ -228,7 +228,7 @@ def register_element(self, qubit, pulse, time_of_flight=0, smearing=0):
             # register flux element
             self.register_flux_element(qubit, pulse.frequency)
 
-    def register_pulse(self, qubit, pulse):
+    def register_pulse(self, qubit, qmpulse):
         """Registers pulse, waveforms and integration weights in QM config.
 
         Args:
@@ -241,23 +241,24 @@ def register_pulse(self, qubit, pulse):
                 instantiation of the Qubit objects. They are named as
                 "drive0", "drive1", "flux0", "readout0", ...
         """
-        if pulse.serial not in self.pulses:
+        pulse = qmpulse.pulse
+        if qmpulse.operation not in self.pulses:
             if pulse.type is PulseType.DRIVE:
                 serial_i = self.register_waveform(pulse, "i")
                 serial_q = self.register_waveform(pulse, "q")
-                self.pulses[pulse.serial] = {
+                self.pulses[qmpulse.operation] = {
                     "operation": "control",
                     "length": pulse.duration,
                     "waveforms": {"I": serial_i, "Q": serial_q},
                 }
                 # register drive pulse in elements
                 self.elements[f"drive{qubit.name}"]["operations"][
-                    pulse.serial
-                ] = pulse.serial
+                    qmpulse.operation
+                ] = qmpulse.operation
 
             elif pulse.type is PulseType.FLUX:
                 serial = self.register_waveform(pulse)
-                self.pulses[pulse.serial] = {
+                self.pulses[qmpulse.operation] = {
                     "operation": "control",
                     "length": pulse.duration,
                     "waveforms": {
@@ -266,14 +267,14 @@ def register_pulse(self, qubit, pulse):
                 }
                 # register flux pulse in elements
                 self.elements[f"flux{qubit.name}"]["operations"][
-                    pulse.serial
-                ] = pulse.serial
+                    qmpulse.operation
+                ] = qmpulse.operation
 
             elif pulse.type is PulseType.READOUT:
                 serial_i = self.register_waveform(pulse, "i")
                 serial_q = self.register_waveform(pulse, "q")
                 self.register_integration_weights(qubit, pulse.duration)
-                self.pulses[pulse.serial] = {
+                self.pulses[qmpulse.operation] = {
                     "operation": "measurement",
                     "length": pulse.duration,
                     "waveforms": {
@@ -289,8 +290,8 @@ def register_pulse(self, qubit, pulse):
                 }
                 # register readout pulse in elements
                 self.elements[f"readout{qubit.name}"]["operations"][
-                    pulse.serial
-                ] = pulse.serial
+                    qmpulse.operation
+                ] = qmpulse.operation
 
             else:
                 raise_error(TypeError, f"Unknown pulse type {pulse.type.name}.")
diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index d81917cda..8ca52e251 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -252,10 +252,9 @@ def fetch_results(result, ro_pulses):
         handles.wait_for_all_values()
         results = {}
         for qmpulse in ro_pulses:
-            pulse = qmpulse.pulse
-            results[pulse.qubit] = results[pulse.serial] = qmpulse.acquisition.fetch(
-                handles
-            )
+            data = qmpulse.acquisition.fetch(handles)
+            for pulse, result in zip(qmpulse.pulses, data):
+                results[pulse.qubit] = results[pulse.serial] = result
         return results
 
     def create_sequence(self, qubits, sequence, sweepers):
@@ -276,6 +275,17 @@ def create_sequence(self, qubits, sequence, sweepers):
 
         duration_sweep_pulses = find_duration_sweeper_pulses(sweepers)
 
+        qmpulses = {}
+
+        def create_qmpulse(pulse, qmpulse_cls):
+            qmpulse = qmpulse_cls(pulse)
+            key = (qmpulse.operation, pulse.qubit)
+            if key not in qmpulses:
+                qmpulses[key] = qmpulse
+            else:
+                qmpulses[key].pulses.append(pulse)
+            return qmpulses[key]
+
         qmsequence = Sequence()
         sort_key = lambda pulse: (pulse.start, pulse.duration)
         for pulse in sorted(sequence.pulses, key=sort_key):
@@ -293,11 +303,11 @@ def create_sequence(self, qubits, sequence, sweepers):
                 or pulse.duration < 16
                 or pulse.serial in duration_sweep_pulses
             ):
-                qmpulse = BakedPulse(pulse)
+                qmpulse = create_qmpulse(pulse, BakedPulse)
                 qmpulse.bake(self.config, durations=[pulse.duration])
             else:
-                qmpulse = QMPulse(pulse)
-                self.config.register_pulse(qubit, pulse)
+                qmpulse = create_qmpulse(pulse, QMPulse)
+                self.config.register_pulse(qubit, qmpulse)
             qmsequence.add(qmpulse)
 
         qmsequence.shift()
@@ -349,3 +359,6 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
 
         result = self.execute_program(experiment)
         return self.fetch_results(result, qmsequence.ro_pulses)
+
+    def split_batches(self, sequences):
+        return [sequences]
diff --git a/src/qibolab/instruments/qm/sequence.py b/src/qibolab/instruments/qm/sequence.py
index 1457a1871..4522455af 100644
--- a/src/qibolab/instruments/qm/sequence.py
+++ b/src/qibolab/instruments/qm/sequence.py
@@ -35,7 +35,7 @@ class QMPulse:
     as defined in the QM config.
     """
 
-    pulse: Pulse
+    pulses: List[Pulse]
     """:class:`qibolab.pulses.Pulse` implemting the current pulse."""
     element: Optional[str] = None
     """Element that the pulse will be played on, as defined in the QM
@@ -73,11 +73,22 @@ class QMPulse:
     elements_to_align: Set[str] = field(default_factory=set)
 
     def __post_init__(self):
-        self.element: str = f"{self.pulse.type.name.lower()}{self.pulse.qubit}"
-        self.operation: str = self.pulse.serial
+        if isinstance(self.pulses, Pulse):
+            self.pulses = [self.pulses]
+
+        pulse_type = self.pulse.type.name.lower()
+        amplitude = format(self.pulse.amplitude, ".6f").rstrip("0").rstrip(".")
+        self.element: str = f"{pulse_type}{self.pulse.qubit}"
+        self.operation: str = (
+            f"{pulse_type}({self.pulse.duration}, {amplitude}, {self.pulse.shape})"
+        )
         self.relative_phase: float = self.pulse.relative_phase / (2 * np.pi)
         self.elements_to_align.add(self.element)
 
+    @property
+    def pulse(self):
+        return self.pulses[0]
+
     def __hash__(self):
         return hash(self.pulse)
 
@@ -115,10 +126,15 @@ def play(self):
     def declare_output(self, options, threshold=None, angle=None):
         average = options.averaging_mode is AveragingMode.CYCLIC
         acquisition_type = options.acquisition_type
+        acquisition_name = f"{self.operation}_{self.pulse.qubit}"
         if acquisition_type is AcquisitionType.RAW:
-            self.acquisition = RawAcquisition(self.pulse.serial, average)
+            self.acquisition = RawAcquisition(
+                acquisition_name, average, len(self.pulses)
+            )
         elif acquisition_type is AcquisitionType.INTEGRATION:
-            self.acquisition = IntegratedAcquisition(self.pulse.serial, average)
+            self.acquisition = IntegratedAcquisition(
+                acquisition_name, average, len(self.pulses)
+            )
         elif acquisition_type is AcquisitionType.DISCRIMINATION:
             if threshold is None or angle is None:
                 raise_error(
@@ -127,7 +143,7 @@ def declare_output(self, options, threshold=None, angle=None):
                     "if threshold and angle are not given.",
                 )
             self.acquisition = ShotsAcquisition(
-                self.pulse.serial, average, threshold, angle
+                acquisition_name, average, len(self.pulses), threshold, angle
             )
         else:
             raise_error(ValueError, f"Invalid acquisition type {acquisition_type}.")
@@ -185,8 +201,8 @@ def bake(self, config: QMConfig, durations: DurationsType):
                         self.calculate_waveform(waveform_i, t),
                         self.calculate_waveform(waveform_q, t),
                     ]
-                segment.add_op(self.pulse.serial, self.element, waveform)
-                segment.play(self.pulse.serial, self.element)
+                segment.add_op(self.operation, self.element, waveform)
+                segment.play(self.operation, self.element)
             self.segments.append(segment)
 
     @property
@@ -248,7 +264,9 @@ def add(self, qmpulse: QMPulse):
         pulse = qmpulse.pulse
         self.pulse_to_qmpulse[pulse.serial] = qmpulse
         if pulse.type is PulseType.READOUT:
-            self.ro_pulses.append(qmpulse)
+            if len(qmpulse.pulses) == 1:
+                # if ``qmpulse`` contains more than one pulses it is already added in ``ro_pulses``
+                self.ro_pulses.append(qmpulse)
 
         previous = self._find_previous(pulse)
         if previous is not None:
@@ -287,7 +305,9 @@ def play(self, relaxation_time=0):
             if qmpulse.wait_cycles is not None:
                 qua.wait(qmpulse.wait_cycles, qmpulse.element)
             if pulse.type is PulseType.READOUT:
+                # Save data to the stream processing
                 qmpulse.acquisition.measure(qmpulse.operation, qmpulse.element)
+                qmpulse.acquisition.save()
             else:
                 if (
                     not isinstance(qmpulse.relative_phase, float)
@@ -304,7 +324,3 @@ def play(self, relaxation_time=0):
 
         if relaxation_time > 0:
             qua.wait(relaxation_time // 4)
-
-        # Save data to the stream processing
-        for qmpulse in self.ro_pulses:
-            qmpulse.acquisition.save()

From 4c169d88ad17a7c8f2168c921fd42177b77d0704 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 9 Jan 2024 15:41:13 +0400
Subject: [PATCH 02/19] fix: Recover old QMPulse to fix pulse scheduling

---
 src/qibolab/instruments/qm/acquisition.py | 64 ++++++++++++++++++++---
 src/qibolab/instruments/qm/controller.py  | 34 ++++++------
 src/qibolab/instruments/qm/sequence.py    | 52 ++----------------
 src/qibolab/instruments/qm/simulator.py   |  2 +-
 src/qibolab/platform.py                   |  5 ++
 5 files changed, 84 insertions(+), 73 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 80249522b..b60fd6efc 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
+from typing import List, Optional
 
 import numpy as np
 from qm import qua
@@ -8,6 +9,8 @@
 from qualang_tools.addons.variables import assign_variables_to_element
 from qualang_tools.units import unit
 
+from qibolab.execution_parameters import AcquisitionType, AveragingMode
+from qibolab.qubits import QubitId
 from qibolab.result import (
     AveragedIntegratedResults,
     AveragedRawWaveformResults,
@@ -30,8 +33,18 @@ class Acquisition(ABC):
     name: str
     """Name of the acquisition used as identifier to download results from the
     instruments."""
+    qubit: QubitId
     average: bool
-    npulses: int
+    threshold: Optional[float] = None
+    """Threshold to be used for classification of single shots."""
+    angle: Optional[float] = None
+    """Angle in the IQ plane to be used for classification of single shots."""
+
+    keys: List[str] = field(default_factory=list)
+
+    @property
+    def npulses(self):
+        return len(self.keys)
 
     @abstractmethod
     def assign_element(self, element):
@@ -175,11 +188,6 @@ class ShotsAcquisition(Acquisition):
     Threshold and angle must be given in order to classify shots.
     """
 
-    threshold: float
-    """Threshold to be used for classification of single shots."""
-    angle: float
-    """Angle in the IQ plane to be used for classification of single shots."""
-
     I: _Variable = field(default_factory=lambda: declare(fixed))
     Q: _Variable = field(default_factory=lambda: declare(fixed))
     """Variables to save the (I, Q) values acquired from a single shot."""
@@ -189,6 +197,12 @@ class ShotsAcquisition(Acquisition):
     """Stream to collect multiple shots."""
 
     def __post_init__(self):
+        if threshold is None or angle is None:
+            raise_error(
+                ValueError,
+                "Cannot use ``AcquisitionType.DISCRIMINATION`` "
+                "if threshold and angle are not given.",
+            )
         self.cos = np.cos(self.angle)
         self.sin = np.sin(self.angle)
 
@@ -237,3 +251,41 @@ def fetch(self, handles):
                 # TODO: calculate std
                 return [AveragedSampleResults(shots)]
             return [SampleResults(shots.astype(int))]
+
+
+ACQUISITION_TYPES = {
+    AcquisitionType.RAW: RawAcquisition,
+    AcquisitionType.INTEGRATION: IntegratedAcquisition,
+    AcquisitionType.DISCRIMINATION: ShotsAcquisition,
+}
+
+
+def declare_acquisitions(ro_pulses, qubits, options):
+    """Declares variables for saving acquisition in the QUA program.
+
+    Args:
+        ro_pulses (list): List of readout pulses in the sequence.
+        qubits (dict): Dictionary containing all the :class:`qibolab.qubits.Qubit`
+            objects of the platform.
+        options (:class:`qibolab.execution_parameters.ExecutionParameters`): Execution
+            options containing acquisition type and averaging mode.
+
+    Returns:
+        Dictionary containing the different :class:`qibolab.instruments.qm.acquisition.Acquisition` objects.
+    """
+    acquisitions = {}
+    for qmpulse in ro_pulses:
+        qubit = qmpulse.pulse.qubit
+        name = f"{qmpulse.operation}_{qubit}"
+        if name not in acquisitions:
+            threshold = qubits[qubit].threshold
+            iq_angle = qubits[qubit].iq_angle
+            average = options.averaging_mode is AveragingMode.CYCLIC
+            acquisition_cls = ACQUISITION_TYPES[options.acquisition_type]
+            acquisition = acquisition_cls(name, qubit, options, threshold, iq_angle)
+            acquisition.assign_element(qmpulse.element)
+            acquisitions[name] = acquisition
+
+        acquisitions[name].keys.append(qmpulse.pulse.serial)
+        qmpulse.acquisition = acquisitions[name]
+    return acquisitions
diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index 8ca52e251..466854bc4 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -11,6 +11,7 @@
 from qibolab.pulses import PulseType
 from qibolab.sweeper import Parameter
 
+from .acquisition import declare_acquisitions
 from .config import SAMPLING_RATE, QMConfig
 from .devices import Octave, OPXplus
 from .ports import OPXIQ
@@ -239,7 +240,7 @@ def execute_program(self, program):
         return machine.execute(program)
 
     @staticmethod
-    def fetch_results(result, ro_pulses):
+    def fetch_results(result, acquisitions):
         """Fetches results from an executed experiment.
 
         Defined as ``@staticmethod`` because it is overwritten
@@ -251,10 +252,10 @@ def fetch_results(result, ro_pulses):
         handles = result.result_handles
         handles.wait_for_all_values()
         results = {}
-        for qmpulse in ro_pulses:
-            data = qmpulse.acquisition.fetch(handles)
-            for pulse, result in zip(qmpulse.pulses, data):
-                results[pulse.qubit] = results[pulse.serial] = result
+        for acquisition in acquisitions.values():
+            data = acquisition.fetch(handles)
+            for serial, result in zip(acquisition.keys, data):
+                results[acquisition.qubit] = results[serial] = result
         return results
 
     def create_sequence(self, qubits, sequence, sweepers):
@@ -287,6 +288,7 @@ def create_qmpulse(pulse, qmpulse_cls):
             return qmpulses[key]
 
         qmsequence = Sequence()
+        ro_pulses = []
         sort_key = lambda pulse: (pulse.start, pulse.duration)
         for pulse in sorted(sequence.pulses, key=sort_key):
             qubit = qubits[pulse.qubit]
@@ -303,15 +305,17 @@ def create_qmpulse(pulse, qmpulse_cls):
                 or pulse.duration < 16
                 or pulse.serial in duration_sweep_pulses
             ):
-                qmpulse = create_qmpulse(pulse, BakedPulse)
+                qmpulse = BakedPulse(pulse)
                 qmpulse.bake(self.config, durations=[pulse.duration])
             else:
-                qmpulse = create_qmpulse(pulse, QMPulse)
+                qmpulse = QMPulse(pulse)
+                if pulse.type is PulseType.READOUT:
+                    ro_pulses.append(qmpulse)
                 self.config.register_pulse(qubit, qmpulse)
             qmsequence.add(qmpulse)
 
         qmsequence.shift()
-        return qmsequence
+        return qmsequence, ro_pulses
 
     def play(self, qubits, couplers, sequence, options):
         return self.sweep(qubits, couplers, sequence, options)
@@ -331,15 +335,11 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
                 self.config.register_port(qubit.flux.port)
                 self.config.register_flux_element(qubit)
 
-        qmsequence = self.create_sequence(qubits, sequence, sweepers)
+        qmsequence, ro_pulses = self.create_sequence(qubits, sequence, sweepers)
         # play pulses using QUA
         with qua.program() as experiment:
             n = declare(int)
-            for qmpulse in qmsequence.ro_pulses:
-                threshold = qubits[qmpulse.pulse.qubit].threshold
-                iq_angle = qubits[qmpulse.pulse.qubit].iq_angle
-                qmpulse.declare_output(options, threshold, iq_angle)
-
+            acquisitions = declare_acquisitions(ro_pulses, qubits, options)
             with for_(n, 0, n < options.nshots, n + 1):
                 sweep(
                     list(sweepers),
@@ -350,15 +350,15 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
                 )
 
             with qua.stream_processing():
-                for qmpulse in qmsequence.ro_pulses:
-                    qmpulse.acquisition.download(*buffer_dims)
+                for acquisition in acquisitions.values():
+                    acquisition.download(*buffer_dims)
 
         if self.script_file_name is not None:
             with open(self.script_file_name, "w") as file:
                 file.write(generate_qua_script(experiment, self.config.__dict__))
 
         result = self.execute_program(experiment)
-        return self.fetch_results(result, qmsequence.ro_pulses)
+        return self.fetch_results(result, acquisitions)
 
     def split_batches(self, sequences):
         return [sequences]
diff --git a/src/qibolab/instruments/qm/sequence.py b/src/qibolab/instruments/qm/sequence.py
index 4522455af..d7127f788 100644
--- a/src/qibolab/instruments/qm/sequence.py
+++ b/src/qibolab/instruments/qm/sequence.py
@@ -4,19 +4,12 @@
 
 import numpy as np
 from numpy import typing as npt
-from qibo.config import raise_error
 from qm import qua
 from qm.qua._dsl import _Variable  # for type declaration only
 from qualang_tools.bakery import baking
 from qualang_tools.bakery.bakery import Baking
 
-from qibolab import AcquisitionType, AveragingMode
-from qibolab.instruments.qm.acquisition import (
-    Acquisition,
-    IntegratedAcquisition,
-    RawAcquisition,
-    ShotsAcquisition,
-)
+from qibolab.instruments.qm.acquisition import Acquisition
 from qibolab.pulses import Pulse, PulseType
 
 from .config import SAMPLING_RATE, QMConfig
@@ -35,8 +28,8 @@ class QMPulse:
     as defined in the QM config.
     """
 
-    pulses: List[Pulse]
-    """:class:`qibolab.pulses.Pulse` implemting the current pulse."""
+    pulse: Pulse
+    """:class:`qibolab.pulses.Pulse` corresponding to the ``QMPulse``."""
     element: Optional[str] = None
     """Element that the pulse will be played on, as defined in the QM
     config."""
@@ -73,9 +66,6 @@ class QMPulse:
     elements_to_align: Set[str] = field(default_factory=set)
 
     def __post_init__(self):
-        if isinstance(self.pulses, Pulse):
-            self.pulses = [self.pulses]
-
         pulse_type = self.pulse.type.name.lower()
         amplitude = format(self.pulse.amplitude, ".6f").rstrip("0").rstrip(".")
         self.element: str = f"{pulse_type}{self.pulse.qubit}"
@@ -85,10 +75,6 @@ def __post_init__(self):
         self.relative_phase: float = self.pulse.relative_phase / (2 * np.pi)
         self.elements_to_align.add(self.element)
 
-    @property
-    def pulse(self):
-        return self.pulses[0]
-
     def __hash__(self):
         return hash(self.pulse)
 
@@ -123,32 +109,6 @@ def play(self):
         """
         qua.play(self.operation, self.element, duration=self.swept_duration)
 
-    def declare_output(self, options, threshold=None, angle=None):
-        average = options.averaging_mode is AveragingMode.CYCLIC
-        acquisition_type = options.acquisition_type
-        acquisition_name = f"{self.operation}_{self.pulse.qubit}"
-        if acquisition_type is AcquisitionType.RAW:
-            self.acquisition = RawAcquisition(
-                acquisition_name, average, len(self.pulses)
-            )
-        elif acquisition_type is AcquisitionType.INTEGRATION:
-            self.acquisition = IntegratedAcquisition(
-                acquisition_name, average, len(self.pulses)
-            )
-        elif acquisition_type is AcquisitionType.DISCRIMINATION:
-            if threshold is None or angle is None:
-                raise_error(
-                    ValueError,
-                    "Cannot use ``AcquisitionType.DISCRIMINATION`` "
-                    "if threshold and angle are not given.",
-                )
-            self.acquisition = ShotsAcquisition(
-                acquisition_name, average, len(self.pulses), threshold, angle
-            )
-        else:
-            raise_error(ValueError, f"Invalid acquisition type {acquisition_type}.")
-        self.acquisition.assign_element(self.element)
-
 
 @dataclass
 class BakedPulse(QMPulse):
@@ -234,8 +194,6 @@ class Sequence:
     qmpulses: List[QMPulse] = field(default_factory=list)
     """List of :class:`qibolab.instruments.qm.QMPulse` objects corresponding to
     the original pulses."""
-    ro_pulses: List[QMPulse] = field(default_factory=list)
-    """List of readout pulses used for registering outputs."""
     pulse_to_qmpulse: Dict[Pulse, QMPulse] = field(default_factory=dict)
     """Map from qibolab pulses to QMPulses (useful when sweeping)."""
     clock: Dict[str, int] = field(default_factory=lambda: collections.defaultdict(int))
@@ -263,10 +221,6 @@ def _find_previous(self, pulse):
     def add(self, qmpulse: QMPulse):
         pulse = qmpulse.pulse
         self.pulse_to_qmpulse[pulse.serial] = qmpulse
-        if pulse.type is PulseType.READOUT:
-            if len(qmpulse.pulses) == 1:
-                # if ``qmpulse`` contains more than one pulses it is already added in ``ro_pulses``
-                self.ro_pulses.append(qmpulse)
 
         previous = self._find_previous(pulse)
         if previous is not None:
diff --git a/src/qibolab/instruments/qm/simulator.py b/src/qibolab/instruments/qm/simulator.py
index 91a7a5f6a..2542aaa72 100644
--- a/src/qibolab/instruments/qm/simulator.py
+++ b/src/qibolab/instruments/qm/simulator.py
@@ -42,7 +42,7 @@ def connect(self):
             self.manager = QuantumMachinesManager(host, int(port))
 
     @staticmethod
-    def fetch_results(result, ro_pulses):
+    def fetch_results(result, acquisitions):
         return result
 
     def execute_program(self, program):
diff --git a/src/qibolab/platform.py b/src/qibolab/platform.py
index 60da8c097..0848792b1 100644
--- a/src/qibolab/platform.py
+++ b/src/qibolab/platform.py
@@ -283,6 +283,11 @@ def execute_pulse_sequences(
         results = defaultdict(list)
         for batch in self._controller.split_batches(sequences):
             sequence, readouts = unroll_sequences(batch, options.relaxation_time)
+
+            print()
+            print(sequence)
+            print()
+
             result = self._execute(sequence, options, **kwargs)
             for serial, new_serials in readouts.items():
                 results[serial].extend(result[ser] for ser in new_serials)

From ecd161157c35afc75e3928f6c49288bb88d40cc0 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 9 Jan 2024 17:28:56 +0400
Subject: [PATCH 03/19] refactor: Remove QMSim class

---
 src/qibolab/instruments/qm/__init__.py    |  1 -
 src/qibolab/instruments/qm/acquisition.py | 30 ++++++--
 src/qibolab/instruments/qm/controller.py  | 86 +++++++++++++----------
 src/qibolab/instruments/qm/simulator.py   | 55 ---------------
 src/qibolab/platform.py                   | 12 ----
 5 files changed, 76 insertions(+), 108 deletions(-)
 delete mode 100644 src/qibolab/instruments/qm/simulator.py

diff --git a/src/qibolab/instruments/qm/__init__.py b/src/qibolab/instruments/qm/__init__.py
index 041a7b916..e053aa970 100644
--- a/src/qibolab/instruments/qm/__init__.py
+++ b/src/qibolab/instruments/qm/__init__.py
@@ -1,3 +1,2 @@
 from .controller import QMController
 from .devices import Octave, OPXplus
-from .simulator import QMSim
diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index b60fd6efc..c9efd1c08 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -197,11 +197,10 @@ class ShotsAcquisition(Acquisition):
     """Stream to collect multiple shots."""
 
     def __post_init__(self):
-        if threshold is None or angle is None:
-            raise_error(
-                ValueError,
+        if self.threshold is None or self.angle is None:
+            raise ValueError(
                 "Cannot use ``AcquisitionType.DISCRIMINATION`` "
-                "if threshold and angle are not given.",
+                "if threshold and angle are not given."
             )
         self.cos = np.cos(self.angle)
         self.sin = np.sin(self.angle)
@@ -289,3 +288,26 @@ def declare_acquisitions(ro_pulses, qubits, options):
         acquisitions[name].keys.append(qmpulse.pulse.serial)
         qmpulse.acquisition = acquisitions[name]
     return acquisitions
+
+
+def fetch_results(result, acquisitions):
+    """Fetches results from an executed experiment.
+
+    Args:
+        result: Result of the executed experiment.
+        acquisition (dict): Dictionary containing :class:`qibolab.instruments.qm.acquisition.Acquisition` objects.
+
+    Returns:
+        Dictionary with the results in the format required by the platform.
+    """
+    # TODO: Update result asynchronously instead of waiting
+    # for all values, in order to allow live plotting
+    # using ``handles.is_processing()``
+    handles = result.result_handles
+    handles.wait_for_all_values()
+    results = {}
+    for acquisition in acquisitions.values():
+        data = acquisition.fetch(handles)
+        for serial, result in zip(acquisition.keys, data):
+            results[acquisition.qubit] = results[serial] = result
+    return results
diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index 466854bc4..61ddcbb20 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -1,17 +1,18 @@
 from dataclasses import dataclass, field
 from typing import Dict, Optional
 
-from qm import generate_qua_script, qua
+from qm import QuantumMachinesManager, SimulationConfig, generate_qua_script, qua
 from qm.octave import QmOctaveConfig
 from qm.qua import declare, for_
-from qm.QuantumMachinesManager import QuantumMachinesManager
+from qm.simulate.credentials import create_credentials
+from qualang_tools.simulator_tools import create_simulator_controller_connections
 
 from qibolab import AveragingMode
 from qibolab.instruments.abstract import Controller
 from qibolab.pulses import PulseType
 from qibolab.sweeper import Parameter
 
-from .acquisition import declare_acquisitions
+from .acquisition import declare_acquisitions, fetch_results
 from .config import SAMPLING_RATE, QMConfig
 from .devices import Octave, OPXplus
 from .ports import OPXIQ
@@ -140,6 +141,22 @@ class QMController(Controller):
     is_connected: bool = False
     """Boolean that shows whether we are connected to the QM manager."""
 
+    simulation_duration: Optional[int] = None
+    """Duration for the simulation in ns.
+
+    If given the simulator will be used instead of actual hardware
+    execution.
+    """
+    cloud: bool = False
+    """If ``True`` the QM cloud simulator is used which does not require access
+    to physical instruments.
+
+    This assumes that a proper cloud address has been given.
+    If ``False`` and ``simulation_duration`` was given, then the built-in simulator
+    of the instruments is used. This requires connection to instruments.
+    Default is ``False``.
+    """
+
     def __post_init__(self):
         super().__init__(self.name, self.address)
         # convert lists to dicts
@@ -148,6 +165,10 @@ def __post_init__(self):
         if not isinstance(self.octaves, dict):
             self.octaves = {instr.name: instr for instr in self.octaves}
 
+        if self.simulation_duration is not None:
+            # convert simulation duration from ns to clock cycles
+            self.simulation_duration //= 4
+
     def ports(self, name, input=False):
         """Provides instrument ports to the user.
 
@@ -182,7 +203,12 @@ def connect(self):
         """Connect to the Quantum Machines manager."""
         host, port = self.address.split(":")
         octave = declare_octaves(self.octaves, host, self.calibration_path)
-        self.manager = QuantumMachinesManager(host=host, port=int(port), octave=octave)
+        credentials = None
+        if self.cloud:
+            credentials = create_credentials()
+        self.manager = QuantumMachinesManager(
+            host=host, port=int(port), octave=octave, credentials=credentials
+        )
 
     def setup(self):
         """Deprecated method."""
@@ -232,31 +258,23 @@ def execute_program(self, program):
 
         Args:
             program: QUA program.
-
-        Returns:
-            TODO
         """
         machine = self.manager.open_qm(self.config.__dict__)
         return machine.execute(program)
 
-    @staticmethod
-    def fetch_results(result, acquisitions):
-        """Fetches results from an executed experiment.
+    def simulate_program(self, program):
+        """Simulates an arbitrary program written in QUA language.
 
-        Defined as ``@staticmethod`` because it is overwritten
-        in :class:`qibolab.instruments.qm.simulator.QMSim`.
+        Args:
+            program: QUA program.
         """
-        # TODO: Update result asynchronously instead of waiting
-        # for all values, in order to allow live plotting
-        # using ``handles.is_processing()``
-        handles = result.result_handles
-        handles.wait_for_all_values()
-        results = {}
-        for acquisition in acquisitions.values():
-            data = acquisition.fetch(handles)
-            for serial, result in zip(acquisition.keys, data):
-                results[acquisition.qubit] = results[serial] = result
-        return results
+        ncontrollers = len(self.config.controllers)
+        controller_connections = create_simulator_controller_connections(ncontrollers)
+        simulation_config = SimulationConfig(
+            duration=self.simulation_duration,
+            controller_connections=controller_connections,
+        )
+        return self.manager.simulate(self.config.__dict__, program, simulation_config)
 
     def create_sequence(self, qubits, sequence, sweepers):
         """Translates a :class:`qibolab.pulses.PulseSequence` to a
@@ -276,17 +294,6 @@ def create_sequence(self, qubits, sequence, sweepers):
 
         duration_sweep_pulses = find_duration_sweeper_pulses(sweepers)
 
-        qmpulses = {}
-
-        def create_qmpulse(pulse, qmpulse_cls):
-            qmpulse = qmpulse_cls(pulse)
-            key = (qmpulse.operation, pulse.qubit)
-            if key not in qmpulses:
-                qmpulses[key] = qmpulse
-            else:
-                qmpulses[key].pulses.append(pulse)
-            return qmpulses[key]
-
         qmsequence = Sequence()
         ro_pulses = []
         sort_key = lambda pulse: (pulse.start, pulse.duration)
@@ -357,8 +364,15 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
             with open(self.script_file_name, "w") as file:
                 file.write(generate_qua_script(experiment, self.config.__dict__))
 
-        result = self.execute_program(experiment)
-        return self.fetch_results(result, acquisitions)
+        if self.simulation_duration is not None:
+            result = self.simulate_program(experiment)
+            results = {}
+            for pulse in ro_pulses:
+                results[pulse.qubit] = results[pulse.serial] = result
+            return results
+        else:
+            result = self.execute_program(experiment)
+            return fetch_results(result, acquisitions)
 
     def split_batches(self, sequences):
         return [sequences]
diff --git a/src/qibolab/instruments/qm/simulator.py b/src/qibolab/instruments/qm/simulator.py
deleted file mode 100644
index 2542aaa72..000000000
--- a/src/qibolab/instruments/qm/simulator.py
+++ /dev/null
@@ -1,55 +0,0 @@
-from dataclasses import dataclass
-
-from qm import SimulationConfig
-from qm.QuantumMachinesManager import QuantumMachinesManager
-from qualang_tools.simulator_tools import create_simulator_controller_connections
-
-from .controller import QMController
-
-
-@dataclass
-class QMSim(QMController):
-    """Instrument for using the Quantum Machines (QM) OPX simulator.
-
-    Args:
-        address (str): Address and port of the simulator.
-        simulation_duration (int): Duration for the simulation in ns.
-        cloud (bool): If ``True`` the QM cloud simulator is used which does not
-            require access to physical instruments. This assumes that a proper
-            cloud address has been given.
-            If ``False`` the simulator built-in the instruments is used.
-            This requires connection to instruments.
-            Default is ``False``.
-    """
-
-    simulation_duration: int = 1000
-    cloud: bool = False
-
-    def __post_init__(self):
-        """Convert simulation duration from ns to clock cycles."""
-        self.simulation_duration //= 4
-        self.settings = None
-
-    def connect(self):
-        host, port = self.address.split(":")
-        if self.cloud:
-            from qm.simulate.credentials import create_credentials
-
-            self.manager = QuantumMachinesManager(
-                host, int(port), credentials=create_credentials()
-            )
-        else:
-            self.manager = QuantumMachinesManager(host, int(port))
-
-    @staticmethod
-    def fetch_results(result, acquisitions):
-        return result
-
-    def execute_program(self, program):
-        ncontrollers = len(self.config.controllers)
-        controller_connections = create_simulator_controller_connections(ncontrollers)
-        simulation_config = SimulationConfig(
-            duration=self.simulation_duration,
-            controller_connections=controller_connections,
-        )
-        return self.manager.simulate(self.config.__dict__, program, simulation_config)
diff --git a/src/qibolab/platform.py b/src/qibolab/platform.py
index 0848792b1..68d29e295 100644
--- a/src/qibolab/platform.py
+++ b/src/qibolab/platform.py
@@ -210,9 +210,6 @@ def _execute(self, sequence, options, **kwargs):
                 )
                 if isinstance(new_result, dict):
                     result.update(new_result)
-                elif new_result is not None:
-                    # currently the result of QMSim is not a dict
-                    result = new_result
 
         return result
 
@@ -283,11 +280,6 @@ def execute_pulse_sequences(
         results = defaultdict(list)
         for batch in self._controller.split_batches(sequences):
             sequence, readouts = unroll_sequences(batch, options.relaxation_time)
-
-            print()
-            print(sequence)
-            print()
-
             result = self._execute(sequence, options, **kwargs)
             for serial, new_serials in readouts.items():
                 results[serial].extend(result[ser] for ser in new_serials)
@@ -348,10 +340,6 @@ def sweep(
                 )
                 if isinstance(new_result, dict):
                     result.update(new_result)
-                elif new_result is not None:
-                    # currently the result of QMSim is not a dict
-                    result = new_result
-
         return result
 
     def __call__(self, sequence, options):

From b4841f75ef3f937c4ee8abc4f103fe5fdb9d3e09 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 9 Jan 2024 17:38:52 +0400
Subject: [PATCH 04/19] fix: Fix simulator results

---
 src/qibolab/instruments/qm/controller.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index 61ddcbb20..9ec2290ea 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -367,7 +367,8 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
         if self.simulation_duration is not None:
             result = self.simulate_program(experiment)
             results = {}
-            for pulse in ro_pulses:
+            for qmpulse in ro_pulses:
+                pulse = qmpulse.pulse
                 results[pulse.qubit] = results[pulse.serial] = result
             return results
         else:

From b935dfbc0541e2befae672ee45f7f8cab0c252fe Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 9 Jan 2024 18:05:27 +0400
Subject: [PATCH 05/19] test: Update tests

---
 src/qibolab/instruments/qm/acquisition.py |  2 +-
 tests/test_instruments_qm.py              | 31 ++++++++++-------------
 tests/test_instruments_qmsim.py           |  8 +++---
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index c9efd1c08..38b6c46f5 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -236,7 +236,7 @@ def download(self, *dimensions):
 
     def fetch(self, handles):
         shots = handles.get(f"{self.name}_shots").fetch_all()
-        if len(self.npulses) > 1:
+        if self.npulses > 1:
             if self.average:
                 # TODO: calculate std
                 return [
diff --git a/tests/test_instruments_qm.py b/tests/test_instruments_qm.py
index f433cc437..9d6055fd8 100644
--- a/tests/test_instruments_qm.py
+++ b/tests/test_instruments_qm.py
@@ -6,10 +6,11 @@
 
 from qibolab import AcquisitionType, ExecutionParameters, create_platform
 from qibolab.instruments.qm import OPXplus, QMController
-from qibolab.instruments.qm.acquisition import Acquisition
+from qibolab.instruments.qm.acquisition import Acquisition, declare_acquisitions
 from qibolab.instruments.qm.controller import controllers_config
 from qibolab.instruments.qm.sequence import BakedPulse, QMPulse, Sequence
 from qibolab.pulses import FluxPulse, Pulse, PulseSequence, ReadoutPulse, Rectangular
+from qibolab.qubits import Qubit
 from qibolab.sweeper import Parameter, Sweeper
 
 from .conftest import set_platform_profile
@@ -18,7 +19,7 @@
 def test_qmpulse():
     pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
     qmpulse = QMPulse(pulse)
-    assert qmpulse.operation == pulse.serial
+    assert qmpulse.operation == "drive(40, 0.05, Rectangular())"
     assert qmpulse.relative_phase == 0
 
 
@@ -27,13 +28,14 @@ def test_qmpulse_declare_output(acquisition_type):
     options = ExecutionParameters(acquisition_type=acquisition_type)
     pulse = Pulse(0, 40, 0.05, int(3e9), 0.0, Rectangular(), "ch0", qubit=0)
     qmpulse = QMPulse(pulse)
+    qubits = {0: Qubit(0, threshold=0.1, iq_angle=0.2)}
     if acquisition_type is AcquisitionType.SPECTROSCOPY:
-        with pytest.raises(ValueError):
+        with pytest.raises(KeyError):
             with qua.program() as _:
-                qmpulse.declare_output(options, 0.1, 0.2)
+                declare_acquisitions([qmpulse], qubits, options)
     else:
         with qua.program() as _:
-            qmpulse.declare_output(options, 0.1, 0.2)
+            declare_acquisitions([qmpulse], qubits, options)
         acquisition = qmpulse.acquisition
         assert isinstance(acquisition, Acquisition)
         if acquisition_type is AcquisitionType.DISCRIMINATION:
@@ -61,7 +63,6 @@ def test_qmsequence():
     qmsequence.add(QMPulse(ro_pulse))
     assert len(qmsequence.pulse_to_qmpulse) == 2
     assert len(qmsequence.qmpulses) == 2
-    assert len(qmsequence.ro_pulses) == 1
 
 
 def test_qmpulse_previous_and_next():
@@ -322,14 +323,11 @@ def test_qm_register_pulse(qmplatform, pulse_type, qubit):
     controller.config.register_element(
         platform.qubits[qubit], pulse, controller.time_of_flight, controller.smearing
     )
-    controller.config.register_pulse(platform.qubits[qubit], pulse)
-    assert controller.config.pulses[pulse.serial] == target_pulse
+    qmpulse = QMPulse(pulse)
+    controller.config.register_pulse(platform.qubits[qubit], qmpulse)
+    assert controller.config.pulses[qmpulse.operation] == target_pulse
     assert target_pulse["waveforms"]["I"] in controller.config.waveforms
     assert target_pulse["waveforms"]["Q"] in controller.config.waveforms
-    assert (
-        controller.config.elements[f"{pulse_type}{qubit}"]["operations"][pulse.serial]
-        == pulse.serial
-    )
 
 
 def test_qm_register_flux_pulse(qmplatform):
@@ -344,14 +342,11 @@ def test_qm_register_flux_pulse(qmplatform):
         "length": pulse.duration,
         "waveforms": {"single": "constant_wf0.005"},
     }
+    qmpulse = QMPulse(pulse)
     controller.config.register_element(platform.qubits[qubit], pulse)
-    controller.config.register_pulse(platform.qubits[qubit], pulse)
-    assert controller.config.pulses[pulse.serial] == target_pulse
+    controller.config.register_pulse(platform.qubits[qubit], qmpulse)
+    assert controller.config.pulses[qmpulse.operation] == target_pulse
     assert target_pulse["waveforms"]["single"] in controller.config.waveforms
-    assert (
-        controller.config.elements[f"flux{qubit}"]["operations"][pulse.serial]
-        == pulse.serial
-    )
 
 
 @pytest.mark.parametrize("duration", [0, 30])
diff --git a/tests/test_instruments_qmsim.py b/tests/test_instruments_qmsim.py
index 27940ba23..74ab6ef25 100644
--- a/tests/test_instruments_qmsim.py
+++ b/tests/test_instruments_qmsim.py
@@ -22,7 +22,6 @@
 
 from qibolab import AcquisitionType, AveragingMode, ExecutionParameters, create_platform
 from qibolab.backends import QibolabBackend
-from qibolab.instruments.qm import QMSim
 from qibolab.pulses import SNZ, FluxPulse, PulseSequence, Rectangular
 from qibolab.sweeper import Parameter, Sweeper
 
@@ -43,10 +42,11 @@ def simulator(request):
         pytest.skip("Skipping QM simulator tests because address was not provided.")
 
     platform = create_platform("qm")
-    duration = request.config.getoption("--simulation-duration")
-    controller = QMSim("qmopx", address, simulation_duration=duration, cloud=True)
+    controller = platform.instruments["qm"]
+    controller.simulation_duration = request.config.getoption("--simulation-duration")
     controller.time_of_flight = 280
-    platform.instruments["qmopx"] = controller
+    # controller.cloud = True
+
     platform.connect()
     platform.setup()
     yield platform

From 7d40211a9e0d1287d7444967741b547650ed3ff5 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Wed, 10 Jan 2024 17:00:40 +0400
Subject: [PATCH 06/19] Fix bug with singleshot discrimination

---
 src/qibolab/instruments/qm/acquisition.py | 2 +-
 tests/test_result_shapes.py               | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 38b6c46f5..ce9a6069d 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -281,7 +281,7 @@ def declare_acquisitions(ro_pulses, qubits, options):
             iq_angle = qubits[qubit].iq_angle
             average = options.averaging_mode is AveragingMode.CYCLIC
             acquisition_cls = ACQUISITION_TYPES[options.acquisition_type]
-            acquisition = acquisition_cls(name, qubit, options, threshold, iq_angle)
+            acquisition = acquisition_cls(name, qubit, average, threshold, iq_angle)
             acquisition.assign_element(qmpulse.element)
             acquisitions[name] = acquisition
 
diff --git a/tests/test_result_shapes.py b/tests/test_result_shapes.py
index 6498936bb..9bb149d3e 100644
--- a/tests/test_result_shapes.py
+++ b/tests/test_result_shapes.py
@@ -29,8 +29,8 @@ def execute(platform, acquisition_type, averaging_mode, sweep=False):
         nshots=NSHOTS, acquisition_type=acquisition_type, averaging_mode=averaging_mode
     )
     if sweep:
-        amp_values = np.linspace(0.01, 0.05, NSWEEP1)
-        freq_values = np.linspace(-2e6, 2e6, NSWEEP2)
+        amp_values = np.arange(0.01, 0.06, 0.01)
+        freq_values = np.arange(-4e6, 4e6, 1e6)
         sweeper1 = Sweeper(Parameter.bias, amp_values, qubits=[platform.qubits[qubit]])
         # sweeper1 = Sweeper(Parameter.amplitude, amp_values, pulses=[qd_pulse])
         sweeper2 = Sweeper(Parameter.frequency, freq_values, pulses=[ro_pulse])

From f3aac36ab114c63f5ff58e6c36b309a95e67b9d7 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:58:02 +0400
Subject: [PATCH 07/19] refactor: Move threshold and angle to ShotsAcquisition

---
 src/qibolab/instruments/qm/acquisition.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index db5abc09d..3b112689b 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -35,10 +35,6 @@ class Acquisition(ABC):
     instruments."""
     qubit: QubitId
     average: bool
-    threshold: Optional[float] = None
-    """Threshold to be used for classification of single shots."""
-    angle: Optional[float] = None
-    """Angle in the IQ plane to be used for classification of single shots."""
 
     keys: List[str] = field(default_factory=list)
 
@@ -188,6 +184,11 @@ class ShotsAcquisition(Acquisition):
     Threshold and angle must be given in order to classify shots.
     """
 
+    threshold: Optional[float] = None
+    """Threshold to be used for classification of single shots."""
+    angle: Optional[float] = None
+    """Angle in the IQ plane to be used for classification of single shots."""
+
     I: _Variable = field(default_factory=lambda: declare(fixed))
     Q: _Variable = field(default_factory=lambda: declare(fixed))
     """Variables to save the (I, Q) values acquired from a single shot."""
@@ -277,11 +278,17 @@ def declare_acquisitions(ro_pulses, qubits, options):
         qubit = qmpulse.pulse.qubit
         name = f"{qmpulse.operation}_{qubit}"
         if name not in acquisitions:
-            threshold = qubits[qubit].threshold
-            iq_angle = qubits[qubit].iq_angle
             average = options.averaging_mode is AveragingMode.CYCLIC
             acquisition_cls = ACQUISITION_TYPES[options.acquisition_type]
-            acquisition = acquisition_cls(name, qubit, average, threshold, iq_angle)
+            if options.acquisition_type is AcquisitionType.DISCRIMINATION:
+                threshold = qubits[qubit].threshold
+                iq_angle = qubits[qubit].iq_angle
+                acquisition = acquisition_cls(
+                    name, qubit, average, threshold=threshold, angle=iq_angle
+                )
+            else:
+                acquisition = acquisition_cls(name, qubit, average)
+
             acquisition.assign_element(qmpulse.element)
             acquisitions[name] = acquisition
 

From 555b12fa4c4dde3723b799211dc04de9fea94c4c Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:59:02 +0400
Subject: [PATCH 08/19] refactor: Drop error

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

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 3b112689b..648237656 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -198,11 +198,6 @@ class ShotsAcquisition(Acquisition):
     """Stream to collect multiple shots."""
 
     def __post_init__(self):
-        if self.threshold is None or self.angle is None:
-            raise ValueError(
-                "Cannot use ``AcquisitionType.DISCRIMINATION`` "
-                "if threshold and angle are not given."
-            )
         self.cos = np.cos(self.angle)
         self.sin = np.sin(self.angle)
 

From 50af0a8f202c1e8666b806560989e33ffc5b24d6 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 5 Feb 2024 18:03:31 +0400
Subject: [PATCH 09/19] refactor: Merge acquisition measure and save

---
 src/qibolab/instruments/qm/acquisition.py | 11 -----------
 src/qibolab/instruments/qm/sequence.py    |  2 --
 2 files changed, 13 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 648237656..0a4fce600 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -61,10 +61,6 @@ def measure(self, operation, element):
             element (str): Element (from ``config``) that the pulse will be applied on.
         """
 
-    @abstractmethod
-    def save(self):
-        """Save acquired results from variables to streams."""
-
     @abstractmethod
     def download(self, *dimensions):
         """Save streams to prepare for fetching from host device.
@@ -93,9 +89,6 @@ def assign_element(self, element):
     def measure(self, operation, element):
         qua.measure(operation, element, self.adc_stream)
 
-    def save(self):
-        pass
-
     def download(self, *dimensions):
         i_stream = self.adc_stream.input1()
         q_stream = self.adc_stream.input2()
@@ -138,8 +131,6 @@ def measure(self, operation, element):
             qua.dual_demod.full("cos", "out1", "sin", "out2", self.I),
             qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.Q),
         )
-
-    def save(self):
         qua.save(self.I, self.I_stream)
         qua.save(self.Q, self.Q_stream)
 
@@ -216,8 +207,6 @@ def measure(self, operation, element):
             self.shot,
             qua.Cast.to_int(self.I * self.cos - self.Q * self.sin > self.threshold),
         )
-
-    def save(self):
         qua.save(self.shot, self.shots)
 
     def download(self, *dimensions):
diff --git a/src/qibolab/instruments/qm/sequence.py b/src/qibolab/instruments/qm/sequence.py
index d7127f788..21b2ded97 100644
--- a/src/qibolab/instruments/qm/sequence.py
+++ b/src/qibolab/instruments/qm/sequence.py
@@ -259,9 +259,7 @@ def play(self, relaxation_time=0):
             if qmpulse.wait_cycles is not None:
                 qua.wait(qmpulse.wait_cycles, qmpulse.element)
             if pulse.type is PulseType.READOUT:
-                # Save data to the stream processing
                 qmpulse.acquisition.measure(qmpulse.operation, qmpulse.element)
-                qmpulse.acquisition.save()
             else:
                 if (
                     not isinstance(qmpulse.relative_phase, float)

From 77d2d7111ddc40a09ee31adb501bc84dc5319054 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 5 Feb 2024 18:11:02 +0400
Subject: [PATCH 10/19] test: remove calibration path from testing platform

---
 tests/dummy_qrc/qm_octave/platform.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/dummy_qrc/qm_octave/platform.py b/tests/dummy_qrc/qm_octave/platform.py
index 19a04a019..3d3d307bf 100644
--- a/tests/dummy_qrc/qm_octave/platform.py
+++ b/tests/dummy_qrc/qm_octave/platform.py
@@ -31,7 +31,6 @@ def create(runcard_path=RUNCARD):
         opxs=opxs,
         octaves=[octave1, octave2, octave3],
         time_of_flight=280,
-        calibration_path=str(runcard_path.parent),
     )
 
     # Create channel objects and map controllers to channels

From 7af1dee3a53d54cc32f00612e2631d9e51ac277c Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Thu, 15 Feb 2024 18:22:15 +0400
Subject: [PATCH 11/19] Set batch size for unrolling

---
 src/qibolab/instruments/qm/controller.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index 13ee9dde6..39d807ba3 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -9,6 +9,7 @@
 
 from qibolab import AveragingMode
 from qibolab.instruments.abstract import Controller
+from qibolab.instruments.unrolling import batch_max_sequences
 from qibolab.pulses import PulseType
 from qibolab.sweeper import Parameter
 
@@ -22,6 +23,7 @@
 OCTAVE_ADDRESS_OFFSET = 11000
 """Offset to be added to Octave addresses, because they must be 11xxx, where
 xxx are the last three digits of the Octave IP address."""
+MAX_BATCH_SIZE = 30
 
 
 def declare_octaves(octaves, host, calibration_path=None):
@@ -363,4 +365,4 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
             return fetch_results(result, acquisitions)
 
     def split_batches(self, sequences):
-        return [sequences]
+        return batch_max_sequences(sequences, MAX_BATCH_SIZE)

From ac31d56ec09c1ca711191ffb9bc1a685acb27a79 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 17:25:30 +0400
Subject: [PATCH 12/19] fix: tests

---
 tests/test_instruments_qm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_instruments_qm.py b/tests/test_instruments_qm.py
index 32e756159..f7970ccae 100644
--- a/tests/test_instruments_qm.py
+++ b/tests/test_instruments_qm.py
@@ -307,7 +307,7 @@ def test_qm_register_pulse(qmplatform, pulse_type, qubit):
         target_pulse = {
             "operation": "measurement",
             "length": pulse.duration,
-            "waveforms": {"I": "constant_wf0.003575", "Q": "constant_wf0.003575"},
+            "waveforms": {"I": "constant_wf0.003575", "Q": "zero_wf"},
             "digital_marker": "ON",
             "integration_weights": {
                 "cos": "cosine_weights1",

From eb853ecdfbb9b355af6a5c7f58714a155241306c Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 17:33:42 +0400
Subject: [PATCH 13/19] chore: naming conventions

---
 src/qibolab/instruments/qm/acquisition.py | 66 +++++++++++------------
 tests/test_instruments_qm.py              |  4 +-
 2 files changed, 35 insertions(+), 35 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 0a4fce600..e9be966b0 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import Optional
 
 import numpy as np
 from qm import qua
@@ -36,7 +36,7 @@ class Acquisition(ABC):
     qubit: QubitId
     average: bool
 
-    keys: List[str] = field(default_factory=list)
+    keys: list[str] = field(default_factory=list)
 
     @property
     def npulses(self):
@@ -90,13 +90,13 @@ def measure(self, operation, element):
         qua.measure(operation, element, self.adc_stream)
 
     def download(self, *dimensions):
-        i_stream = self.adc_stream.input1()
-        q_stream = self.adc_stream.input2()
+        istream = self.adc_stream.input1()
+        qstream = self.adc_stream.input2()
         if self.average:
-            i_stream = i_stream.average()
-            q_stream = q_stream.average()
-        i_stream.save(f"{self.name}_I")
-        q_stream.save(f"{self.name}_Q")
+            istream = istream.average()
+            qstream = qstream.average()
+        istream.save(f"{self.name}_I")
+        qstream.save(f"{self.name}_Q")
 
     def fetch(self, handles):
         ires = handles.get(f"{self.name}_I").fetch_all()
@@ -113,41 +113,41 @@ def fetch(self, handles):
 class IntegratedAcquisition(Acquisition):
     """QUA variables used for integrated acquisition."""
 
-    I: _Variable = field(default_factory=lambda: declare(fixed))
-    Q: _Variable = field(default_factory=lambda: declare(fixed))
+    i: _Variable = field(default_factory=lambda: declare(fixed))
+    q: _Variable = field(default_factory=lambda: declare(fixed))
     """Variables to save the (I, Q) values acquired from a single shot."""
-    I_stream: _ResultSource = field(default_factory=lambda: declare_stream())
-    Q_stream: _ResultSource = field(default_factory=lambda: declare_stream())
+    istream: _ResultSource = field(default_factory=lambda: declare_stream())
+    qstream: _ResultSource = field(default_factory=lambda: declare_stream())
     """Streams to collect the results of all shots."""
 
     def assign_element(self, element):
-        assign_variables_to_element(element, self.I, self.Q)
+        assign_variables_to_element(element, self.i, self.q)
 
     def measure(self, operation, element):
         qua.measure(
             operation,
             element,
             None,
-            qua.dual_demod.full("cos", "out1", "sin", "out2", self.I),
-            qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.Q),
+            qua.dual_demod.full("cos", "out1", "sin", "out2", self.i),
+            qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.q),
         )
-        qua.save(self.I, self.I_stream)
-        qua.save(self.Q, self.Q_stream)
+        qua.save(self.I, self.istream)
+        qua.save(self.Q, self.qstream)
 
     def download(self, *dimensions):
-        Istream = self.I_stream
-        Qstream = self.Q_stream
+        istream = self.istream
+        qstream = self.qstream
         if self.npulses > 1:
-            Istream = Istream.buffer(self.npulses)
-            Qstream = Qstream.buffer(self.npulses)
+            istream = istream.buffer(self.npulses)
+            qstream = qstream.buffer(self.npulses)
         for dim in dimensions:
-            Istream = Istream.buffer(dim)
-            Qstream = Qstream.buffer(dim)
+            istream = istream.buffer(dim)
+            qstream = qstream.buffer(dim)
         if self.average:
-            Istream = Istream.average()
-            Qstream = Qstream.average()
-        Istream.save(f"{self.name}_I")
-        Qstream.save(f"{self.name}_Q")
+            istream = istream.average()
+            qstream = qstream.average()
+        istream.save(f"{self.name}_I")
+        qstream.save(f"{self.name}_Q")
 
     def fetch(self, handles):
         ires = handles.get(f"{self.name}_I").fetch_all()
@@ -180,8 +180,8 @@ class ShotsAcquisition(Acquisition):
     angle: Optional[float] = None
     """Angle in the IQ plane to be used for classification of single shots."""
 
-    I: _Variable = field(default_factory=lambda: declare(fixed))
-    Q: _Variable = field(default_factory=lambda: declare(fixed))
+    i: _Variable = field(default_factory=lambda: declare(fixed))
+    q: _Variable = field(default_factory=lambda: declare(fixed))
     """Variables to save the (I, Q) values acquired from a single shot."""
     shot: _Variable = field(default_factory=lambda: declare(int))
     """Variable for calculating an individual shots."""
@@ -193,19 +193,19 @@ def __post_init__(self):
         self.sin = np.sin(self.angle)
 
     def assign_element(self, element):
-        assign_variables_to_element(element, self.I, self.Q, self.shot)
+        assign_variables_to_element(element, self.i, self.q, self.shot)
 
     def measure(self, operation, element):
         qua.measure(
             operation,
             element,
             None,
-            qua.dual_demod.full("cos", "out1", "sin", "out2", self.I),
-            qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.Q),
+            qua.dual_demod.full("cos", "out1", "sin", "out2", self.i),
+            qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.q),
         )
         qua.assign(
             self.shot,
-            qua.Cast.to_int(self.I * self.cos - self.Q * self.sin > self.threshold),
+            qua.Cast.to_int(self.i * self.cos - self.q * self.sin > self.threshold),
         )
         qua.save(self.shot, self.shots)
 
diff --git a/tests/test_instruments_qm.py b/tests/test_instruments_qm.py
index f7970ccae..ce7d1fbf6 100644
--- a/tests/test_instruments_qm.py
+++ b/tests/test_instruments_qm.py
@@ -47,8 +47,8 @@ def test_qmpulse_declare_output(acquisition_type):
         elif acquisition_type is AcquisitionType.INTEGRATION:
             assert isinstance(acquisition.I, qua._dsl._Variable)
             assert isinstance(acquisition.Q, qua._dsl._Variable)
-            assert isinstance(acquisition.I_stream, qua._dsl._ResultSource)
-            assert isinstance(acquisition.Q_stream, qua._dsl._ResultSource)
+            assert isinstance(acquisition.istream, qua._dsl._ResultSource)
+            assert isinstance(acquisition.qstream, qua._dsl._ResultSource)
         elif acquisition_type is AcquisitionType.RAW:
             assert isinstance(acquisition.adc_stream, qua._dsl._ResultSource)
 

From d5b8b45738b829e76f26d798d91dbf75a51bc6d5 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 18:09:19 +0400
Subject: [PATCH 14/19] refactor: simplify fetching

---
 src/qibolab/instruments/qm/acquisition.py | 35 ++++++-----------------
 src/qibolab/result.py                     |  2 +-
 2 files changed, 9 insertions(+), 28 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index e9be966b0..c79e5a31c 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -131,8 +131,8 @@ def measure(self, operation, element):
             qua.dual_demod.full("cos", "out1", "sin", "out2", self.i),
             qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.q),
         )
-        qua.save(self.I, self.istream)
-        qua.save(self.Q, self.qstream)
+        qua.save(self.i, self.istream)
+        qua.save(self.q, self.qstream)
 
     def download(self, *dimensions):
         istream = self.istream
@@ -153,19 +153,10 @@ def fetch(self, handles):
         ires = handles.get(f"{self.name}_I").fetch_all()
         qres = handles.get(f"{self.name}_Q").fetch_all()
         signal = ires + 1j * qres
+        res_cls = AveragedIntegratedResults if self.average else IntegratedResults
         if self.npulses > 1:
-            if self.average:
-                # TODO: calculate std
-                return [
-                    AveragedIntegratedResults(signal[..., i])
-                    for i in range(self.npulses)
-                ]
-            return [IntegratedResults(signal[..., i]) for i in range(self.npulses)]
-        else:
-            if self.average:
-                # TODO: calculate std
-                return [AveragedIntegratedResults(signal)]
-            return [IntegratedResults(signal)]
+            return [res_cls(signal[..., i]) for i in range(self.npulses)]
+        return [res_cls(signal)]
 
 
 @dataclass
@@ -221,20 +212,10 @@ def download(self, *dimensions):
 
     def fetch(self, handles):
         shots = handles.get(f"{self.name}_shots").fetch_all()
+        res_cls = AveragedSampleResults if self.average else SampleResults
         if self.npulses > 1:
-            if self.average:
-                # TODO: calculate std
-                return [
-                    AveragedSampleResults(shots[..., i]) for i in range(self.npulses)
-                ]
-            return [
-                SampleResults(shots[..., i].astype(int)) for i in range(self.npulses)
-            ]
-        else:
-            if self.average:
-                # TODO: calculate std
-                return [AveragedSampleResults(shots)]
-            return [SampleResults(shots.astype(int))]
+            return [res_cls(shots[..., i]) for i in range(self.npulses)]
+        return [res_cls(shots.astype(int))]
 
 
 ACQUISITION_TYPES = {
diff --git a/src/qibolab/result.py b/src/qibolab/result.py
index 1ee07539a..a72757898 100644
--- a/src/qibolab/result.py
+++ b/src/qibolab/result.py
@@ -108,7 +108,7 @@ class SampleResults:
     """
 
     def __init__(self, data: np.ndarray):
-        self.samples: npt.NDArray[np.uint32] = data
+        self.samples: npt.NDArray[np.uint32] = np.array(data).astype(int)
 
     def __add__(self, data):
         return self.__class__(np.append(self.samples, data.samples))

From 0aaf828df50f48fa390d051bad0616b7c0f4d2bd Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 18:15:12 +0400
Subject: [PATCH 15/19] refactor: simplify declare_acquisition

---
 src/qibolab/instruments/qm/acquisition.py | 14 ++++++--------
 tests/test_instruments_qm.py              |  4 ++--
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index c79e5a31c..896832adf 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -244,15 +244,13 @@ def declare_acquisitions(ro_pulses, qubits, options):
         name = f"{qmpulse.operation}_{qubit}"
         if name not in acquisitions:
             average = options.averaging_mode is AveragingMode.CYCLIC
-            acquisition_cls = ACQUISITION_TYPES[options.acquisition_type]
+            kwargs = {}
             if options.acquisition_type is AcquisitionType.DISCRIMINATION:
-                threshold = qubits[qubit].threshold
-                iq_angle = qubits[qubit].iq_angle
-                acquisition = acquisition_cls(
-                    name, qubit, average, threshold=threshold, angle=iq_angle
-                )
-            else:
-                acquisition = acquisition_cls(name, qubit, average)
+                kwargs["threshold"] = qubits[qubit].threshold
+                kwargs["angle"] = qubits[qubit].iq_angle
+            acquisition = ACQUISITION_TYPES[options.acquisition_type](
+                name, qubit, average, **kwargs
+            )
 
             acquisition.assign_element(qmpulse.element)
             acquisitions[name] = acquisition
diff --git a/tests/test_instruments_qm.py b/tests/test_instruments_qm.py
index ce7d1fbf6..b013801de 100644
--- a/tests/test_instruments_qm.py
+++ b/tests/test_instruments_qm.py
@@ -45,8 +45,8 @@ def test_qmpulse_declare_output(acquisition_type):
             assert isinstance(acquisition.shot, qua._dsl._Variable)
             assert isinstance(acquisition.shots, qua._dsl._ResultSource)
         elif acquisition_type is AcquisitionType.INTEGRATION:
-            assert isinstance(acquisition.I, qua._dsl._Variable)
-            assert isinstance(acquisition.Q, qua._dsl._Variable)
+            assert isinstance(acquisition.i, qua._dsl._Variable)
+            assert isinstance(acquisition.q, qua._dsl._Variable)
             assert isinstance(acquisition.istream, qua._dsl._ResultSource)
             assert isinstance(acquisition.qstream, qua._dsl._ResultSource)
         elif acquisition_type is AcquisitionType.RAW:

From 27893f87d54baf03be93f1d54172609888e877ac Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 23:57:28 +0400
Subject: [PATCH 16/19] refactor: Lift result object creation to base
 Acquisition

---
 src/qibolab/instruments/qm/acquisition.py | 37 +++++++++++++++--------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 896832adf..1b5470966 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -38,6 +38,12 @@ class Acquisition(ABC):
 
     keys: list[str] = field(default_factory=list)
 
+    RESULT_CLS = IntegratedResults
+    """Result object type that corresponds to this acquisition type."""
+    AVERAGED_RESULT_CLS = AveragedIntegratedResults
+    """Averaged result object type that corresponds to this acquisition
+    type."""
+
     @property
     def npulses(self):
         return len(self.keys)
@@ -73,6 +79,13 @@ def download(self, *dimensions):
     def fetch(self):
         """Fetch downloaded streams to host device."""
 
+    def result(self, data):
+        """Creates Qibolab result object that is returned to the platform."""
+        res_cls = self.AVERAGED_RESULT_CLS if self.average else self.RESULT_CLS
+        if self.npulses > 1:
+            return [res_cls(data[..., i]) for i in range(self.npulses)]
+        return [res_cls(data)]
+
 
 @dataclass
 class RawAcquisition(Acquisition):
@@ -83,6 +96,9 @@ class RawAcquisition(Acquisition):
     )
     """Stream to collect raw ADC data."""
 
+    _result_cls = RawWaveformResults
+    _averaged_result_cls = AveragedRawWaveformResults
+
     def assign_element(self, element):
         pass
 
@@ -104,9 +120,7 @@ def fetch(self, handles):
         # convert raw ADC signal to volts
         u = unit()
         signal = u.raw2volts(ires) + 1j * u.raw2volts(qres)
-        if self.average:
-            return [AveragedRawWaveformResults(signal)]
-        return [RawWaveformResults(signal)]
+        return self.result(signal)
 
 
 @dataclass
@@ -120,6 +134,9 @@ class IntegratedAcquisition(Acquisition):
     qstream: _ResultSource = field(default_factory=lambda: declare_stream())
     """Streams to collect the results of all shots."""
 
+    _result_cls = IntegratedResults
+    _averaged_result_cls = AveragedIntegratedResults
+
     def assign_element(self, element):
         assign_variables_to_element(element, self.i, self.q)
 
@@ -152,11 +169,7 @@ def download(self, *dimensions):
     def fetch(self, handles):
         ires = handles.get(f"{self.name}_I").fetch_all()
         qres = handles.get(f"{self.name}_Q").fetch_all()
-        signal = ires + 1j * qres
-        res_cls = AveragedIntegratedResults if self.average else IntegratedResults
-        if self.npulses > 1:
-            return [res_cls(signal[..., i]) for i in range(self.npulses)]
-        return [res_cls(signal)]
+        return self.result(ires + 1j * qres)
 
 
 @dataclass
@@ -179,6 +192,9 @@ class ShotsAcquisition(Acquisition):
     shots: _ResultSource = field(default_factory=lambda: declare_stream())
     """Stream to collect multiple shots."""
 
+    _result_cls = SampleResults
+    _averaged_result_cls = AveragedSampleResults
+
     def __post_init__(self):
         self.cos = np.cos(self.angle)
         self.sin = np.sin(self.angle)
@@ -212,10 +228,7 @@ def download(self, *dimensions):
 
     def fetch(self, handles):
         shots = handles.get(f"{self.name}_shots").fetch_all()
-        res_cls = AveragedSampleResults if self.average else SampleResults
-        if self.npulses > 1:
-            return [res_cls(shots[..., i]) for i in range(self.npulses)]
-        return [res_cls(shots.astype(int))]
+        return self.result(shots)
 
 
 ACQUISITION_TYPES = {

From 2d0420e9b5f2e9fc734f77d69d77d46c5832faeb Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 20 Feb 2024 00:01:39 +0400
Subject: [PATCH 17/19] fix: names of result_cls

---
 src/qibolab/instruments/qm/acquisition.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 1b5470966..4b7c31794 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -96,8 +96,8 @@ class RawAcquisition(Acquisition):
     )
     """Stream to collect raw ADC data."""
 
-    _result_cls = RawWaveformResults
-    _averaged_result_cls = AveragedRawWaveformResults
+    RESULT_CLS = RawWaveformResults
+    AVERAGED_RESULT_CLS = AveragedRawWaveformResults
 
     def assign_element(self, element):
         pass
@@ -134,8 +134,8 @@ class IntegratedAcquisition(Acquisition):
     qstream: _ResultSource = field(default_factory=lambda: declare_stream())
     """Streams to collect the results of all shots."""
 
-    _result_cls = IntegratedResults
-    _averaged_result_cls = AveragedIntegratedResults
+    RESULT_CLS = IntegratedResults
+    AVERAGED_RESULT_CLS = AveragedIntegratedResults
 
     def assign_element(self, element):
         assign_variables_to_element(element, self.i, self.q)
@@ -192,8 +192,8 @@ class ShotsAcquisition(Acquisition):
     shots: _ResultSource = field(default_factory=lambda: declare_stream())
     """Stream to collect multiple shots."""
 
-    _result_cls = SampleResults
-    _averaged_result_cls = AveragedSampleResults
+    RESULT_CLS = SampleResults
+    AVERAGED_RESULT_CLS = AveragedSampleResults
 
     def __post_init__(self):
         self.cos = np.cos(self.angle)

From 5f77a7101d9b0e08d9f4bd2d9fa8756dfd68c72b Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 20 Feb 2024 00:05:01 +0400
Subject: [PATCH 18/19] fix: cast samples to uint32

Co-authored-by: Alessandro Candido <candido.ale@gmail.com>
---
 src/qibolab/result.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibolab/result.py b/src/qibolab/result.py
index a72757898..8124bc9d4 100644
--- a/src/qibolab/result.py
+++ b/src/qibolab/result.py
@@ -108,7 +108,7 @@ class SampleResults:
     """
 
     def __init__(self, data: np.ndarray):
-        self.samples: npt.NDArray[np.uint32] = np.array(data).astype(int)
+        self.samples: npt.NDArray[np.uint32] = np.array(data).astype(np.uint32)
 
     def __add__(self, data):
         return self.__class__(np.append(self.samples, data.samples))

From 12841b3fe2511e5f01c6564a7a89a0240d43e794 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Tue, 20 Feb 2024 14:18:39 +0400
Subject: [PATCH 19/19] Convert acquisitions to list

---
 src/qibolab/instruments/qm/acquisition.py | 8 ++++----
 src/qibolab/instruments/qm/controller.py  | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index 4b7c31794..f3ceafbb2 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -249,7 +249,7 @@ def declare_acquisitions(ro_pulses, qubits, options):
             options containing acquisition type and averaging mode.
 
     Returns:
-        Dictionary containing the different :class:`qibolab.instruments.qm.acquisition.Acquisition` objects.
+        List of all :class:`qibolab.instruments.qm.acquisition.Acquisition` objects.
     """
     acquisitions = {}
     for qmpulse in ro_pulses:
@@ -261,16 +261,16 @@ def declare_acquisitions(ro_pulses, qubits, options):
             if options.acquisition_type is AcquisitionType.DISCRIMINATION:
                 kwargs["threshold"] = qubits[qubit].threshold
                 kwargs["angle"] = qubits[qubit].iq_angle
+
             acquisition = ACQUISITION_TYPES[options.acquisition_type](
                 name, qubit, average, **kwargs
             )
-
             acquisition.assign_element(qmpulse.element)
             acquisitions[name] = acquisition
 
         acquisitions[name].keys.append(qmpulse.pulse.serial)
         qmpulse.acquisition = acquisitions[name]
-    return acquisitions
+    return list(acquisitions.values())
 
 
 def fetch_results(result, acquisitions):
@@ -286,7 +286,7 @@ def fetch_results(result, acquisitions):
     handles = result.result_handles
     handles.wait_for_all_values()  # for async replace with ``handles.is_processing()``
     results = {}
-    for acquisition in acquisitions.values():
+    for acquisition in acquisitions:
         data = acquisition.fetch(handles)
         for serial, result in zip(acquisition.keys, data):
             results[acquisition.qubit] = results[serial] = result
diff --git a/src/qibolab/instruments/qm/controller.py b/src/qibolab/instruments/qm/controller.py
index 114e71367..aea41b4cb 100644
--- a/src/qibolab/instruments/qm/controller.py
+++ b/src/qibolab/instruments/qm/controller.py
@@ -346,7 +346,7 @@ def sweep(self, qubits, couplers, sequence, options, *sweepers):
                 )
 
             with qua.stream_processing():
-                for acquisition in acquisitions.values():
+                for acquisition in acquisitions:
                     acquisition.download(*buffer_dims)
 
         if self.script_file_name is not None: