Repository: qiboteam/qibo

Issue:
Circuit.light_cone() resets phases of parametrized gates
Example:

```python3
from qibo import gates as gt
from qibo.models import Circuit
from numpy import asarray, pi
from numpy.random import rand

qubits = 7
depth = 2

circuit = Circuit(qubits)
for _ in range(depth):
    for qubit in range(qubits):
        circuit.add(gt.U3(qubit, 0, 0, 0))
    cnots_even = asarray([[k, k + 1] for k in range(0, qubits - 1, 2)])
    cnots_odd = asarray([[k, k + 1] for k in range(1, qubits - 1, 2)])
    for row in cnots_even:
        circuit.add(gt.CNOT(*row))
    for row in cnots_odd:
        circuit.add(gt.CNOT(*row))

phases = list()
for _ in range(depth):
    p = rand(qubits, 3)
    p[:, 0] = p[:, 0] * pi
    p[:, 1] = p[:, 1] * 2*pi
    p[:, 1] = p[:, 2] * 2*pi
    phases.append(p)

phases = asarray(phases).flatten().reshape(qubits * depth, 3)
circuit.set_parameters(phases)
print(circuit.get_parameters())

cone = circuit.copy(True)
cone.add((gt.Y(k) for k in [1, 2]))
cone = cone.light_cone(1, 2)[0]
print(cone.get_parameters())
```

I have tried `cone.set_parameters(phases)` after the deep copy, but it doesn't solve the issue. No matter what, phases are set to 0 after `light_cone()`

Comments:

Patch:
diff --git a/src/qibo/gates/abstract.py b/src/qibo/gates/abstract.py
index 00c79562b..831f7a379 100644
--- a/src/qibo/gates/abstract.py
+++ b/src/qibo/gates/abstract.py
@@ -340,6 +340,11 @@ class ParametrizedGate(Gate):
         for gate in self.device_gates:  # pragma: no cover
             gate.parameters = x
 
+    def on_qubits(self, qubit_map):
+        gate = super().on_qubits(qubit_map)
+        gate.parameters = self.parameters
+        return gate
+
     def substitute_symbols(self):
         params = list(self._parameters)
         for i, param in self.symbolic_parameters.items():
diff --git a/src/qibo/gates/gates.py b/src/qibo/gates/gates.py
index e08b185c6..0fb66e0e3 100644
--- a/src/qibo/gates/gates.py
+++ b/src/qibo/gates/gates.py
@@ -1264,6 +1264,7 @@ class Unitary(ParametrizedGate):
         if self.is_controlled_by:
             controls = (qubit_map.get(i) for i in self.control_qubits)
             gate = gate.controlled_by(*controls)
+        gate.parameters = self.parameters
         return gate
 
     def _dagger(self):


Test patch:
diff --git a/src/qibo/tests/test_models_circuit_parametrized.py b/src/qibo/tests/test_models_circuit_parametrized.py
index 6fb28f60f..3314c49c7 100644
--- a/src/qibo/tests/test_models_circuit_parametrized.py
+++ b/src/qibo/tests/test_models_circuit_parametrized.py
@@ -190,6 +190,27 @@ def test_set_parameters_with_gate_fusion(backend, trainable):
     backend.assert_circuitclose(fused_c, c)
 
 
+@pytest.mark.parametrize("trainable", [True, False])
+def test_set_parameters_with_light_cone(backend, trainable):
+    """Check updating parameters of light cone circuit."""
+    params = np.random.random(4)
+    c = Circuit(4)
+    c.add(gates.RX(0, theta=params[0], trainable=trainable))
+    c.add(gates.RY(1, theta=params[1]))
+    c.add(gates.CZ(0, 1))
+    c.add(gates.RX(2, theta=params[2]))
+    c.add(gates.RY(3, theta=params[3], trainable=trainable))
+    c.add(gates.CZ(2, 3))
+    if trainable:
+        c.set_parameters(np.random.random(4))
+    else:
+        c.set_parameters(np.random.random(2))
+    target_state = backend.execute_circuit(c)
+    lc, _ = c.light_cone(1, 2)
+    final_state = backend.execute_circuit(lc)
+    backend.assert_allclose(final_state, target_state)
+
+
 def test_variable_theta():
     """Check that parametrized gates accept `tf.Variable` parameters."""
     try: