Merge pull request #738 from qiboteam/qmunrolling

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 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,24 @@ 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)
+
+    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)
+
     @abstractmethod
     def assign_element(self, element):
         """Assign acquisition variables to the corresponding QM controlled.
@@ -50,10 +67,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.
@@ -66,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):
@@ -76,82 +96,80 @@ class RawAcquisition(Acquisition):
     )
     """Stream to collect raw ADC data."""
 
+    RESULT_CLS = RawWaveformResults
+    AVERAGED_RESULT_CLS = AveragedRawWaveformResults
+
     def assign_element(self, element):
         pass
 
     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()
+        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.serial}_I")
-        q_stream.save(f"{self.serial}_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.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)
-        if self.average:
-            return AveragedRawWaveformResults(ires + 1j * qres)
-        return RawWaveformResults(ires + 1j * qres)
+        signal = u.raw2volts(ires) + 1j * u.raw2volts(qres)
+        return self.result(signal)
 
 
 @dataclass
 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."""
 
+    RESULT_CLS = IntegratedResults
+    AVERAGED_RESULT_CLS = AveragedIntegratedResults
+
     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),
         )
-
-    def save(self):
-        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)
         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.serial}_I")
-        Qstream.save(f"{self.serial}_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.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()
+        return self.result(ires + 1j * qres)
 
 
 @dataclass
@@ -161,53 +179,115 @@ 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))
-    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."""
     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)
 
     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),
         )
-
-    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()
+        return self.result(shots)
+
+
+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:
+        List of all :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
+            kwargs = {}
+            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 list(acquisitions.values())
+
+
+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:
+        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}.")