Repository: qiboteam/qibo

Issue:
`Circuit.unitary` breaks in the presente of two-qubit gate initialized by `Gate.controlled_by`
Code to reproduce bug:

```python
from qibo import Circuit, gates

circuit = Circuit(2)
circuit.add(gates.H(0).controlled_by(1))

circuit.unitary()
```

Discussion:

Patch:
diff --git a/src/qibo/backends/abstract.py b/src/qibo/backends/abstract.py
index 4a725c22e..8e4175b8e 100644
--- a/src/qibo/backends/abstract.py
+++ b/src/qibo/backends/abstract.py
@@ -133,11 +133,6 @@ class Backend(abc.ABC):
         """Fuse matrices of multiple gates."""
         raise_error(NotImplementedError)
 
-    @abc.abstractmethod
-    def control_matrix(self, gate):  # pragma: no cover
-        """ "Calculate full matrix representation of a controlled gate."""
-        raise_error(NotImplementedError)
-
     @abc.abstractmethod
     def apply_gate(self, gate, state, nqubits):  # pragma: no cover
         """Apply a gate to state vector."""
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 7248b2348..5760a62bd 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -3,6 +3,7 @@ import math
 
 import numpy as np
 from scipy import sparse
+from scipy.linalg import block_diag
 
 from qibo import __version__
 from qibo.backends import einsum_utils
@@ -109,14 +110,13 @@ class NumpyBackend(Backend):
         _matrix = getattr(self.matrices, name)
         if callable(_matrix):
             _matrix = _matrix(2 ** len(gate.target_qubits))
-
         return self.cast(_matrix, dtype=_matrix.dtype)
 
     def matrix_parametrized(self, gate):
         """Convert a parametrized gate to its matrix representation in the computational basis."""
         name = gate.__class__.__name__
-        matrix = getattr(self.matrices, name)(*gate.parameters)
-        return self.cast(matrix, dtype=matrix.dtype)
+        _matrix = getattr(self.matrices, name)(*gate.parameters)
+        return self.cast(_matrix, dtype=_matrix.dtype)
 
     def matrix_fused(self, fgate):
         rank = len(fgate.target_qubits)
@@ -127,6 +127,13 @@ class NumpyBackend(Backend):
             # small tensor calculations
             # explicit to_numpy see https://github.com/qiboteam/qibo/issues/928
             gmatrix = self.to_numpy(gate.matrix(self))
+            # add controls if controls were instantiated using
+            # the ``Gate.controlled_by`` method
+            num_controls = len(gate.control_qubits)
+            if num_controls > 0:
+                gmatrix = block_diag(
+                    np.eye(2 ** len(gate.qubits) - len(gmatrix)), gmatrix
+                )
             # Kronecker product with identity is needed to make the
             # original matrix have shape (2**rank x 2**rank)
             eye = np.eye(2 ** (rank - len(gate.qubits)))
@@ -148,30 +155,6 @@ class NumpyBackend(Backend):
 
         return self.cast(matrix.toarray())
 
-    def control_matrix(self, gate):
-        if len(gate.control_qubits) > 1:
-            raise_error(
-                NotImplementedError,
-                "Cannot calculate controlled "
-                "unitary for more than two "
-                "control qubits.",
-            )
-        matrix = gate.matrix(self)
-        shape = matrix.shape
-        if shape != (2, 2):
-            raise_error(
-                ValueError,
-                "Cannot use ``control_unitary`` method on "
-                + f"gate matrix of shape {shape}.",
-            )
-        zeros = self.np.zeros((2, 2), dtype=self.dtype)
-        zeros = self.cast(zeros, dtype=zeros.dtype)
-        identity = self.np.eye(2, dtype=self.dtype)
-        identity = self.cast(identity, dtype=identity.dtype)
-        part1 = self.np.concatenate([identity, zeros], axis=0)
-        part2 = self.np.concatenate([zeros, matrix], axis=0)
-        return self.np.concatenate([part1, part2], axis=1)
-
     def apply_gate(self, gate, state, nqubits):
         state = self.cast(state)
         state = self.np.reshape(state, nqubits * (2,))
diff --git a/src/qibo/gates/abstract.py b/src/qibo/gates/abstract.py
index 84f66c963..f222c027e 100644
--- a/src/qibo/gates/abstract.py
+++ b/src/qibo/gates/abstract.py
@@ -138,8 +138,7 @@ class Gate:
 
     @property
     def qubits(self) -> Tuple[int]:
-        """Tuple with ids of all qubits (control and target) that the gate
-        acts."""
+        """Tuple with ids of all qubits (control and target) that the gate acts."""
         return self.control_qubits + self.target_qubits
 
     @property
@@ -197,8 +196,7 @@ class Gate:
 
     @staticmethod
     def _find_repeated(qubits: Sequence[int]) -> int:
-        """Finds the first qubit id that is repeated in a sequence of qubit
-        ids."""
+        """Finds the first qubit id that is repeated in a sequence of qubit ids."""
         temp_set = set()
         for qubit in qubits:
             if qubit in temp_set:
@@ -213,14 +211,14 @@ class Gate:
         if common:
             raise_error(
                 ValueError,
-                f"{set(self._target_qubits) & set(self._control_qubits)} qubits are both targets and controls "
+                f"{set(self._target_qubits) & set(self._control_qubits)}"
+                + "qubits are both targets and controls "
                 + f"for gate {self.__class__.__name__}.",
             )
 
     @property
     def parameters(self):
-        """Returns a tuple containing the current value of gate's
-        parameters."""
+        """Returns a tuple containing the current value of gate's parameters."""
         return self._parameters
 
     def commutes(self, gate: "Gate") -> bool:
@@ -230,7 +228,7 @@ class Gate:
             gate: Gate to check if it commutes with the current gate.
 
         Returns:
-            ``True`` if the gates commute, otherwise ``False``.
+            bool: ``True`` if the gates commute, ``False`` otherwise.
         """
         if isinstance(gate, SpecialGate):  # pragma: no cover
             return False
@@ -297,8 +295,7 @@ class Gate:
         action of dagger will be lost.
 
         Returns:
-            A :class:`qibo.gates.Gate` object representing the dagger of
-            the original gate.
+            :class:`qibo.gates.Gate`: object representing the dagger of the original gate.
         """
         new_gate = self._dagger()
         new_gate.is_controlled_by = self.is_controlled_by
@@ -322,12 +319,24 @@ class Gate:
     def controlled_by(self, *qubits: int) -> "Gate":
         """Controls the gate on (arbitrarily many) qubits.
 
+        To see how this method affects the underlying matrix representation of a gate,
+        please see the documentation of :meth:`qibo.gates.Gate.matrix`.
+
+        .. note::
+            Some gate classes default to another gate class depending on the number of controls
+            present. For instance, an :math:`1`-controlled :class:`qibo.gates.X` gate
+            will default to a :class:`qibo.gates.CNOT` gate, while a :math:`2`-controlled
+            :class:`qibo.gates.X` gate defaults to a :class:`qibo.gates.TOFFOLI` gate.
+            Other gates affected by this method are: :class:`qibo.gates.Y`, :class:`qibo.gates.Z`,
+            :class:`qibo.gates.RX`, :class:`qibo.gates.RY`, :class:`qibo.gates.RZ`,
+            :class:`qibo.gates.U1`, :class:`qibo.gates.U2`, and :class:`qibo.gates.U3`.
+
         Args:
             *qubits (int): Ids of the qubits that the gate will be controlled on.
 
         Returns:
-            A :class:`qibo.gates.Gate` object in with the corresponding
-            gate being controlled in the given qubits.
+            :class:`qibo.gates.Gate`: object in with the corresponding
+                gate being controlled in the given qubits.
         """
         if qubits:
             self.is_controlled_by = True
@@ -343,7 +352,7 @@ class Gate:
             free: Ids of free qubits to use for the gate decomposition.
 
         Returns:
-            List with gates that have the same effect as applying the original gate.
+            list: gates that have the same effect as applying the original gate.
         """
         # TODO: Implement this method for all gates not supported by OpenQASM.
         # Currently this is implemented only for multi-controlled X gates.
@@ -354,6 +363,28 @@ class Gate:
     def matrix(self, backend=None):
         """Returns the matrix representation of the gate.
 
+        If gate has controlled qubits inserted by :meth:`qibo.gates.Gate.controlled_by`,
+        then :meth:`qibo.gates.Gate.matrix` returns the matrix of the original gate.
+
+        .. code-block:: python
+
+            from qibo import gates
+
+            gate = gates.SWAP(3, 4).controlled_by(0, 1, 2)
+            print(gate.matrix())
+
+        To return the full matrix that takes the control qubits into account,
+        one should use :meth:`qibo.models.Circuit.unitary`, e.g.
+
+        .. code-block:: python
+
+            from qibo import Circuit, gates
+
+            nqubits = 5
+            circuit = Circuit(nqubits)
+            circuit.add(gates.SWAP(3, 4).controlled_by(0, 1, 2))
+            print(circuit.unitary())
+
         Args:
             backend (:class:`qibo.backends.abstract.Backend`, optional): backend
                 to be used in the execution. If ``None``, it uses
@@ -374,7 +405,7 @@ class Gate:
         """This function returns the eigenvalues of the gate's generator.
 
         Returns:
-            (float) eigenvalue of the generator.
+            float: eigenvalue of the generator.
         """
 
         raise_error(
diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 67ec3a670..930626cdf 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -737,13 +737,17 @@ class Circuit:
     def gates_of_type(self, gate: Union[str, type]) -> List[Tuple[int, gates.Gate]]:
         """Finds all gate objects of specific type or name.
 
+        This method can be affected by how :meth:`qibo.gates.Gate.controlled_by`
+        behaves with certain gates. To see how :meth:`qibo.gates.Gate.controlled_by`
+        affects gates, we refer to the documentation of :meth:`qibo.gates.Gate.controlled_by`.
+
         Args:
-            gate (str, type): The QASM name of a gate or the corresponding gate class.
+            gate (str or type): The name of a gate or the corresponding gate class.
 
         Returns:
-            List with all gates that are in the circuit and have the same type
-            with the given ``gate``. The list contains tuples ``(i, g)`` where
-            ``i`` is the index of the gate ``g`` in the circuit's gate queue.
+            list: gates that are in the circuit and have the same type as ``gate``.
+                The list contains tuples ``(k, g)`` where ``k`` is the index of the gate
+                ``g`` in the circuit's gate queue.
         """
         if isinstance(gate, str):
             return [(i, g) for i, g in enumerate(self.queue) if g.name == gate]


Test patch:
diff --git a/tests/test_backends.py b/tests/test_backends.py
index 876e931b2..c59654f9d 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -89,38 +89,6 @@ def test_matrix_rotations(backend, gate, target_matrix):
     backend.assert_allclose(gate.matrix(backend), target_matrix(theta))
 
 
-def test_control_matrix(backend):
-    theta = 0.1234
-    rotation = np.array(
-        [
-            [np.cos(theta / 2.0), -np.sin(theta / 2.0)],
-            [np.sin(theta / 2.0), np.cos(theta / 2.0)],
-        ]
-    )
-    target_matrix = np.eye(4, dtype=rotation.dtype)
-    target_matrix[2:, 2:] = rotation
-    gate = gates.RY(0, theta).controlled_by(1)
-    backend.assert_allclose(gate.matrix(backend), target_matrix)
-
-    gate = gates.RY(0, theta).controlled_by(1, 2)
-    with pytest.raises(NotImplementedError):
-        matrix = backend.control_matrix(gate)
-
-
-def test_control_matrix_unitary(backend):
-    u = np.random.random((2, 2))
-    gate = gates.Unitary(u, 0).controlled_by(1)
-    matrix = backend.control_matrix(gate)
-    target_matrix = np.eye(4, dtype=np.complex128)
-    target_matrix[2:, 2:] = u
-    backend.assert_allclose(matrix, target_matrix)
-
-    u = np.random.random((16, 16))
-    gate = gates.Unitary(u, 0, 1, 2, 3).controlled_by(4)
-    with pytest.raises(ValueError):
-        matrix = backend.control_matrix(gate)
-
-
 def test_plus_density_matrix(backend):
     matrix = backend.plus_density_matrix(4)
     target_matrix = np.ones((16, 16)) / 16
diff --git a/tests/test_gates_gates.py b/tests/test_gates_gates.py
index 372bf0b20..34b215129 100644
--- a/tests/test_gates_gates.py
+++ b/tests/test_gates_gates.py
@@ -521,8 +521,7 @@ def test_cnot(backend, applyx):
 
 @pytest.mark.parametrize("seed_observable", list(range(1, 10 + 1)))
 @pytest.mark.parametrize("seed_state", list(range(1, 10 + 1)))
-@pytest.mark.parametrize("controlled_by", [False, True])
-def test_cy(backend, controlled_by, seed_state, seed_observable):
+def test_cy(backend, seed_state, seed_observable):
     nqubits = 2
     initial_state = random_statevector(2**nqubits, seed=seed_state, backend=backend)
     matrix = np.array(
@@ -544,15 +543,10 @@ def test_cy(backend, controlled_by, seed_state, seed_observable):
         initial_state=initial_state,
     )
 
-    if controlled_by:
-        gate = gates.Y(1).controlled_by(0)
-    else:
-        gate = gates.CY(0, 1)
+    gate = gates.CY(0, 1)
 
     final_state = apply_gates(backend, [gate], initial_state=initial_state)
 
-    assert gate.name == "cy"
-
     backend.assert_allclose(final_state, target_state)
 
     # testing random expectation value due to global phase difference
@@ -566,15 +560,15 @@ def test_cy(backend, controlled_by, seed_state, seed_observable):
         @ backend.cast(target_state),
     )
 
-    assert gates.CY(0, 1).qasm_label == "cy"
-    assert gates.CY(0, 1).clifford
-    assert gates.CY(0, 1).unitary
+    assert gate.name == "cy"
+    assert gate.qasm_label == "cy"
+    assert gate.clifford
+    assert gate.unitary
 
 
 @pytest.mark.parametrize("seed_observable", list(range(1, 10 + 1)))
 @pytest.mark.parametrize("seed_state", list(range(1, 10 + 1)))
-@pytest.mark.parametrize("controlled_by", [False, True])
-def test_cz(backend, controlled_by, seed_state, seed_observable):
+def test_cz(backend, seed_state, seed_observable):
     nqubits = 2
     initial_state = random_statevector(2**nqubits, seed=seed_state, backend=backend)
     matrix = np.eye(4)
@@ -590,15 +584,10 @@ def test_cz(backend, controlled_by, seed_state, seed_observable):
         initial_state=initial_state,
     )
 
-    if controlled_by:
-        gate = gates.Z(1).controlled_by(0)
-    else:
-        gate = gates.CZ(0, 1)
+    gate = gates.CZ(0, 1)
 
     final_state = apply_gates(backend, [gate], initial_state=initial_state)
 
-    assert gate.name == "cz"
-
     backend.assert_allclose(final_state, target_state)
 
     # testing random expectation value due to global phase difference
@@ -612,9 +601,10 @@ def test_cz(backend, controlled_by, seed_state, seed_observable):
         @ backend.cast(target_state),
     )
 
-    assert gates.CZ(0, 1).qasm_label == "cz"
-    assert gates.CZ(0, 1).clifford
-    assert gates.CZ(0, 1).unitary
+    assert gate.name == "cz"
+    assert gate.qasm_label == "cz"
+    assert gate.clifford
+    assert gate.unitary
 
 
 def test_csx(backend):
@@ -1457,8 +1447,6 @@ def test_controlled_u1(backend):
     target_state = np.zeros_like(final_state)
     target_state[1] = np.exp(1j * theta)
     backend.assert_allclose(final_state, target_state)
-    gate = gates.U1(0, theta).controlled_by(1)
-    assert gate.__class__.__name__ == "CU1"
 
 
 def test_controlled_u2(backend):
@@ -1574,23 +1562,6 @@ def test_controlled_unitary(backend):
     backend.assert_allclose(final_state, target_state)
 
 
-def test_controlled_unitary_matrix(backend):
-    nqubits = 2
-    initial_state = random_statevector(2**nqubits, backend=backend)
-
-    matrix = np.random.random((2, 2))
-    gate = gates.Unitary(matrix, 1).controlled_by(0)
-
-    target_state = apply_gates(backend, [gate], nqubits, initial_state)
-
-    u = backend.control_matrix(gate)
-    u = backend.cast(u, dtype=u.dtype)
-
-    final_state = np.dot(u, initial_state)
-
-    backend.assert_allclose(final_state, target_state)
-
-
 ###############################################################################
 
 ################################# Test dagger #################################
diff --git a/tests/test_models_circuit_features.py b/tests/test_models_circuit_features.py
index f20b983c9..b8a1ade5e 100644
--- a/tests/test_models_circuit_features.py
+++ b/tests/test_models_circuit_features.py
@@ -58,6 +58,28 @@ def test_circuit_unitary_and_inverse_with_noise_channel(backend):
         circuit.invert()
 
 
+def test_circuit_unitary_non_trivial(backend):
+    target = np.array(
+        [
+            [1, 0, 0, 0, 0, 0, 0, 0],
+            [0, 1, 0, 0, 0, 0, 0, 0],
+            [0, 0, 1, 0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 1, 0],
+            [0, 0, 0, 0, 1, 0, 0, 0],
+            [0, 0, 0, 0, 0, 1, 0, 0],
+            [0, 0, 0, 1, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 1],
+        ],
+        dtype=complex,
+    )
+    target = backend.cast(target, dtype=target.dtype)
+    nqubits = 3
+    circuit = Circuit(nqubits)
+    circuit.add(gates.SWAP(0, 2).controlled_by(1))
+    unitary = circuit.unitary(backend)
+    backend.assert_allclose(unitary, target)
+
+
 @pytest.mark.parametrize("compile", [False, True])
 def test_circuit_vs_gate_execution(backend, compile):
     """Check consistency between executing circuit and stand alone gates."""