Sequence unrolling with QM

src/qibolab/instruments/qm/__init__.py

@@ -1,3 +1,2 @@
 from .controller import QMController
 from .devices import Octave, OPXplus
-from .simulator import QMSim

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,
@@ -27,10 +30,18 @@ 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."""
+    qubit: QubitId
     average: bool
 
+    keys: List[str] = field(default_factory=list)
+
+    @property
+    def npulses(self):
+        return len(self.keys)
+
     @abstractmethod
     def assign_element(self, element):
         """Assign acquisition variables to the corresponding QM controlled.
@@ -50,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.
@@ -82,28 +89,24 @@ 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()
         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
@@ -128,30 +131,41 @@ 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)
 
     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
@@ -161,9 +175,9 @@ class ShotsAcquisition(Acquisition):
     Threshold and angle must be given in order to classify shots.
     """
 
-    threshold: float
+    threshold: Optional[float] = None
     """Threshold to be used for classification of single shots."""
-    angle: float
+    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))
@@ -193,21 +207,95 @@ 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):
         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 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))]
+
+
+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:
+            average = options.averaging_mode is AveragingMode.CYCLIC
+            acquisition_cls = ACQUISITION_TYPES[options.acquisition_type]
+            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
+
+        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.
+    """
+    handles = result.result_handles
+    handles.wait_for_all_values()  # for async replace with ``handles.is_processing()``
+    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

src/qibolab/instruments/qm/config.py

@@ -230,7 +230,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:
@@ -243,23 +243,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": {
@@ -268,14 +269,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": {
@@ -291,8 +292,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}.")