Added channel filtering to pulse converter

qiskit_dynamics/pulse/pulse_to_signals.py

@@ -14,7 +14,7 @@
 Pulse schedule to Signals converter.
 """
 
-from typing import List
+from typing import List, Optional, Union
 import numpy as np
 
 from qiskit.pulse import (
@@ -38,7 +38,9 @@ class InstructionToSignals:
     the :meth:`get_signals` method on a schedule.
     """
 
-    def __init__(self, dt: float, carriers: List[float] = None):
+    def __init__(
+        self, dt: float, carriers: List[float] = None, channels: Optional[List[str]] = None
+    ):
         """Initialize pulse schedule to signals converter.
 
         Args:
@@ -48,10 +50,35 @@ def __init__(self, dt: float, carriers: List[float] = None):
             carriers: a list of carrier frequencies. If it is not None there
                 must be at least as many carrier frequencies as there are
                 channels in the schedules that will be converted.
+            channels: A list of channels that the :meth:`get_signals` method should return.
+                This argument will cause :meth:`get_signals` to return the signals in the
+                same order as the channels. Channels present in the schedule but absent
+                from channels will not be included in the returned object. If None is given
+                (the default) then all channels present in the pulse schedule are returned.
+
+        Raises:
+            QiskitError: If the number of channels and carriers does not match when both are given.
         """
 
         self._dt = dt
-        self._carriers = carriers
+        self._channels = channels
+        self._carriers = {}
+
+        # If both are given we tie them together in a dict to ensure consistency.
+        if channels is not None and carriers is not None:
+            if len(channels) != len(carriers):
+                raise QiskitError(
+                    "The number of required channels and carries does not match: "
+                    f"len({channels}) != len({carriers})."
+                )
+
+            for idx, chan in enumerate(channels):
+                self._carriers[chan] = carriers[idx]
+
+        # If only carriers is given we map using the index of the carriers.
+        elif channels is None and carriers is not None:
+            for idx, carrier in enumerate(carriers):
+                self._carriers[idx] = carrier
 
     def get_signals(self, schedule: Schedule) -> List[DiscreteSignal]:
         """
@@ -65,14 +92,17 @@ def get_signals(self, schedule: Schedule) -> List[DiscreteSignal]:
             qiskit.QiskitError: if not enough frequencies supplied
         """
 
-        if self._carriers and len(self._carriers) < len(schedule.channels):
-            raise QiskitError("Not enough carrier frequencies supplied.")
+        if self._channels is None:
+            if self._carriers and len(self._carriers) < len(schedule.channels):
+                raise QiskitError("Not enough carrier frequencies supplied.")
 
         signals, phases, frequency_shifts = {}, {}, {}
 
         for idx, chan in enumerate(schedule.channels):
-            if self._carriers:
+            if self._carriers and not self._channels:
                 carrier_freq = self._carriers[idx]
+            elif self._carriers and self._channels:
+                carrier_freq = self._carriers.get(chan.name, 0.0)
             else:
                 carrier_freq = 0.0
 
@@ -129,7 +159,19 @@ def get_signals(self, schedule: Schedule) -> List[DiscreteSignal]:
                     samples=np.zeros(max_duration - sig.duration, dtype=complex),
                 )
 
-        return list(signals.values())
+        # filter the channels
+        if self._channels is None:
+            return list(signals.values())
+
+        return_signals = []
+        for chan_name in self._channels:
+            signal = signals.get(
+                chan_name, DiscreteSignal(samples=[], dt=self._dt, name=chan_name, carrier_freq=0.0)
+            )
+
+            return_signals.append(signal)
+
+        return return_signals
 
     @staticmethod
     def get_awg_signals(

requirements-dev.txt

@@ -10,3 +10,5 @@ reno>=3.4.0
 nbsphinx
 qutip
 matplotlib<3.5
+
+ddt~=1.4.2

test/dynamics/signals/test_pulse_to_signals.py

@@ -13,11 +13,15 @@
 Tests to convert from pulse schedules to signals.
 """
 
+from ddt import ddt, data, unpack
 import numpy as np
 
+import qiskit.pulse as pulse
 from qiskit.pulse import (
     Schedule,
     DriveChannel,
+    ControlChannel,
+    MeasureChannel,
     Play,
     Drag,
     ShiftFrequency,
@@ -28,6 +32,9 @@
     Constant,
     Waveform,
 )
+from qiskit.pulse.transforms.canonicalization import block_to_schedule
+from qiskit.exceptions import QiskitError
+
 from qiskit_dynamics.pulse import InstructionToSignals
 from qiskit_dynamics.signals import DiscreteSignal
 
@@ -134,3 +141,68 @@ def test_uneven_pulse_length(self):
 
         self.assertTrue(signals[0].carrier_freq == 2.0)
         self.assertTrue(signals[1].carrier_freq == 3.0)
+
+
+@ddt
+class TestPulseToSignalsFiltering(QiskitDynamicsTestCase):
+    """Test the extraction of signals when specifying channels."""
+
+    def setUp(self):
+        """Setup the tests."""
+
+        super().setUp()
+
+        # Drags on all qubits, then two CRs, then readout all qubits.
+        with pulse.build(name="test schedule") as schedule:
+            with pulse.align_sequential():
+                with pulse.align_left():
+                    for chan_idx in [0, 1, 2, 3]:
+                        pulse.play(Drag(160, 0.5, 40, 0.1), DriveChannel(chan_idx))
+
+                with pulse.align_sequential():
+                    for chan_idx in [0, 1]:
+                        pulse.play(GaussianSquare(660, 0.2, 40, 500), ControlChannel(chan_idx))
+
+                with pulse.align_left():
+                    for chan_idx in [0, 1, 2, 3]:
+                        pulse.play(GaussianSquare(660, 0.2, 40, 500), MeasureChannel(chan_idx))
+
+        self._schedule = block_to_schedule(schedule)
+
+    @unpack
+    @data(
+        ([5.0, 5.1, 5.0, 5.1], ["d0", "d2", "u0", "u1"]),
+        ([5.0, 5.1, 5.0, 5.1], ["m0", "m1", "m2", "m3"]),
+        ([5.0, 5.1, 5.0, 5.1], ["m0", "m1", "d0", "d1"]),
+        ([5.0], ["d1"]),
+        ([5.0], ["d123"]),
+    )
+    def test_channel_combinations(self, carriers, channels):
+        """Test that we can filter out channels in the right order and number."""
+
+        converter = InstructionToSignals(dt=0.222, carriers=carriers, channels=channels)
+
+        signals = converter.get_signals(self._schedule)
+
+        self.assertEqual(len(signals), len(channels))
+        for idx, chan_name in enumerate(channels):
+            self.assertEqual(signals[idx].name, chan_name)
+
+    def test_empty_signal(self):
+        """Test that requesting a channel that is not in the schedule gives and empty signal."""
+
+        converter = InstructionToSignals(dt=0.222, carriers=[1.0], channels=["d123"])
+
+        signals = converter.get_signals(self._schedule)
+
+        self.assertEqual(len(signals), 1)
+        self.assertEqual(signals[0].duration, 0)
+
+    def test_malformed_input_args(self):
+        """Test that we get errors if the carriers and channels do not match."""
+
+        with self.assertRaises(QiskitError):
+            InstructionToSignals(dt=1, carriers=[1], channels=["d0", "d1"])
+
+        with self.assertRaises(QiskitError):
+            InstructionToSignals(dt=1, carriers=[1, 2], channels=["d0"])