Repository: qiboteam/qibo

Issue:
Noise model doesn't work with circuits including measurements
Currently when generating the noisy circuit from a `NoiseModel` the measurements gate are removed.

```python
from qibo import gates, models, set_backend
from qibo.noise import PauliError, ThermalRelaxationError, NoiseModel, ResetError

pauli = PauliError(0, 0.2, 0.3)
thermal = ThermalRelaxationError(2, 1.5, 0.3)
noise = NoiseModel() 
noise.add(pauli, gates.X, 1)
noise.add(thermal, gates.Z, (0,1))

circuit = models.Circuit(3)
circuit.add(gates.CNOT(0,1))
circuit.add(gates.Z(1))
circuit.add(gates.X(1))
circuit.add(gates.X(2))
circuit.add(gates.Z(2))
circuit.add(gates.M(0,1,2))

noisy_circuit = noise.apply(circuit)
print(noisy_circuit.draw())
# output
# q0: ─o────────
# q1: ─X─Z─X─PN─
# q2: ───X─Z────
```
I'll open a PR shortly to fix this issue.


Discussion:

Patch:
diff --git a/src/qibo/abstractions/circuit.py b/src/qibo/abstractions/circuit.py
index 90f54dd26..47e53f6d1 100644
--- a/src/qibo/abstractions/circuit.py
+++ b/src/qibo/abstractions/circuit.py
@@ -417,13 +417,9 @@ class AbstractCircuit(ABC):
         # Create new circuit with noise gates inside
         noisy_circuit = self.__class__(self.nqubits)
         for i, gate in enumerate(self.queue):
-            # Do not use `circuit.add` here because these gates are already
-            # added in the original circuit
-            noisy_circuit.queue.append(gate)
+            noisy_circuit.add(gate)
             for noise_gate in noise_gates[i]:
                 noisy_circuit.add(noise_gate)
-        noisy_circuit.parametrized_gates = list(self.parametrized_gates)
-        noisy_circuit.trainable_gates = list(self.trainable_gates)
         noisy_circuit.measurement_tuples = dict(self.measurement_tuples)
         noisy_circuit.measurement_gate = self.measurement_gate
         return noisy_circuit
diff --git a/src/qibo/noise.py b/src/qibo/noise.py
index b86176fa1..725c9658a 100644
--- a/src/qibo/noise.py
+++ b/src/qibo/noise.py
@@ -93,10 +93,9 @@ class NoiseModel():
                 to the initial circuit with noise gates added according
                 to the noise model.
         """
-
-        circ = circuit.__class__(**circuit.init_kwargs)
+        noisy_circuit = circuit.__class__(**circuit.init_kwargs)
         for gate in circuit.queue:
-            circ.add(gate)
+            noisy_circuit.add(gate)
             if gate.__class__ in self.errors:
                 error, qubits = self.errors.get(gate.__class__)
                 if qubits is None:
@@ -104,5 +103,7 @@ class NoiseModel():
                 else:
                     qubits = tuple(set(gate.qubits) & set(qubits))
                 for q in qubits:
-                    circ.add(error.channel(q, *error.options))
-        return circ
\ No newline at end of file
+                    noisy_circuit.add(error.channel(q, *error.options))
+        noisy_circuit.measurement_tuples = dict(circuit.measurement_tuples)
+        noisy_circuit.measurement_gate = circuit.measurement_gate
+        return noisy_circuit
\ No newline at end of file


Test patch:
diff --git a/src/qibo/tests/test_abstract_circuit.py b/src/qibo/tests/test_abstract_circuit.py
index e10abdef1..cb62be68f 100644
--- a/src/qibo/tests/test_abstract_circuit.py
+++ b/src/qibo/tests/test_abstract_circuit.py
@@ -1,6 +1,6 @@
 """Test all methods defined in `qibo/abstractions/circuit.py`."""
 import pytest
-import qibo
+from qibo import __version__
 from qibo.abstractions import gates
 from qibo.config import raise_error
 from qibo.abstractions.circuit import AbstractCircuit
@@ -267,7 +267,7 @@ def test_circuit_light_cone():
     c.add(gates.CZ(i, i + 1) for i in range(1, nqubits - 1, 2))
     c.add(gates.CZ(0, nqubits - 1))
     sc, qubit_map = c.light_cone(4, 5)
-    target_qasm = f"""// Generated by QIBO {qibo.__version__}
+    target_qasm = f"""// Generated by QIBO {__version__}
 OPENQASM 2.0;
 include "qelib1.inc";
 qreg q[6];
diff --git a/src/qibo/tests/test_noise.py b/src/qibo/tests/test_noise.py
index 34a83ec2d..18d44b9ee 100644
--- a/src/qibo/tests/test_noise.py
+++ b/src/qibo/tests/test_noise.py
@@ -6,7 +6,8 @@ from qibo.noise import *
 from qibo.tests.utils import random_density_matrix, random_state
 
 @pytest.mark.parametrize("density_matrix", [False, True])
-def test_pauli_error(backend, density_matrix):
+@pytest.mark.parametrize("nshots", [None, 10, 100])
+def test_pauli_error(backend, density_matrix, nshots):
 
     pauli = PauliError(0, 0.2, 0.3)
     noise = NoiseModel()
@@ -20,6 +21,7 @@ def test_pauli_error(backend, density_matrix):
     circuit.add(gates.X(1))
     circuit.add(gates.X(2))
     circuit.add(gates.Z(2))
+    circuit.add(gates.M(0, 1, 2))
 
     target_circuit = Circuit(3, density_matrix=density_matrix)
     target_circuit.add(gates.CNOT(0,1))
@@ -31,22 +33,28 @@ def test_pauli_error(backend, density_matrix):
     target_circuit.add(gates.PauliNoiseChannel(1, 0, 0.2, 0.3))
     target_circuit.add(gates.X(2))
     target_circuit.add(gates.Z(2))
+    target_circuit.add(gates.M(0, 1, 2))
 
     initial_psi = random_density_matrix(3) if density_matrix else random_state(3)
     np.random.seed(123)
     K.set_seed(123)
-    final_state = noise.apply(circuit)(initial_state=np.copy(initial_psi))
+    final_state = noise.apply(circuit)(initial_state=np.copy(initial_psi), nshots=nshots)
+    final_state_samples = final_state.samples() if nshots else None
     np.random.seed(123)
     K.set_seed(123)
-    target_final_state = target_circuit(initial_state=np.copy(initial_psi))
+    target_final_state = target_circuit(initial_state=np.copy(initial_psi), nshots=nshots)
+    target_final_state_samples = target_final_state.samples() if nshots else None
 
-    K.assert_allclose(final_state, target_final_state)
+    if nshots is None:
+        K.assert_allclose(final_state, target_final_state)
+    else:
+        K.assert_allclose(final_state_samples, target_final_state_samples)
 
 
 @pytest.mark.parametrize("density_matrix", [False, True])
 def test_thermal_error(backend, density_matrix):
 
-    thermal = ThermalRelaxationError(1, 1, 0.3)
+    thermal = ThermalRelaxationError(2, 1, 0.3)
     noise = NoiseModel()
     noise.add(thermal, gates.X, 1)
     noise.add(thermal, gates.CNOT)
@@ -61,12 +69,12 @@ def test_thermal_error(backend, density_matrix):
 
     target_circuit = Circuit(3, density_matrix=density_matrix)
     target_circuit.add(gates.CNOT(0,1))
-    target_circuit.add(gates.ThermalRelaxationChannel(0, 1, 1, 0.3))
-    target_circuit.add(gates.ThermalRelaxationChannel(1, 1, 1, 0.3))
+    target_circuit.add(gates.ThermalRelaxationChannel(0, 2, 1, 0.3))
+    target_circuit.add(gates.ThermalRelaxationChannel(1, 2, 1, 0.3))
     target_circuit.add(gates.Z(1))
-    target_circuit.add(gates.ThermalRelaxationChannel(1, 1, 1, 0.3))
+    target_circuit.add(gates.ThermalRelaxationChannel(1, 2, 1, 0.3))
     target_circuit.add(gates.X(1))
-    target_circuit.add(gates.ThermalRelaxationChannel(1, 1, 1, 0.3))
+    target_circuit.add(gates.ThermalRelaxationChannel(1, 2, 1, 0.3))
     target_circuit.add(gates.X(2))
     target_circuit.add(gates.Z(2))