Add controlled unitary gate

.github/CHANGELOG.md

@@ -47,6 +47,9 @@
 
 <h3>Improvements</h3>
 
+- Added the `ControlledQubitUnitary` operation.
+  [(#1069)](https://github.com/PennyLaneAI/pennylane/pull/1069)
+
 * Most layers in Pytorch or Keras accept arbitrary dimension inputs, where each dimension barring
   the last (in the case where the actual weight function of the layer operates on one-dimensional
   vectors) is broadcast over. This is now also supported by KerasLayer and TorchLayer.

doc/introduction/operations.rst

@@ -81,6 +81,7 @@ Qubit gates
     ~pennylane.Toffoli
     ~pennylane.CSWAP
     ~pennylane.QubitUnitary
+    ~pennylane.ControlledQubitUnitary
     ~pennylane.DiagonalQubitUnitary
     ~pennylane.QFT
 

pennylane/devices/default_mixed.py

@@ -62,6 +62,7 @@ class DefaultMixed(QubitDevice):
         "BasisState",
         "QubitStateVector",
         "QubitUnitary",
+        "ControlledQubitUnitary",
         "DiagonalQubitUnitary",
         "PauliX",
         "PauliY",

pennylane/devices/default_qubit.py

@@ -97,6 +97,7 @@ class DefaultQubit(QubitDevice):
         "BasisState",
         "QubitStateVector",
         "QubitUnitary",
+        "ControlledQubitUnitary",
         "DiagonalQubitUnitary",
         "PauliX",
         "PauliY",

pennylane/devices/tests/test_gates.py

@@ -56,6 +56,7 @@
     "ControlledPhaseShift": qml.ControlledPhaseShift(0, wires=[0, 1]),
     "QubitStateVector": qml.QubitStateVector(np.array([1.0, 0.0]), wires=[0]),
     "QubitUnitary": qml.QubitUnitary(np.eye(2), wires=[0]),
+    "ControlledQubitUnitary": qml.ControlledQubitUnitary(np.eye(2), control_wires=[1], wires=[0]),
     "RX": qml.RX(0, wires=[0]),
     "RY": qml.RY(0, wires=[0]),
     "RZ": qml.RZ(0, wires=[0]),

pennylane/ops/qubit.py

@@ -22,12 +22,14 @@
 import math
 
 import numpy as np
+from scipy.linalg import block_diag
 
 import pennylane as qml
 from pennylane.operation import AnyWires, DiagonalOperation, Observable, Operation
 from pennylane.templates import template
 from pennylane.templates.state_preparations import BasisStatePreparation, MottonenStatePreparation
 from pennylane.utils import expand, pauli_eigs
+from pennylane.wires import Wires
 
 INV_SQRT2 = 1 / math.sqrt(2)
 
@@ -1601,6 +1603,64 @@ def _matrix(cls, *params):
         return U
 
 
+class ControlledQubitUnitary(QubitUnitary):
+    r"""ControlledQubitUnitary(U, control_wires, wires)
+    Apply an arbitrary fixed unitary to ``wires`` with control from the ``control_wires``.
+
+    **Details:**
+
+    * Number of wires: Any (the operation can act on any number of wires)
+    * Number of parameters: 1
+    * Gradient recipe: None
+
+    Args:
+        U (array[complex]): square unitary matrix
+        control_wires (Union[Wires, Sequence[int], or int]): the control wire(s)
+        wires (Union[Wires, Sequence[int], or int]): the wire(s) the unitary acts on
+
+    **Example**
+
+    The following shows how a single-qubit unitary can be applied to wire ``2`` with control on
+    both wires ``0`` and ``1``:
+
+    >>> U = np.array([[ 0.94877869,  0.31594146], [-0.31594146,  0.94877869]])
+    >>> qml.ControlledQubitUnitary(U, control_wires=[0, 1], wires=2)
+    """
+    num_params = 1
+    num_wires = AnyWires
+    par_domain = "A"
+    grad_method = None
+
+    def __init__(self, *params, control_wires=None, wires=None, do_queue=True):
+        if control_wires is None:
+            raise ValueError("Must specify control wires")
+
+        wires = Wires(wires)
+        control_wires = Wires(control_wires)
+
+        if Wires.shared_wires([wires, control_wires]):
+            raise ValueError(
+                "The control wires must be different from the wires specified to apply the unitary on."
+            )
+
+        U = params[0]
+        target_dim = 2 ** len(wires)
+        if len(U) != target_dim:
+            raise ValueError(f"Input unitary must be of shape {(target_dim, target_dim)}")
+
+        wires = control_wires + wires
+
+        # Given that the controlled wires are listed before the target wires, we need to create a
+        # block-diagonal matrix of shape ((I, 0), (0, U)) where U acts on the target wires and I
+        # acts on the control wires.
+        padding = 2 ** len(wires) - len(U)
+        CU = block_diag(np.eye(padding), U)
+
+        params = list(params)
+        params[0] = CU
+        super().__init__(*params, wires=wires, do_queue=do_queue)
+
+
 class DiagonalQubitUnitary(DiagonalOperation):
     r"""DiagonalQubitUnitary(D, wires)
     Apply an arbitrary fixed diagonal unitary matrix.
@@ -1914,6 +1974,7 @@ def diagonalizing_gates(self):
     "BasisState",
     "QubitStateVector",
     "QubitUnitary",
+    "ControlledQubitUnitary",
     "DiagonalQubitUnitary",
     "QFT",
 }

tests/ops/test_qubit_ops.py

@@ -18,6 +18,7 @@
 import functools
 import numpy as np
 from numpy.linalg import multi_dot
+from scipy.stats import unitary_group
 
 import pennylane as qml
 from pennylane.wires import Wires
@@ -1295,3 +1296,117 @@ def test_identity_eigvals(tol):
     res = qml.Identity._eigvals()
     expected = np.array([1, 1])
     assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestControlledQubitUnitary:
+    """Tests for the ControlledQubitUnitary operation"""
+
+    X = np.array([[0, 1], [1, 0]])
+
+    def test_matrix(self):
+        """Test if ControlledQubitUnitary returns the correct matrix for a control-control-X
+        (Toffoli) gate"""
+        mat = qml.ControlledQubitUnitary(X, control_wires=[0, 1], wires=2).matrix
+        mat2 = qml.Toffoli(wires=[0, 1, 2]).matrix
+        assert np.allclose(mat, mat2)
+
+    def test_no_control(self):
+        """Test if ControlledQubitUnitary raises an error if control wires are not specified"""
+        with pytest.raises(ValueError, match="Must specify control wires"):
+            qml.ControlledQubitUnitary(X, wires=2)
+
+    def test_shared_control(self):
+        """Test if ControlledQubitUnitary raises an error if control wires are shared with wires"""
+        with pytest.raises(ValueError, match="The control wires must be different from the wires"):
+            qml.ControlledQubitUnitary(X, control_wires=[0, 2], wires=2)
+
+    def test_wrong_shape(self):
+        """Test if ControlledQubitUnitary raises a ValueError if a unitary of shape inconsistent
+        with wires is provided"""
+        with pytest.raises(ValueError, match=r"Input unitary must be of shape \(2, 2\)"):
+            qml.ControlledQubitUnitary(np.eye(4), control_wires=[0, 1], wires=2)
+
+    @pytest.mark.parametrize("target_wire", range(3))
+    def test_toffoli(self, target_wire):
+        """Test if ControlledQubitUnitary acts like a Toffoli gate when the input unitary is a
+        single-qubit X. This test allows the target wire to be any of the three wires."""
+        control_wires = list(range(3))
+        del control_wires[target_wire]
+
+        # pick some random unitaries (with a fixed seed) to make the circuit less trivial
+        U1 = unitary_group.rvs(8, random_state=1)
+        U2 = unitary_group.rvs(8, random_state=2)
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def f1():
+            qml.QubitUnitary(U1, wires=range(3))
+            qml.ControlledQubitUnitary(X, control_wires=control_wires, wires=target_wire)
+            qml.QubitUnitary(U2, wires=range(3))
+            return qml.state()
+
+        @qml.qnode(dev)
+        def f2():
+            qml.QubitUnitary(U1, wires=range(3))
+            qml.Toffoli(wires=control_wires + [target_wire])
+            qml.QubitUnitary(U2, wires=range(3))
+            return qml.state()
+
+        state_1 = f1()
+        state_2 = f2()
+
+        assert np.allclose(state_1, state_2)
+
+    def test_arbitrary_multiqubit(self):
+        """Test if ControlledQubitUnitary applies correctly for a 2-qubit unitary with 2-qubit
+        control, where the control and target wires are not ordered."""
+        control_wires = [1, 3]
+        target_wires = [2, 0]
+
+        # pick some random unitaries (with a fixed seed) to make the circuit less trivial
+        U1 = unitary_group.rvs(16, random_state=1)
+        U2 = unitary_group.rvs(16, random_state=2)
+
+        # the two-qubit unitary
+        U = unitary_group.rvs(4, random_state=3)
+
+        # the 4-qubit representation of the unitary if the control wires were [0, 1] and the target
+        # wires were [2, 3]
+        U_matrix = np.eye(16, dtype=np.complex128)
+        U_matrix[12:16, 12:16] = U
+
+        # We now need to swap wires so that the control wires are [1, 3] and the target wires are
+        # [2, 0]
+        swap = qml.SWAP.matrix
+
+        # initial wire permutation: 0123
+        # target wire permutation: 1302
+        swap1 = np.kron(swap, np.eye(4))  # -> 1023
+        swap2 = np.kron(np.eye(4), swap)  # -> 1032
+        swap3 = np.kron(np.kron(np.eye(2), swap), np.eye(2))  # -> 1302
+        swap4 = np.kron(np.eye(4), swap)  # -> 1320
+
+        all_swap = swap4 @ swap3 @ swap2 @ swap1
+        U_matrix = all_swap.T @ U_matrix @ all_swap
+
+        dev = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(dev)
+        def f1():
+            qml.QubitUnitary(U1, wires=range(4))
+            qml.ControlledQubitUnitary(U, control_wires=control_wires, wires=target_wires)
+            qml.QubitUnitary(U2, wires=range(4))
+            return qml.state()
+
+        @qml.qnode(dev)
+        def f2():
+            qml.QubitUnitary(U1, wires=range(4))
+            qml.QubitUnitary(U_matrix, wires=range(4))
+            qml.QubitUnitary(U2, wires=range(4))
+            return qml.state()
+
+        state_1 = f1()
+        state_2 = f2()
+
+        assert np.allclose(state_1, state_2)

tests/test_operation.py

@@ -123,6 +123,9 @@ def test_heisenberg(self, test_class, tol):
     def test_operation_init(self, test_class, monkeypatch):
         "Operation subclass initialization."
 
+        if test_class == qml.ControlledQubitUnitary:
+            pytest.skip("ControlledQubitUnitary alters the input params and wires in its __init__")
+
         n = test_class.num_params
         w = test_class.num_wires
         ww = list(range(w))
@@ -199,6 +202,19 @@ def test_operation_init(self, test_class, monkeypatch):
             with pytest.raises(ValueError, match="Unknown parameter domain"):
                 test_class(*pars, wires=ww)
 
+    def test_controlled_qubit_unitary_init(self):
+        """Test for the init of ControlledQubitUnitary"""
+        control_wires = [3, 2]
+        target_wires = [1, 0]
+        U = qml.CRX._matrix(0.4)
+
+        op = qml.ControlledQubitUnitary(U, control_wires=control_wires, wires=target_wires)
+        target_data = [np.block([[np.eye(12), np.zeros((12, 4))], [np.zeros((4, 12)), U]])]
+
+        assert op.name == qml.ControlledQubitUnitary.__name__
+        assert np.allclose(target_data, op.data)
+        assert op._wires == Wires(control_wires) + Wires(target_wires)
+
     @pytest.fixture(scope="function")
     def qnode_for_inverse(self, mock_device):
         """Provides a QNode for the subsequent tests of inv"""