[WIP] Injecting cals in calibration experiment

qiskit_experiments/calibration_management/transpiler.py

@@ -0,0 +1,144 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2021.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Transpiler pass for calibration experiments."""
+
+from typing import Dict, List, Tuple, Union
+
+from qiskit import QuantumCircuit
+from qiskit.dagcircuit import DAGCircuit
+from qiskit.pulse.schedule import ScheduleBlock
+from qiskit.transpiler.basepasses import TransformationPass
+from qiskit.transpiler.passmanager import PassManager
+
+from qiskit_experiments.calibration_management.calibrations import Calibrations
+from qiskit_experiments.exceptions import CalibrationError
+from qiskit_experiments.framework.base_experiment import BaseExperiment
+
+
+class CalibrationsAdder(TransformationPass):
+    """A transformation pass to add calibrations for standard circuit instructions.
+
+    This transpiler pass is intended to be run in the :meth:`circuits` method of the
+    experiment classes before the main transpiler pass. It's only goal is to extract
+    the needed pulse schedules from an instance of Calibrations and attach them to the
+    template circuit. This has a couple of challenges. Note that if no mapping is provided
+    this transpiler pass assumes that the name of the schedule in the calibrations is the
+    same as the name of the gate instruction.
+    """
+
+    def __init__(
+            self,
+            calibrations: Calibrations,
+            instructions_map: Dict[str, str],
+            qubit_layout: Dict[int, int],
+    ):
+        """Initialize the pass.
+
+        Args:
+            calibrations: An instance of calibration from which to fetch the schedules.
+            instructions_map: A map of circuit instruction names (keys) to schedule names stored
+                in the calibrations (values). If an entry is not found the pass will assume that
+                the instruction in the circuit and the schedule have the same name.
+            qubit_layout: The initial layout that will be used.
+        """
+        super().__init__()
+        self._qubit_layout = qubit_layout
+        self._cals = calibrations
+        self._instructions_map = instructions_map
+
+    def _get_calibration(self, gate: str, qubits: Tuple[int, ...]) -> Union[ScheduleBlock, None]:
+        """Get a schedule from the internally stored calibrations.
+
+        Args:
+            gate: Name of the gate for which to get the schedule.
+            qubits: The qubits for which to get the parameters.
+
+        Returns:
+            The schedule if one is found otherwise return None.
+        """
+
+        # Extract the gate to schedule and any parameter name mappings.
+        sched_name = self._instructions_map.get(gate, gate)
+
+        # Try and get a schedule, if there is none then return None.
+        try:
+            return self._cals.get_schedule(sched_name, qubits)
+        except CalibrationError:
+            return None
+
+    def run(self, dag: DAGCircuit) -> DAGCircuit:
+        """Run the calibration adder pass on `dag`.
+
+        Args:
+            dag: DAG to schedule.
+
+        Returns:
+            A DAG with calibrations added to it.
+        """
+        bit_indices = {bit: index for index, bit in enumerate(dag.qubits)}
+
+        for node in dag.nodes():
+            if node.type == "op":
+
+                # Get the qubit indices in the circuit.
+                qubits = tuple(bit_indices[qarg] for qarg in node.qargs)
+
+                # Get the physical qubits that they will remap to.
+                qubits = tuple(self._qubit_layout[qubit] for qubit in qubits)
+
+                schedule = self._get_calibration(node.op.name, qubits)
+
+                # Permissive stance: if we don't find a schedule we continue.
+                # The call to the transpiler that happens before running the
+                # experiment will either complain or force us to use the
+                # backend gates.
+                if schedule is None:
+                    continue
+
+                if len(set(qubits) & set(ch.index for ch in schedule.channels)) == 0:
+                    raise CalibrationError(
+                        f"None of the qubits {qubits} are contained in the channels of "
+                        f"the schedule named {schedule.name} for gate {node.op}."
+                    )
+
+                dag.add_calibration(node.op, qubits, schedule)
+
+        return dag
+
+
+def inject_calibrations(
+    circuits: Union[QuantumCircuit, List[QuantumCircuit]],
+    experiment: BaseExperiment
+) -> Union[QuantumCircuit, List[QuantumCircuit]]:
+    """Inject calibrations from a :class:`Calibrations` instance into a circuit.
+
+    This function only adds calibrations if it can find them in the calibrations.
+
+    Args:
+        circuits: The circuit or list of circuits into which to inject calibrations.
+        experiment: The experiment object that
+
+    Returns:
+         A quantum circuit with the relevant calibrations added to it.
+    """
+    layout = {idx: qubit for idx, qubit in enumerate(experiment.physical_qubits)}
+
+    # Identify the available schedule data.
+    calibrations = experiment.experiment_options.get("calibrations", None)
+
+    # Run the transpiler pass according to the available data
+    if calibrations is not None:
+        inst_maps = experiment.experiment_options.get("instruction_name_map", None) or dict()
+        return PassManager(CalibrationsAdder(calibrations, inst_maps, layout)).run(circuits)
+
+    return circuits

qiskit_experiments/library/calibration/fine_amplitude.py

@@ -27,6 +27,7 @@
     FineAmplitudeAnalysis,
 )
 from qiskit_experiments.exceptions import CalibrationError
+from qiskit_experiments.calibration_management.transpiler import inject_calibrations
 
 
 class FineAmplitude(BaseExperiment):
@@ -122,6 +123,7 @@ def _default_experiment_options(cls) -> Options:
         options.add_sx = False
         options.add_xp_circuit = True
         options.sx_schedule = None
+        options.calibrations = None
 
         return options
 
@@ -205,6 +207,7 @@ def circuits(self, backend: Optional[Backend] = None) -> List[QuantumCircuit]:
         if schedule is None:
             raise CalibrationError("No schedule set for fine amplitude calibration.")
 
+        # TODO Get ride of this?
         if self.physical_qubits[0] not in set(ch.index for ch in schedule.channels):
             raise CalibrationError(
                 f"User provided schedule {schedule.name} does not contain a channel "
@@ -269,6 +272,8 @@ def circuits(self, backend: Optional[Backend] = None) -> List[QuantumCircuit]:
 
             circuits.append(circuit)
 
+        circuits = inject_calibrations(circuits, self)
+
         return circuits
 
 

qiskit_experiments/library/calibration/rabi.py

@@ -20,11 +20,13 @@
 from qiskit.qobj.utils import MeasLevel
 from qiskit.providers import Backend
 import qiskit.pulse as pulse
+from qiskit.pulse import ScheduleBlock
 from qiskit.providers.options import Options
 
 from qiskit_experiments.framework import BaseExperiment
 from qiskit_experiments.library.calibration.analysis.oscillation_analysis import OscillationAnalysis
 from qiskit_experiments.exceptions import CalibrationError
+from qiskit_experiments.calibration_management.transpiler import inject_calibrations
 
 
 class Rabi(BaseExperiment):
@@ -72,6 +74,8 @@ def _default_experiment_options(cls) -> Options:
             sigma=40,
             amplitudes=np.linspace(-0.95, 0.95, 51),
             schedule=None,
+            calibrations=None,
+            instruction_name_map=None,
         )
 
     @classmethod
@@ -90,37 +94,40 @@ def __init__(self, qubit: int):
         - duration: The duration of the rabi pulse in samples, the default is 160 samples.
         - sigma: The standard deviation of the pulse, the default is duration 40.
         - amplitudes: The amplitude that are scanned in the experiment, default  is
-        np.linspace(-0.95, 0.95, 51)
+          np.linspace(-0.95, 0.95, 51).
 
         Args:
             qubit: The qubit on which to run the Rabi experiment.
         """
         super().__init__([qubit])
 
-    def _template_circuit(self, amp_param) -> QuantumCircuit:
+    def _template_circuit(self, sched: ScheduleBlock) -> QuantumCircuit:
         """Return the template quantum circuit."""
-        gate = Gate(name=self.__rabi_gate_name__, num_qubits=1, params=[amp_param])
+        gate = Gate(name=self.__rabi_gate_name__, num_qubits=1, params=list(sched.parameters))
 
         circuit = QuantumCircuit(1)
         circuit.append(gate, (0,))
         circuit.measure_active()
 
+        circuit.add_calibration(
+            self.__rabi_gate_name__, (self.physical_qubits[0],), sched, list(sched.parameters)
+        )
+
         return circuit
 
-    def _default_gate_schedule(self, backend: Optional[Backend] = None):
+    def default_schedules(self, backend: Optional[Backend] = None):
         """Create the default schedule for the Rabi gate."""
-        amp = Parameter("amp")
-        with pulse.build(backend=backend, name="rabi") as default_schedule:
+        with pulse.build(backend=backend, name=self.__rabi_gate_name__) as schedule:
             pulse.play(
                 pulse.Gaussian(
                     duration=self.experiment_options.duration,
-                    amp=amp,
+                    amp=Parameter("amp"),
                     sigma=self.experiment_options.sigma,
                 ),
                 pulse.DriveChannel(self.physical_qubits[0]),
             )
 
-        return default_schedule
+        return schedule
 
     def circuits(self, backend: Optional[Backend] = None) -> List[QuantumCircuit]:
         """Create the circuits for the Rabi experiment.
@@ -138,30 +145,30 @@ def circuits(self, backend: Optional[Backend] = None) -> List[QuantumCircuit]:
                   that matches the qubit on which to run the Rabi experiment.
                 - If the user provided schedule has more than one free parameter.
         """
-        schedule = self.experiment_options.get("schedule", None)
 
-        if schedule is None:
-            schedule = self._default_gate_schedule(backend=backend)
-        else:
-            if self.physical_qubits[0] not in set(ch.index for ch in schedule.channels):
-                raise CalibrationError(
-                    f"User provided schedule {schedule.name} does not contain a channel "
-                    "for the qubit on which to run Rabi."
-                )
+        # 1. Get the schedules for the custom gates.
+        schedule = self.experiment_options.schedule or self.default_schedules(backend)
+
+        if self.physical_qubits[0] not in set(ch.index for ch in schedule.channels):
+            raise CalibrationError(
+                f"Schedule {schedule.name} does not have a channel "
+                f"for the qubit on which to run {self.__class__.__name__}."
+            )
 
         if len(schedule.parameters) != 1:
-            raise CalibrationError("Schedule in Rabi must have exactly one free parameter.")
+            raise CalibrationError(
+                f"Schedule in {self.__class__.__name__} must have exactly one free parameter."
+            )
 
-        param = next(iter(schedule.parameters))
+        # 2. Create template circuit and attach the calibration.
+        circuit = self._template_circuit(schedule)
 
-        # Create template circuit
-        circuit = self._template_circuit(param)
-        circuit.add_calibration(
-            self.__rabi_gate_name__, (self.physical_qubits[0],), schedule, params=[param]
-        )
+        # 3. Inject calibrations for standard gates.
+        circuit = inject_calibrations(circuit, self)
 
-        # Create the circuits to run
+        # 4. Assign parameter values to create circuits to run on.
         circs = []
+        param = next(iter(schedule.parameters))
         for amp in self.experiment_options.amplitudes:
             amp = np.round(amp, decimals=6)
             assigned_circ = circuit.assign_parameters({param: amp}, inplace=False)
@@ -213,18 +220,12 @@ def _default_experiment_options(cls) -> Options:
             ef_rabi.set_experiment_options(schedule=rabi_schedule)
 
         """
-        return Options(
-            duration=160,
-            sigma=40,
-            amplitudes=np.linspace(-0.95, 0.95, 51),
-            schedule=None,
-            normalization=True,
-            frequency_shift=None,
-        )
+        options = super()._default_experiment_options()
+        options.frequency_shift=None
+        return options
 
-    def _default_gate_schedule(self, backend: Optional[Backend] = None):
-        """Create the default schedule for the EFRabi gate with a frequency shift to the 1-2
-        transition."""
+    def default_schedules(self, backend: Optional[Backend] = None):
+        """Create the default schedule with a frequency shift to the 1-2 transition."""
 
         if self.experiment_options.frequency_shift is None:
             try:
@@ -244,7 +245,6 @@ def _default_gate_schedule(self, backend: Optional[Backend] = None):
                     "to be set manually through EFRabi.set_experiment_options(frequency_shift=..)."
                 ) from att_err
 
-        amp = Parameter("amp")
         with pulse.build(backend=backend, name=self.__rabi_gate_name__) as default_schedule:
             with pulse.frequency_offset(
                 self.experiment_options.frequency_shift,
@@ -253,19 +253,17 @@ def _default_gate_schedule(self, backend: Optional[Backend] = None):
                 pulse.play(
                     pulse.Gaussian(
                         duration=self.experiment_options.duration,
-                        amp=amp,
+                        amp=Parameter("amp"),
                         sigma=self.experiment_options.sigma,
                     ),
                     pulse.DriveChannel(self.physical_qubits[0]),
                 )
 
         return default_schedule
 
-    def _template_circuit(self, amp_param) -> QuantumCircuit:
+    def _template_circuit(self, sched: ScheduleBlock) -> QuantumCircuit:
         """Return the template quantum circuit."""
-        circuit = QuantumCircuit(1)
+        circuit = QuantumCircuit(1, 1)
         circuit.x(0)
-        circuit.append(Gate(name=self.__rabi_gate_name__, num_qubits=1, params=[amp_param]), (0,))
-        circuit.measure_active()
 
-        return circuit
+        return circuit.compose(super()._template_circuit(sched))

test/calibration/experiments/test_rabi.py

@@ -27,6 +27,7 @@
 from qiskit_experiments.library import Rabi, EFRabi
 
 from qiskit_experiments.library.calibration.analysis.oscillation_analysis import OscillationAnalysis
+from qiskit_experiments.calibration_management.calibrations import Calibrations
 from qiskit_experiments.data_processing.data_processor import DataProcessor
 from qiskit_experiments.data_processing.nodes import Probability
 from qiskit_experiments.test.mock_iq_backend import MockIQBackend
@@ -190,6 +191,30 @@ def test_user_schedule(self):
         assigned_sched = my_schedule.assign_parameters({amp: 0.5}, inplace=False)
         self.assertEqual(circs[0].calibrations["Rabi"][((2,), (0.5,))], assigned_sched)
 
+    def test_transpile_x_gate(self):
+        """Test that we can use the calibrations to transpile in the x gate."""
+        cals = Calibrations()
+        qubit = 3
+
+        chan = Parameter("ch0")
+        amp = Parameter("amp")
+
+        with pulse.build(name="x") as xp_sched:
+            pulse.play(pulse.Gaussian(123, amp, 25), pulse.DriveChannel(chan))
+
+        cals.add_schedule(xp_sched)
+        cals.add_parameter_value(0.2, "amp", schedule="x")
+
+        rabi = EFRabi(qubit)
+        rabi.set_experiment_options(calibrations=cals, frequency_shift=-330e6)
+
+        circs = rabi.circuits(RabiBackend())
+        print(circs[0].calibrations)
+
+        expected = xp_sched.assign_parameters({amp: 0.2, chan: 3}, inplace=False)
+
+        self.assertEqual(circs[0].calibrations["x"][(3, ), ()], expected)
+
 
 class TestRabiAnalysis(QiskitTestCase):
     """Class to test the fitting."""