Merge pull request #979 from qiboteam/qm-duration-sweeps

QM duration sweeper using multiple waveforms

src/qibolab/instruments/qm/controller.py

@@ -11,18 +11,18 @@
 from qm.simulate.credentials import create_credentials
 from qualang_tools.simulator_tools import create_simulator_controller_connections
 
-from qibolab.components import AcquireChannel, Config, DcChannel, IqChannel
+from qibolab.components import AcquireChannel, Channel, Config, DcChannel, IqChannel
 from qibolab.execution_parameters import ExecutionParameters
 from qibolab.identifier import ChannelId
 from qibolab.instruments.abstract import Controller
-from qibolab.pulses.pulse import Acquisition, Delay, VirtualZ, _Readout
+from qibolab.pulses.pulse import Acquisition, Align, Delay, Pulse, _Readout
 from qibolab.sequence import PulseSequence
-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
@@ -112,6 +112,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
@@ -291,6 +296,24 @@ def configure_channels(
             channel = self.channels[str(channel_id)]
             self.configure_channel(channel, configs)
 
+    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
+        #    or pulse.id in pulses_to_bake
+        # ):
+        #    qmpulse = BakedPulse(pulse, element)
+        #    qmpulse.bake(self.config, durations=[pulse.duration])
+        # else:
+        name = str(channel.name)
+        if isinstance(channel, DcChannel):
+            return self.config.register_dc_pulse(name, pulse)
+        if channel.acquisition is None:
+            return self.config.register_iq_pulse(name, pulse)
+        return self.config.register_acquisition_pulse(name, pulse)
+
     def register_pulses(self, configs: dict[str, Config], sequence: PulseSequence):
         """Adds all pulses except measurements of a given sequence in the QM
         ``config``.
@@ -299,21 +322,26 @@ def register_pulses(self, configs: dict[str, Config], sequence: PulseSequence):
             acquisitions (dict): Map from measurement instructions to acquisition objects.
         """
         for channel_id, pulse in sequence:
-            if not isinstance(pulse, (Acquisition, Delay, VirtualZ)):
-                name = str(channel_id)
-                channel = self.channels[name].logical_channel
-                # if (
-                #    pulse.duration % 4 != 0
-                #    or pulse.duration < 16
-                #    or pulse.id in pulses_to_bake
-                # ):
-                #    qmpulse = BakedPulse(pulse, element)
-                #    qmpulse.bake(self.config, durations=[pulse.duration])
-                # else:
-                if isinstance(channel, DcChannel):
-                    self.config.register_dc_pulse(name, pulse)
-                elif channel.acquisition is None:
-                    self.config.register_iq_pulse(name, pulse)
+            if isinstance(pulse, Pulse):
+                channel = self.channels[str(channel_id)].logical_channel
+                self.register_pulse(channel, 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, (Align, Delay)):
+                continue
+
+            op = operation(pulse)
+            channel_name = str(args.sequence.pulse_channels(pulse.id)[0])
+            channel = self.channels[channel_name].logical_channel
+            for value in sweeper.values:
+                sweep_pulse = pulse.model_copy(update={"duration": value})
+                sweep_op = self.register_pulse(channel, sweep_pulse)
+                args.parameters[op].pulses.append((value, sweep_op))
 
     def register_acquisitions(
         self,
@@ -395,7 +423,13 @@ def play(
         self.configure_channels(configs, sequence.channels)
         self.register_pulses(configs, sequence)
         acquisitions = self.register_acquisitions(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(
@@ -405,8 +439,6 @@ def play(
 
         if self.script_file_name is not None:
             script = generate_qua_script(experiment, asdict(self.config))
-            for _, pulse in sequence:
-                script = script.replace(operation(pulse), str(pulse))
             with open(self.script_file_name, "w") as file:
                 file.write(script)
 

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.sequence 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[tuple[float, 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,23 +6,30 @@
 
 from qibolab.components import Config
 from qibolab.execution_parameters import AcquisitionType, ExecutionParameters
-from qibolab.pulses import Delay, Pulse
-from qibolab.sequence import PulseSequence
+from qibolab.pulses import Align, Delay, Pulse, VirtualZ
 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
+from .sweepers import INT_TYPE, NORMALIZERS, SWEEPER_METHODS, normalize_phase
 
 
 def _delay(pulse: Delay, element: str, parameters: Parameters):
     # TODO: How to play delays on multiple elements?
     if parameters.duration is None:
-        duration = int(pulse.duration) // 4 + 1
+        duration = int(pulse.duration) // 4
     else:
         duration = parameters.duration
-    qua.wait(duration, element)
+    qua.wait(duration + 1, element)
+
+
+def _play_multiple_waveforms(element: str, parameters: Parameters):
+    """Sweeping pulse duration using distinctly uploaded waveforms."""
+    with qua.switch_(parameters.duration, unsafe=True):
+        for value, sweep_op in parameters.pulses:
+            with qua.case_(value // 4):
+                qua.play(sweep_op, element)
 
 
 def _play(
@@ -36,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)
@@ -51,6 +61,12 @@ def play(args: ExecutionArguments):
     Should be used inside a ``program()`` context.
     """
     qua.align()
+
+    # keep track of ``Align`` command that were already played
+    # because the same ``Align`` will appear on multiple channels
+    # in the sequence
+    processed_aligns = set()
+
     for channel_id, pulse in args.sequence:
         element = str(channel_id)
         op = operation(pulse)
@@ -60,6 +76,12 @@ def play(args: ExecutionArguments):
         elif isinstance(pulse, Pulse):
             acquisition = args.acquisitions.get((op, element))
             _play(op, element, params, acquisition)
+        elif isinstance(pulse, VirtualZ):
+            qua.frame_rotation_2pi(normalize_phase(pulse.phase), element)
+        elif isinstance(pulse, Align) and pulse.id not in processed_aligns:
+            channel_ids = args.sequence.pulse_channels(pulse.id)
+            qua.align(*(str(ch) for ch in channel_ids))
+            processed_aligns.add(pulse.id)
 
     if args.relaxation_time > 0:
         qua.wait(args.relaxation_time // 4)
@@ -110,25 +132,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

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

@@ -150,15 +150,26 @@ def _duration(
         args.parameters[operation(pulse)].duration = variable
 
 
-INT_TYPE = {Parameter.frequency, Parameter.duration}
+def normalize_phase(values):
+    """Normalize phase from [0, 2pi] to [0, 1]."""
+    return values / (2 * np.pi)
+
+
+def normalize_duration(values):
+    """Convert duration from ns to clock cycles (clock cycle = 4ns)."""
+    return (values // 4).astype(int)
+
+
+INT_TYPE = {Parameter.frequency, Parameter.duration, Parameter.duration_interpolated}
 """Sweeper parameters for which we need ``int`` variable type.
 
 The rest parameters need ``fixed`` type.
 """
 
 NORMALIZERS = {
-    Parameter.relative_phase: lambda values: values / (2 * np.pi),
-    Parameter.duration: lambda values: (values // 4).astype(int),
+    Parameter.relative_phase: normalize_phase,
+    Parameter.duration: normalize_duration,
+    Parameter.duration_interpolated: normalize_duration,
 }
 """Functions to normalize sweeper values.
 
@@ -169,6 +180,7 @@ def _duration(
     Parameter.frequency: _frequency,
     Parameter.amplitude: _amplitude,
     Parameter.duration: _duration,
+    Parameter.duration_interpolated: _duration,
     Parameter.relative_phase: _relative_phase,
     Parameter.bias: _bias,
 }

src/qibolab/pulses/__init__.py

@@ -1,2 +1,2 @@
 from .envelope import *
-from .pulse import Acquisition, Delay, Pulse, PulseLike, VirtualZ
+from .pulse import Acquisition, Align, Delay, Pulse, PulseLike, VirtualZ

src/qibolab/pulses/plot.py

@@ -143,6 +143,8 @@ def sequence(ps: PulseSequence, freq: dict[str, float], filename=None):
         import matplotlib.pyplot as plt
         from matplotlib import gridspec
 
+        # compile ``Align`` to delays as it is not supported here
+        ps = ps.align_to_delays()
         num_pulses = len(ps)
         _ = plt.figure(figsize=(14, 2 * num_pulses), dpi=200)
         gs = gridspec.GridSpec(ncols=1, nrows=num_pulses)

src/qibolab/pulses/pulse.py

@@ -138,6 +138,13 @@ def id(self) -> int:
         return self.acquisition.id
 
 
+class Align(_PulseLike):
+    """Brings different channels at the same point in time."""
+
+    kind: Literal["align"] = "align"
+
+
 PulseLike = Annotated[
-    Union[Pulse, Delay, VirtualZ, Acquisition, _Readout], Field(discriminator="kind")
+    Union[Align, Pulse, Delay, VirtualZ, Acquisition, _Readout],
+    Field(discriminator="kind"),
 ]