Drawing a ControlledQubitUnitary (#1174)

* ControlledUnitary drawing and testing * changelog Co-authored-by: Nathan Killoran <co9olguy@users.noreply.github.com>
PennyLaneAI · Mar 29, 2021 · 3212c71 · 3212c71
1 parent d0d4f81
commit 3212c71
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 5 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -672,6 +672,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where using the circuit drawer with a ``ControlledQubitUnitary``
+  operation raised an error.
+  [(#1174)](https://github.com/PennyLaneAI/pennylane/pull/1174)
+
 * Fixes a bug and a test where the ``QuantumTape.is_sampled`` attribute was not
   being updated.
   [(#1126)](https://github.com/PennyLaneAI/pennylane/pull/1126)

diff --git a/pennylane/circuit_drawer/representation_resolver.py b/pennylane/circuit_drawer/representation_resolver.py
@@ -134,6 +134,23 @@ def _format_matrix_operation(operation, symbol, cache):
 
         return "{}{}".format(symbol, idx)
 
+    @staticmethod
+    def _format_controlled_qubit_unitary(operation, symbol, cache):
+        """Format an operation that corresponds to a single matrix with controls.
+
+        Args:
+            operation (~.Operation): Operation that shall be formatted
+            symbol (str): The symbol that should be used to identify matrices
+            cache (List[numpy.ndarray]): The cache of already known matrices
+
+        Returns:
+            str: The formatted operation
+        """
+        mat = operation.U
+        idx = RepresentationResolver.index_of_array_or_append(mat, cache)
+
+        return "{}{}".format(symbol, idx)
+
     @staticmethod
     def _format_matrix_arguments(params, symbol, cache):
         """Format a sequence of matrix parameters.
@@ -351,6 +368,14 @@ def operator_representation(self, op, wire):
                 op, "U", self.unitary_matrix_cache
             )
 
+        elif base_name == "ControlledQubitUnitary":
+            if wire in op.control_wires:
+                return self.charset.CONTROL
+
+            representation = RepresentationResolver._format_controlled_qubit_unitary(
+                op, "U", self.unitary_matrix_cache
+            )
+
         elif base_name == "Hermitian":
             representation = RepresentationResolver._format_matrix_operation(
                 op, "H", self.hermitian_matrix_cache

diff --git a/pennylane/ops/qubit.py b/pennylane/ops/qubit.py
@@ -1907,6 +1907,11 @@ def __init__(self, *params, control_wires=None, wires=None, control_values=None,
         target_dim = 2 ** len(wires)
         if len(U) != target_dim:
             raise ValueError(f"Input unitary must be of shape {(target_dim, target_dim)}")
+
+        # Saving for the circuit drawer
+        self.control_wires = control_wires
+        self.U = U
+
         wires = control_wires + wires
 
         # If control values unspecified, we control on the all-ones string

diff --git a/tests/circuit_drawer/test_circuit_drawer.py b/tests/circuit_drawer/test_circuit_drawer.py
@@ -274,6 +274,7 @@ def qfunc(a, b, c, angles):
         qml.CRY(0.3589, wires=[3, 1]).inv()
         qml.CSWAP(wires=[4, 2, 1]).inv()
         qml.QubitUnitary(np.eye(2), wires=[2])
+        qml.ControlledQubitUnitary(np.eye(2), control_wires=[0, 1], wires=[2])
         qml.Toffoli(wires=[0, 2, 1])
         qml.CNOT(wires=[0, 2])
         qml.PauliZ(wires=[1])
@@ -326,11 +327,11 @@ def drawn_parameterized_qubit_circuit_with_variable_names():
 def drawn_parameterized_qubit_circuit_with_values():
     """The rendered circuit representation of the above qubit circuit with variable values."""
     return (
-        " 0: ──RX(0.1)───────────────╭C─────RX(0.4)──────────────────────────────────────╭C─────╭C───────╭C──╭C────────────╭C──╭SWAP⁻¹───╭SWAP───┤ ⟨Y⟩       \n"
-        + " 1: ──RX(0.2)────Z──────────╰X⁻¹──╭RY(0.2)──RX(2)────╭RY(0.359)⁻¹──╭SWAP⁻¹──────├X──Z──│───Z⁻¹──╰Z──│─────╭X──╭C──│───│─────────├SWAP───┤ Var[H]    \n"
-        + " 2: ──Rϕ(0.567)──RX(0.6)⁻¹────────│──────────────────│─────────────├SWAP⁻¹──U0──╰C─────╰X───────────╰Z⁻¹──╰C──│───╰X──╰SWAP⁻¹───│───────┤ Sample[X] \n"
-        + " 3: ──────────────────────────────╰C────────RZ(0.2)──╰C────────────│──────────────────────────────────────────╰X───────RZ(0.2)──│──────╭┤ ⟨H0⟩      \n"
-        + " 4: ───────────────────────────────────────────────────────────────╰C───────────────────────────────────────────────────────────╰C─────╰┤ ⟨H0⟩      \n"
+        " 0: ──RX(0.1)───────────────╭C─────RX(0.4)──────────────────────────────────────╭C───╭C─────╭C───────╭C──╭C────────────╭C──╭SWAP⁻¹───╭SWAP───┤ ⟨Y⟩       \n"
+        + " 1: ──RX(0.2)────Z──────────╰X⁻¹──╭RY(0.2)──RX(2)────╭RY(0.359)⁻¹──╭SWAP⁻¹──────├C───├X──Z──│───Z⁻¹──╰Z──│─────╭X──╭C──│───│─────────├SWAP───┤ Var[H]    \n"
+        + " 2: ──Rϕ(0.567)──RX(0.6)⁻¹────────│──────────────────│─────────────├SWAP⁻¹──U0──╰U0──╰C─────╰X───────────╰Z⁻¹──╰C──│───╰X──╰SWAP⁻¹───│───────┤ Sample[X] \n"
+        + " 3: ──────────────────────────────╰C────────RZ(0.2)──╰C────────────│───────────────────────────────────────────────╰X───────RZ(0.2)──│──────╭┤ ⟨H0⟩      \n"
+        + " 4: ───────────────────────────────────────────────────────────────╰C────────────────────────────────────────────────────────────────╰C─────╰┤ ⟨H0⟩      \n"
         + "U0 =\n"
         + "[[1. 0.]\n"
         + " [0. 1.]]\n"