refactor: modify QM duration sweeps to upload multiple waveforms

src/qibolab/instruments/qm/controller.py

@@ -12,12 +12,12 @@
 from qibolab.execution_parameters import ExecutionParameters
 from qibolab.instruments.abstract import Controller
 from qibolab.pulses import Delay, Pulse, PulseSequence, VirtualZ
-from qibolab.sweeper import ParallelSweepers, Parameter
+from qibolab.sweeper import ParallelSweepers, Parameter, Sweeper
 from qibolab.unrolling import Bounds
 
 from .components import QmChannel
 from .config import SAMPLING_RATE, QmConfig, operation
-from .program import create_acquisition, program
+from .program import ExecutionArguments, create_acquisition, program
 
 OCTAVE_ADDRESS_OFFSET = 11000
 """Offset to be added to Octave addresses, because they must be 11xxx, where
@@ -99,6 +99,11 @@ def fetch_results(result, acquisitions):
     }
 
 
+def find_duration_sweepers(sweepers: list[ParallelSweepers]) -> list[Sweeper]:
+    """Find duration sweepers in order to register multiple pulses."""
+    return [s for ps in sweepers for s in ps if s.parameter is Parameter.duration]
+
+
 @dataclass
 class QmController(Controller):
     """:class:`qibolab.instruments.abstract.Controller` object for controlling
@@ -251,8 +256,9 @@ def configure_channels(
             channel = self.channels[channel_name]
             self.configure_channel(channel, configs)
 
-    def register_pulse(self, channel: Channel, pulse: Pulse):
-        """Add pulse in the QM ``config``."""
+    def register_pulse(self, channel: Channel, pulse: Pulse) -> str:
+        """Add pulse in the QM ``config`` and return corresponding
+        operation."""
         # if (
         #    pulse.duration % 4 != 0
         #    or pulse.duration < 16
@@ -267,6 +273,23 @@ def register_pulse(self, channel: Channel, pulse: Pulse):
             return self.config.register_iq_pulse(channel.name, pulse)
         return self.config.register_acquisition_pulse(channel.name, pulse)
 
+    def register_duration_sweeper_pulses(
+        self, args: ExecutionArguments, sweeper: Sweeper
+    ):
+        """Register pulse with many different durations, in order to sweep
+        duration."""
+        for pulse in sweeper.pulses:
+            if isinstance(pulse, Delay):
+                continue
+
+            op = operation(pulse)
+            channel_name = args.sequence.pulse_channel(pulse.id)
+            channel = self.channels[channel_name].logical_channel
+            for value in sweeper.values:
+                sweep_pulse = pulse.model_copy(updates={"duration": value})
+                sweep_op = self.register_pulse(channel, new_pulse)
+                args.parameters[op].pulses.append((value, sweep_op))
+
     def register_pulses(
         self,
         configs: dict[str, Config],
@@ -283,15 +306,11 @@ def register_pulses(
             if isinstance(pulse, (Delay, VirtualZ)):
                 continue
 
-            channel = self.channels[channel_name]
-            logical_channel = channel.logical_channel
-            op = self.register_pulse(logical_channel, pulse)
-
-            if (
-                isinstance(logical_channel, IqChannel)
-                and logical_channel.acquisition is not None
-            ):
-                acq_config = configs[logical_channel.acquisition]
+            channel = self.channels[channel_name].logical_channel
+            op = self.register_pulse(channel, pulse)
+
+            if isinstance(channel, IqChannel) and channel.acquisition is not None:
+                acq_config = configs[channel.acquisition]
                 self.config.register_integration_weights(
                     channel_name, pulse.duration, acq_config.kernel
                 )
@@ -345,7 +364,13 @@ def play(
         sequence = sequences[0]
         self.configure_channels(configs, sequence.channels)
         acquisitions = self.register_pulses(configs, sequence, options)
-        experiment = program(configs, sequence, options, acquisitions, sweepers)
+
+        args = ExecutionArguments(sequence, acquisitions, options.relaxation_time)
+
+        for sweeper in find_duration_sweepers(sweepers):
+            self.register_duration_sweeper_pulses(args, sweeper)
+
+        experiment = program(configs, args, options, sweepers)
 
         if self.manager is None:
             warnings.warn(

src/qibolab/instruments/qm/program/__init__.py

@@ -1,2 +1,3 @@
 from .acquisition import Acquisitions, create_acquisition
+from .arguments import ExecutionArguments
 from .instructions import program

src/qibolab/instruments/qm/program/arguments.py

@@ -6,7 +6,7 @@
 
 from qibolab.pulses import PulseSequence
 
-from .acquisition import Acquisition
+from .acquisition import Acquisitions
 
 
 @dataclass
@@ -16,6 +16,7 @@ class Parameters:
     duration: Optional[_Variable] = None
     amplitude: Optional[_Variable] = None
     phase: Optional[_Variable] = None
+    pulses: list[str] = field(default_factory=list)
 
 
 @dataclass
@@ -27,7 +28,7 @@ class ExecutionArguments:
     """
 
     sequence: PulseSequence
-    acquisitions: dict[tuple[str, str], Acquisition]
+    acquisitions: Acquisitions
     relaxation_time: int = 0
     parameters: dict[str, Parameters] = field(
         default_factory=lambda: defaultdict(Parameters)

src/qibolab/instruments/qm/program/instructions.py

@@ -6,11 +6,11 @@
 
 from qibolab.components import Config
 from qibolab.execution_parameters import AcquisitionType, ExecutionParameters
-from qibolab.pulses import Delay, PulseSequence
+from qibolab.pulses import Delay
 from qibolab.sweeper import ParallelSweepers
 
 from ..config import operation
-from .acquisition import Acquisition, Acquisitions
+from .acquisition import Acquisition
 from .arguments import ExecutionArguments, Parameters
 from .sweepers import INT_TYPE, NORMALIZERS, SWEEPER_METHODS
 
@@ -24,6 +24,14 @@ def _delay(pulse: Delay, element: str, parameters: Parameters):
     qua.wait(duration, element)
 
 
+def _play_multiple_waveforms(element: str, parameters: Parameters):
+    """Sweeping pulse duration using distinctly uploaded waveforms."""
+    with qua.switch_(parameters.duration):
+        for value, sweep_op in parameters.pulses:
+            with qua.case_(value):
+                qua.play(sweep_op, element)
+
+
 def _play(
     op: str,
     element: str,
@@ -35,10 +43,13 @@ def _play(
     if parameters.amplitude is not None:
         op = op * parameters.amplitude
 
-    if acquisition is not None:
-        acquisition.measure(op)
+    if len(parameters.pulses) > 0:
+        _play_multiple_waveforms(element, parameters)
     else:
-        qua.play(op, element, duration=parameters.duration)
+        if acquisition is not None:
+            acquisition.measure(op)
+        else:
+            qua.play(op, element, duration=parameters.duration)
 
     if parameters.phase is not None:
         qua.reset_frame(element)
@@ -108,25 +119,23 @@ def sweep(
 
 def program(
     configs: dict[str, Config],
-    sequence: PulseSequence,
+    args: ExecutionArguments,
     options: ExecutionParameters,
-    acquisitions: Acquisitions,
     sweepers: list[ParallelSweepers],
 ):
     """QUA program implementing the required experiment."""
     with qua.program() as experiment:
         n = declare(int)
         # declare acquisition variables
-        for acquisition in acquisitions.values():
+        for acquisition in args.acquisitions.values():
             acquisition.declare()
         # execute pulses
-        args = ExecutionArguments(sequence, acquisitions, options.relaxation_time)
         with for_(n, 0, n < options.nshots, n + 1):
             sweep(list(sweepers), configs, args)
         # download acquisitions
         has_iq = options.acquisition_type is AcquisitionType.INTEGRATION
         buffer_dims = options.results_shape(sweepers)[::-1][int(has_iq) :]
         with qua.stream_processing():
-            for acquisition in acquisitions.values():
+            for acquisition in args.acquisitions.values():
                 acquisition.download(*buffer_dims)
     return experiment