Add pulse support to the PennyLane-Braket plugin for the OQC device

doc/devices/braket_remote.rst

@@ -99,10 +99,52 @@ from :mod:`braket.pennylane_plugin.ops <.ops>`:
     braket.pennylane_plugin.GPi2
     braket.pennylane_plugin.MS
 
+Pulse Programming
+~~~~~~~~~~~~~~~~~
+
+The PennyLane-Braket plugin provides pulse-level control for the OQC Lucy QPU through PennyLane's `ParametrizedEvolution <https://docs.pennylane.ai/en/latest/code/api/pennylane.pulse.ParametrizedEvolution.html>`_
+operation. Compatible pulse Hamiltonians can be defined using the `qml.pulse.transmon_drive <https://docs.pennylane.ai/en/latest/code/api/pennylane.pulse.transmon_drive.html>`_
+function and used to create ``ParametrizedEvolution``'s using `qml.evolve <https://docs.pennylane.ai/en/stable/code/api/pennylane.evolve.html>`_:
+
+.. code-block:: python
+
+    duration = 15
+    def amp(p, t):
+        return qml.pulse.pwc(duration)(p, t)
+
+    dev = qml.device("braket.aws.qubit", wires=8, device_arn="arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy")
+
+    drive = qml.pulse.transmon.transmon_drive(amplitude=amp, phase=0, freq=4.8, wires=[0])
+
+    @qml.qnode(dev)
+    def circuit(params, t):
+        qml.evolve(drive)(params, t)
+        return qml.expval(qml.PauliZ(wires=0))
+
+Note that the ``freq`` argument of ``qml.pulse.transmon_drive`` is specified in GHz, and for
+hardware upload the amplitude will be interpreted as an output power for control hardware in volts.
+The ``phase`` must be specified in radians.
+
+The pulse settings for the device can be obtained using the ``pulse_settings`` property. These settings can be used to describe the transmon
+interaction Hamiltonian using `qml.pulse.transmon_interaction <https://docs.pennylane.ai/en/latest/code/api/pennylane.pulse.transmon_interaction.html>`_:
+
+    .. code-block:: python
+
+        dev = qml.device("braket.aws.qubit", wires=8, device_arn="arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy")
+        pulse_settings = dev.pulse_settings
+        couplings = [0.01]*len(connections)
+        H = qml.pulse.transmon_interaction(**pulse_settings, coupling=couplings)
+
+By passing ``pulse_settings`` from the remote device to ``qml.pulse.transmon_interaction``, an ``H`` Hamiltonian term is created using
+the constants specific to the hardware. This is relevant for simulating the hardware in PennyLane on the ``default.qubit`` device.
+
+Note that the user must supply coupling coefficients, as these are not available from the hardware backend. On the order of 10 MHz
+(0.01 GHz) is in a realistic range.
 
 Gradient computation on Braket with a QAOA Hamiltonian
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Currently, PennyLane will compute grouping indices for QAOA Hamiltonians and use them to split the Hamiltonian into multiple expectation values. If you wish to use `SV1’s adjoint differentiation capability<https://docs.aws.amazon.com/braket/latest/developerguide/hybrid.html>` when running QAOA from PennyLane, you will need reconstruct the cost Hamiltonian to remove the grouping indices from the cost Hamiltonian, like so:
+
+Currently, PennyLane will compute grouping indices for QAOA Hamiltonians and use them to split the Hamiltonian into multiple expectation values. If you wish to use `SV1’s adjoint differentiation capability <https://docs.aws.amazon.com/braket/latest/developerguide/hybrid.html>`_ when running QAOA from PennyLane, you will need reconstruct the cost Hamiltonian to remove the grouping indices from the cost Hamiltonian, like so:
 
 .. code-block:: python
 

doc/requirements.txt

@@ -1,4 +1,4 @@
 sphinx
 sphinx-automodapi
-pennylane>=0.23.0
+pennylane>=0.31.0
 pennylane-sphinx-theme

src/braket/pennylane_plugin/braket_device.py

@@ -420,12 +420,15 @@ def apply(
                 else:
                     param_names.append(None)
                 param_index += 1
+
+            dev_wires = self.map_wires(operation.wires).tolist()
             gate = translate_operation(
                 operation,
                 use_unique_params=bool(trainable_indices) or use_unique_params,
                 param_names=param_names,
+                device=self._device,
             )
-            dev_wires = self.map_wires(operation.wires).tolist()
+
             ins = Instruction(gate, dev_wires)
             circuit.add_instruction(ins)
 
@@ -708,6 +711,107 @@ def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires):
 
         return outcomes
 
+    def _check_pulse_frequency_validity(self, ev):
+        """Confirm that, for each waveform on the ParametrizedEvolution operator, the frequency
+        setting is a constant, and the value is within the frequency range for the relevant frame;
+        if not, raise an error"""
+
+        # confirm all frequency values are constant (or the qml.pulse.constant function)
+        callable_freqs = [
+            pulse.frequency
+            for pulse in ev.H.pulses
+            if (callable(pulse.frequency) and pulse.frequency != qml.pulse.constant)
+        ]
+
+        if callable_freqs:
+            raise RuntimeError(
+                "Expected all frequencies to be constants or qml.pulse.constant, "
+                "but received callable(s)"
+            )
+
+        # confirm all frequencies are within permitted difference from center frequency
+        param_idx = 0
+        for pulse in ev.H.pulses:
+            freq = pulse.frequency
+            # track the index for parameters in case we need to evaluate qml.pulse.constant
+            if callable(pulse.amplitude):
+                param_idx += 1
+            if callable(pulse.phase):
+                param_idx += 1
+            if callable(pulse.frequency):
+                # if frequency is callable, its qml.pulse.constant and equal to its parameter
+                freq = ev.parameters[param_idx]
+                param_idx += 1
+
+            wires = self.map_wires(pulse.wires).tolist()
+            freq_min = 3  # GHz
+            freq_max = 8
+
+            if not (freq_min < freq < freq_max):
+                raise RuntimeError(
+                    f"Frequency range for wire(s) {wires} is between {freq_min} "
+                    f"and {freq_max} GHz, but received {freq} GHz."
+                )
+
+    def _validate_pulse_parameters(self, ev):
+        """Validates pulse input (ParametrizedEvolution) before converting to a PulseGate"""
+
+        # note: the pulse upload on the AWS service checks at task creation that the max amplitude
+        # is not exceeded, so that check has not been included here
+
+        # confirm frequencies are constant and within the permitted frequency range for the channel
+        self._check_pulse_frequency_validity(ev)
+
+        # confirm all phase values are constant (or the qml.pulse.constant function)
+        callable_phase = [
+            pulse.phase
+            for pulse in ev.H.pulses
+            if (callable(pulse.phase) and pulse.phase != qml.pulse.constant)
+        ]
+
+        if callable_phase:
+            raise RuntimeError(
+                "Expected all phases to be constants or qml.pulse.constant, "
+                "but received callable(s)"
+            )
+
+        # ensure each ParametrizedEvolution/PulseGate contains at most one waveform per frame/wire
+        wires_used = []
+        for pulse in ev.H.pulses:
+            for wire in pulse.wires:
+                if wire in wires_used:
+                    raise RuntimeError(
+                        f"Multiple waveforms assigned to wire {wire} in the same "
+                        f"ParametrizedEvolution gate"
+                    )
+                wires_used.append(wire)
+
+    def check_validity(self, queue, observables):
+        """Check validity of pulse operations before running the standard check_validity function
+
+        Checks whether the operations and observables in queue are all supported by the device. Runs
+        the standard check_validity function for a PennyLane device, and an additional check to
+        validate any pulse-operations in the form of a ParametrizedEvolution operation.
+
+        Args:
+            queue (Iterable[~.operation.Operation]): quantum operation objects which are intended
+                to be applied on the device
+            observables (Iterable[~.operation.Observable]): observables which are intended
+                to be evaluated on the device
+
+        Raises:
+            DeviceError: if there are operations in the queue or observables that the device does
+                not support
+            RuntimeError: if there are ParametrizedEvolution operations in the queue that are not
+                supported because of invalid pulse parameters
+        """
+
+        super().check_validity(queue, observables)
+
+        for op in queue:
+            if isinstance(op, qml.pulse.ParametrizedEvolution):
+                self._validate_pulse_parameters(op)
+
     def capabilities(self=None):
         """Add support for AG on sv1"""
         # normally, we'd just call super().capabilities() here, but super()
@@ -754,6 +858,109 @@ def execute_and_gradients(self, circuits, **kwargs):
             jacs.append(new_jac)
         return res, jacs
 
+    def _is_single_qubit_01_frame(self, f_string, wire=None):
+        """Defines the condition for selecting frames addressing the qubit (01)
+        drive based on frame name"""
+        if self._device.arn == "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy":
+            if wire is not None:
+                return f_string == f"q{wire}_drive"
+            return "drive" in f_string
+        else:
+            raise NotImplementedError(
+                f"Single-qubit drive frame for pulse control not defined for "
+                f"device {self._device.arn}"
+            )
+
+    def _is_single_qubit_12_frame(self, f_string, wire=None):
+        """Defines the condition for selecting frames addressing excitation to
+        the second excited state based on frame name"""
+        if self._device.arn == "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy":
+            if wire is not None:
+                return f_string == f"q{wire}_second_state"
+            return "second_state" in f_string
+        else:
+            raise NotImplementedError(
+                f"Second excitation drive frame for pulse control not defined for "
+                f"device {self._device.arn}"
+            )
+
+    def _get_frames(self, filter, wires):
+        """Takes a filter defining how the relevant frames are labelled, and returns all the frames
+        that fit, i.e.:
+
+        cond = lambda frame_id, wire: f"q{wire}_drive" == frame_id
+        frames = self._get_frames(cond, wires=[0, 1, 2])
+
+        would return all the frames with ids "q0_drive" "q1_drive", and "q2_drive", stored
+        in a dictionary with keys [0, 1, 2] identifying the qubit number.
+        """
+        if not self._device.arn == "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy":
+            raise NotImplementedError(
+                f"Accessing drive frame for pulse control is not defined for "
+                f"device {self._device.arn}"
+            )
+
+        frames = {}
+        for wire in wires:
+            for frame, info in self._device.properties.pulse.dict()["frames"].items():
+                if filter(frame, wire):
+                    frames[wire] = info
+
+        return frames
+
+    @property
+    def pulse_settings(self):
+        """Dictionary of constants set by the hardware (qubit resonant frequencies,
+        inter-qubit connection graph, wires and anharmonicities).
+
+        Used to enable initializing hardware-consistent Hamiltonians by returning
+        values that would need to be passed, i.e.:
+
+            >>> dev_remote = qml.device('braket.aws.qubit',
+            >>>                          wires=8,
+            >>>                          arn='arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy')
+            >>> pulse_settings = dev_remote.pulse_settings
+            >>> H_int = qml.pulse.transmon_interaction(**pulse_settings, coupling=0.02)
+
+        By passing the ``pulse_settings`` from the remote device to ``transmon_interaction``, an
+        ``H_int`` Hamiltonian term is created using the constants specific to the hardware.
+        This is relevant for simulating the hardware in PennyLane on the ``default.qubit`` device.
+
+        Note that the user must supply coupling coefficients, as these are not available from the
+        hardware backend.
+        """
+        if not self._device.arn == "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy":
+            raise NotImplementedError(
+                f"The pulse_settings property for pulse control is not defined for "
+                f"device {self._device.arn}"
+            )
+
+        device_info = self._device.properties.paradigm
+        wires = [i for i in range(device_info.qubitCount)]
+
+        drive_frames_01 = self._get_frames(filter=self._is_single_qubit_01_frame, wires=wires)
+        drive_frames_12 = self._get_frames(filter=self._is_single_qubit_12_frame, wires=wires)
+
+        qubit_freq = [drive_frames_01[wire]["frequency"] * 1e-9 for wire in wires]  # Hz to GHz
+
+        connections = []
+        for q1, connected_qubits in device_info.connectivity.connectivityGraph.items():
+            for q2 in connected_qubits:
+                connection = (int(q1), int(q2))
+                connections.append(connection)
+
+        anharmonicity = [
+            (drive_frames_01[wire]["frequency"] - drive_frames_12[wire]["frequency"]) * 1e-9
+            for wire in wires
+        ]
+
+        return {
+            "qubit_freq": qubit_freq,
+            "connections": connections,
+            "wires": wires,
+            "anharmonicity": anharmonicity,
+        }
+
 
 class BraketLocalQubitDevice(BraketQubitDevice):
     r"""Amazon Braket LocalSimulator qubit device for PennyLane.