Repository: Quantum-TII/qibo (python)

Issue:
the .on_qubits() feature fails when the quantum circuit has .controlled_by() gates.
**Describe the bug**
When trying to add a `Circuit` with `.controlled_by()` gates into a bigger circuit with the `.on_qubits()` method, an error occurs.

**To Reproduce**
```Python
circ1 = Circuit(1)
circ1.add(gates.H(0))

circ2 = Circuit(2)
circ2.add(gates.CNOT(0,1))
circ2.add(gates.H(1).controlled_by(0))

circ3 = Circuit(2)
circ3.add(circ1.on_qubits(*[0]))
circ3.add(circ2.on_qubits(*range(2)))
```

Seems to only happens when `controlled_by` gates are applied, as the CNOT gate seems to go though fine when doing:

```Python
for gate in circ2.on_qubits(*range(2)):
      print(gate)
```


Pull request description:
Fix `.on_qubits()` with `.controlled_by()` gates
Fixes #348. I added some tests that cover the problematic case but it would be good to also extensively test the example of @igres26 before merging this.

Patch:
diff --git a/src/qibo/abstractions/abstract_gates.py b/src/qibo/abstractions/abstract_gates.py
index 09dec8cd7..8a772fe66 100644
--- a/src/qibo/abstractions/abstract_gates.py
+++ b/src/qibo/abstractions/abstract_gates.py
@@ -190,7 +190,15 @@ class Gate:
         Args:
             q (int): Qubit index (or indeces) that the new gate should act on.
         """
-        return self.__class__(*q, **self.init_kwargs)
+        if self.is_controlled_by:
+            targets = (q[i] for i in self.target_qubits)
+            controls = (q[i] for i in self.control_qubits)
+            gate = self.__class__(*targets, **self.init_kwargs)
+            gate = gate.controlled_by(*controls)
+        else:
+            qubits = (q[i] for i in self.qubits)
+            gate = self.__class__(*qubits, **self.init_kwargs)
+        return gate
 
     def _dagger(self) -> "Gate":
         """Helper method for :meth:`qibo.abstractions.gates.Gate.dagger`."""
diff --git a/src/qibo/abstractions/circuit.py b/src/qibo/abstractions/circuit.py
index d819c98b9..82b543cf4 100644
--- a/src/qibo/abstractions/circuit.py
+++ b/src/qibo/abstractions/circuit.py
@@ -180,7 +180,7 @@ class AbstractCircuit(ABC):
                                     "the circuit contains {} qubits."
                                     "".format(len(q), self.nqubits))
         for gate in self.queue:
-            yield gate.on_qubits(*(q[i] for i in gate.qubits))
+            yield gate.on_qubits(*q)
 
     def copy(self, deep: bool = False):
         """Creates a copy of the current ``circuit`` as a new ``Circuit`` model.
diff --git a/src/qibo/abstractions/gates.py b/src/qibo/abstractions/gates.py
index 050cd4c92..d1daaa456 100644
--- a/src/qibo/abstractions/gates.py
+++ b/src/qibo/abstractions/gates.py
@@ -1147,8 +1147,12 @@ class Unitary(ParametrizedGate):
 
     def on_qubits(self, *q) -> "Gate":
         args = [self.init_args[0]]
-        args.extend(q)
-        return self.__class__(*args, **self.init_kwargs)
+        args.extend((q[i] for i in self.target_qubits))
+        gate = self.__class__(*args, **self.init_kwargs)
+        if self.is_controlled_by:
+            controls = (q[i] for i in self.control_qubits)
+            gate = gate.controlled_by(*controls)
+        return gate
 
     @abstractmethod
     def _dagger(self) -> "Gate": # pragma: no cover


Test patch:
diff --git a/src/qibo/tests_new/test_abstract_circuit.py b/src/qibo/tests_new/test_abstract_circuit.py
index 06bfe8fb9..007193535 100644
--- a/src/qibo/tests_new/test_abstract_circuit.py
+++ b/src/qibo/tests_new/test_abstract_circuit.py
@@ -226,6 +226,7 @@ def test_circuit_on_qubits():
     c = Circuit(3)
     c.add([gates.H(0), gates.X(1), gates.Y(2)])
     c.add([gates.CNOT(0, 1), gates.CZ(1, 2)])
+    c.add(gates.H(1).controlled_by(0, 2))
     new_gates = list(c.on_qubits(2, 5, 4))
     assert new_gates[0].target_qubits == (2,)
     assert new_gates[1].target_qubits == (5,)
@@ -234,6 +235,8 @@ def test_circuit_on_qubits():
     assert new_gates[3].control_qubits == (2,)
     assert new_gates[4].target_qubits == (4,)
     assert new_gates[4].control_qubits == (5,)
+    assert new_gates[5].target_qubits == (5,)
+    assert new_gates[5].control_qubits == (2, 4)
 
 
 def test_circuit_on_qubits_errors():
diff --git a/src/qibo/tests_new/test_abstract_gates.py b/src/qibo/tests_new/test_abstract_gates.py
index 3eceb828e..08d947010 100644
--- a/src/qibo/tests_new/test_abstract_gates.py
+++ b/src/qibo/tests_new/test_abstract_gates.py
@@ -392,6 +392,15 @@ def test_controlled_by():
         gate = gate.controlled_by(2)
 
 
+def test_on_qubits_controlled_by():
+    gate = gates.H(0).controlled_by(1, 2)
+    gate = gate.on_qubits(5, 4, 6)
+    assert gate.target_qubits == (5,)
+    assert gate.control_qubits == (4, 6)
+    assert isinstance(gate, gates.H)
+    assert gate.is_controlled_by
+
+
 def test_decompose():
     decomp_gates = gates.H(0).decompose(1)
     assert len(decomp_gates) == 1
diff --git a/src/qibo/tests_new/test_core_circuit_features.py b/src/qibo/tests_new/test_core_circuit_features.py
index 37e6af019..efb494319 100644
--- a/src/qibo/tests_new/test_core_circuit_features.py
+++ b/src/qibo/tests_new/test_core_circuit_features.py
@@ -181,26 +181,66 @@ def test_circuit_on_qubits_double_execution(backend, accelerators, distribute_sm
     qibo.set_backend(original_backend)
 
 
-def test_circuit_on_qubits_with_unitary_execution(backend, accelerators):
+def test_circuit_on_qubits_controlled_by_execution(backend, accelerators):
+    original_backend = qibo.get_backend()
+    qibo.set_backend(backend)
+
+    smallc = Circuit(3)
+    smallc.add(gates.RX(0, theta=0.1).controlled_by(1, 2))
+    smallc.add(gates.RY(1, theta=0.2).controlled_by(0))
+    smallc.add(gates.RX(2, theta=0.3).controlled_by(1, 0))
+    smallc.add(gates.RZ(1, theta=0.4).controlled_by(0, 2))
+
+    largec = Circuit(6, accelerators=accelerators)
+    largec.add(gates.H(i) for i in range(6))
+    largec.add(smallc.on_qubits(1, 4, 3))
+
+    targetc = Circuit(6)
+    targetc.add(gates.H(i) for i in range(6))
+    targetc.add(gates.RX(1, theta=0.1).controlled_by(3, 4))
+    targetc.add(gates.RY(4, theta=0.2).controlled_by(1))
+    targetc.add(gates.RX(3, theta=0.3).controlled_by(1, 4))
+    targetc.add(gates.RZ(4, theta=0.4).controlled_by(1, 3))
+
+    assert largec.depth == targetc.depth
+    np.testing.assert_allclose(largec(), targetc())
+    qibo.set_backend(original_backend)
+
+
+@pytest.mark.parametrize("controlled", [False, True])
+def test_circuit_on_qubits_with_unitary_execution(backend, accelerators, controlled):
     original_backend = qibo.get_backend()
     qibo.set_backend(backend)
 
     unitaries = np.random.random((2, 2, 2))
     smallc = Circuit(2)
-    smallc.add((gates.Unitary(u, i) for i, u in enumerate(unitaries)))
+    if controlled:
+        smallc.add(gates.Unitary(unitaries[0], 0).controlled_by(1))
+        smallc.add(gates.Unitary(unitaries[1], 1).controlled_by(0))
+    else:
+        smallc.add(gates.Unitary(unitaries[0], 0))
+        smallc.add(gates.Unitary(unitaries[1], 1))
     smallc.add(gates.CNOT(0, 1))
 
     largec = Circuit(4, accelerators=accelerators)
-    largec.add(gates.RY(1, theta=0.1))
-    largec.add(gates.RY(2, theta=0.2))
-    largec.add(smallc.on_qubits(0, 3))
+    largec.add(gates.RY(0, theta=0.1))
+    largec.add(gates.RY(1, theta=0.2))
+    largec.add(gates.RY(2, theta=0.3))
+    largec.add(gates.RY(3, theta=0.2))
+    largec.add(smallc.on_qubits(3, 0))
 
     targetc = Circuit(4)
-    targetc.add(gates.RY(1, theta=0.1))
-    targetc.add(gates.RY(2, theta=0.2))
-    targetc.add(gates.Unitary(unitaries[0], 0))
-    targetc.add(gates.Unitary(unitaries[1], 3))
-    targetc.add(gates.CNOT(0, 3))
+    targetc.add(gates.RY(0, theta=0.1))
+    targetc.add(gates.RY(1, theta=0.2))
+    targetc.add(gates.RY(2, theta=0.3))
+    targetc.add(gates.RY(3, theta=0.2))
+    if controlled:
+        targetc.add(gates.Unitary(unitaries[0], 3).controlled_by(0))
+        targetc.add(gates.Unitary(unitaries[1], 0).controlled_by(3))
+    else:
+        targetc.add(gates.Unitary(unitaries[0], 3))
+        targetc.add(gates.Unitary(unitaries[1], 0))
+    targetc.add(gates.CNOT(3, 0))
     assert largec.depth == targetc.depth
     np.testing.assert_allclose(largec(), targetc())
     qibo.set_backend(original_backend)