Merge pull request #974 from qiboteam/fix-qm-sequence

Update QM driver for new sequence layout

src/qibolab/compilers/compiler.py

@@ -15,7 +15,6 @@
     rz_rule,
     z_rule,
 )
-from qibolab.components import Channel
 from qibolab.platform import Platform
 from qibolab.pulses import Delay, PulseSequence
 from qibolab.qubits import QubitId
@@ -151,9 +150,6 @@ def compile(
         measurement_map = {}
         channel_clock = defaultdict(float)
 
-        def find_max(channels: list[Channel]):
-            return max(channel_clock[ch.name] for ch in channels)
-
         def qubit_clock(el: QubitId):
             elements = platform.qubits if el in platform.qubits else platform.couplers
             return max(channel_clock[ch.name] for ch in elements[el].channels)

src/qibolab/instruments/qm/controller.py

@@ -8,7 +8,7 @@
 from qm.simulate.credentials import create_credentials
 from qualang_tools.simulator_tools import create_simulator_controller_connections
 
-from qibolab.components import Channel, Config, DcChannel, IqChannel
+from qibolab.components import Channel, ChannelId, Config, DcChannel, IqChannel
 from qibolab.execution_parameters import ExecutionParameters
 from qibolab.instruments.abstract import Controller
 from qibolab.pulses import Delay, Pulse, PulseSequence, VirtualZ
@@ -240,8 +240,19 @@ def configure_channel(self, channel: QmChannel, configs: dict[str, Config]):
         else:
             raise TypeError(f"Unknown channel type: {type(channel)}.")
 
+    def configure_channels(
+        self,
+        configs: dict[str, Config],
+        channels: set[ChannelId],
+    ):
+        """Register channels participating in the sequence in the QM
+        ``config``."""
+        for channel_name in channels:
+            channel = self.channels[channel_name]
+            self.configure_channel(channel, configs)
+
     def register_pulse(self, channel: Channel, pulse: Pulse):
-        """Add pulse in the ``config``."""
+        """Add pulse in the QM ``config``."""
         # if (
         #    pulse.duration % 4 != 0
         #    or pulse.duration < 16
@@ -256,46 +267,45 @@ 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_pulses(self, configs, sequence, integration_setup, options):
-        """Adds all pulses of a given :class:`qibolab.pulses.PulseSequence` in
-        the ``config``.
+    def register_pulses(
+        self,
+        configs: dict[str, Config],
+        sequence: PulseSequence,
+        options: ExecutionParameters,
+    ):
+        """Adds all pulses of a given sequence in the QM ``config``.
 
         Returns:
             acquisitions (dict): Map from measurement instructions to acquisition objects.
         """
         acquisitions = {}
-        for name, channel_sequence in sequence.items():
-            channel = self.channels[name]
-            self.configure_channel(channel, configs)
-
-            for pulse in channel_sequence:
-                if isinstance(pulse, (Delay, VirtualZ)):
-                    continue
-
-                logical_channel = channel.logical_channel
-                op = self.register_pulse(logical_channel, pulse)
-
-                if (
-                    isinstance(logical_channel, IqChannel)
-                    and logical_channel.acquisition is not None
-                ):
-                    kernel, threshold, iq_angle = integration_setup[
-                        logical_channel.acquisition
-                    ]
-                    self.config.register_integration_weights(
-                        name, pulse.duration, kernel
+        for channel_name, pulse in sequence:
+            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]
+                self.config.register_integration_weights(
+                    channel_name, pulse.duration, acq_config.kernel
+                )
+                if (op, channel_name) in acquisitions:
+                    acquisition = acquisitions[(op, channel_name)]
+                else:
+                    acquisition = acquisitions[(op, channel_name)] = create_acquisition(
+                        op,
+                        channel_name,
+                        options,
+                        acq_config.threshold,
+                        acq_config.iq_angle,
                     )
-                    if (op, name) in acquisitions:
-                        acquisition = acquisitions[(op, name)]
-                    else:
-                        acquisition = acquisitions[(op, name)] = create_acquisition(
-                            op,
-                            name,
-                            options,
-                            threshold,
-                            iq_angle,
-                        )
-                    acquisition.keys.append(pulse.id)
+                acquisition.keys.append(pulse.id)
 
         return acquisitions
 
@@ -319,7 +329,6 @@ def play(
         configs: dict[str, Config],
         sequences: list[PulseSequence],
         options: ExecutionParameters,
-        integration_setup,
         sweepers: list[ParallelSweepers],
     ):
         if len(sequences) > 1:
@@ -334,9 +343,8 @@ def play(
                 self.configure_channel(channel, configs)
 
         sequence = sequences[0]
-        acquisitions = self.register_pulses(
-            configs, sequence, integration_setup, options
-        )
+        self.configure_channels(configs, sequence.channels)
+        acquisitions = self.register_pulses(configs, sequence, options)
         experiment = program(configs, sequence, options, acquisitions, sweepers)
 
         if self.manager is None:
@@ -347,19 +355,17 @@ def play(
 
         if self.script_file_name is not None:
             script = generate_qua_script(experiment, asdict(self.config))
-            for pulses in sequence.values():
-                for pulse in pulses:
-                    script = script.replace(operation(pulse), str(pulse))
+            for _, pulse in sequence:
+                script = script.replace(operation(pulse), str(pulse))
             with open(self.script_file_name, "w") as file:
                 file.write(script)
 
         if self.simulation_duration is not None:
             result = self.simulate_program(experiment)
             results = {}
-            for channel_name, pulses in sequence.items():
+            for channel_name, pulse in sequence:
                 if self.channels[channel_name].logical_channel.acquisition is not None:
-                    for pulse in pulses:
-                        results[pulse.id] = result
+                    results[pulse.id] = result
             return results
 
         result = self.execute_program(experiment)

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

@@ -18,7 +18,7 @@
 def _delay(pulse: Delay, element: str, parameters: Parameters):
     # TODO: How to play delays on multiple elements?
     if parameters.duration is None:
-        duration = pulse.duration // 4 + 1
+        duration = int(pulse.duration) // 4 + 1
     else:
         duration = parameters.duration
     qua.wait(duration, element)
@@ -50,15 +50,14 @@ def play(args: ExecutionArguments):
     Should be used inside a ``program()`` context.
     """
     qua.align()
-    for element, pulses in args.sequence.items():
-        for pulse in pulses:
-            op = operation(pulse)
-            params = args.parameters[op]
-            if isinstance(pulse, Delay):
-                _delay(pulse, element, params)
-            else:
-                acquisition = args.acquisitions.get((op, element))
-                _play(op, element, params, acquisition)
+    for element, pulse in args.sequence:
+        op = operation(pulse)
+        params = args.parameters[op]
+        if isinstance(pulse, Delay):
+            _delay(pulse, element, params)
+        else:
+            acquisition = args.acquisitions.get((op, element))
+            _play(op, element, params, acquisition)
 
     if args.relaxation_time > 0:
         qua.wait(args.relaxation_time // 4)