From 2b535775b591ba13dca2ac2aa0c7fe6ba71db5d2 Mon Sep 17 00:00:00 2001
From: Stavros <35475381+stavros11@users.noreply.github.com>
Date: Fri, 17 Feb 2023 20:38:34 +0400
Subject: [PATCH 01/10] Implement Gate.basis_rotation

---
 src/qibo/gates/abstract.py         |  6 ++++++
 src/qibo/gates/gates.py            |  9 +++++++++
 src/qibo/gates/measurements.py     | 15 ++++++++++++++-
 src/qibo/tests/test_gates_gates.py | 16 ++++++++++++++++
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/qibo/gates/abstract.py b/src/qibo/gates/abstract.py
index 42fb80445c..135605e16a 100644
--- a/src/qibo/gates/abstract.py
+++ b/src/qibo/gates/abstract.py
@@ -270,6 +270,12 @@ def generator_eigenvalue(self):
             f"Generator eigenvalue is not implemented for {self.name}",
         )
 
+    def basis_rotation(self):
+        """Transformation required to rotate the basis for measuring the gate."""
+        raise_error(
+            NotImplementedError, f"Basis rotation is not implemented for {self.name}"
+        )
+
     @property
     def matrix(self):
         from qibo.backends import GlobalBackend
diff --git a/src/qibo/gates/gates.py b/src/qibo/gates/gates.py
index a455988cf1..d296083d41 100644
--- a/src/qibo/gates/gates.py
+++ b/src/qibo/gates/gates.py
@@ -215,6 +215,9 @@ def decompose(self, *free, use_toffolis=True):
         decomp_gates.extend(decomp_gates)
         return decomp_gates
 
+    def basis_rotation(self):
+        return H(self.target_qubits[0])
+
 
 class Y(Gate):
     """The Pauli Y gate.
@@ -229,6 +232,12 @@ def __init__(self, q):
         self.target_qubits = (q,)
         self.init_args = [q]
 
+    def basis_rotation(self):
+        from qibo import matrices
+
+        matrix = (matrices.Y + matrices.Z) / math.sqrt(2)
+        return Unitary(matrix, self.target_qubits[0], trainable=False)
+
 
 class Z(Gate):
     """The Pauli Z gate.
diff --git a/src/qibo/gates/measurements.py b/src/qibo/gates/measurements.py
index 834fde706c..b2da8e19a6 100644
--- a/src/qibo/gates/measurements.py
+++ b/src/qibo/gates/measurements.py
@@ -2,11 +2,12 @@
 
 from qibo.config import raise_error
 from qibo.gates.abstract import Gate
+from qibo.gates.gates import Z
 from qibo.states import MeasurementResult
 
 
 class M(Gate):
-    """The Measure Z gate.
+    """The measure gate.
 
     Args:
         *q (int): id numbers of the qubits to measure.
@@ -20,6 +21,9 @@ class M(Gate):
             performed. Can be used only for single shot measurements.
             If ``True`` the collapsed state vector is returned. If ``False``
             the measurement result is returned.
+        basis (:class:`qibo.gates.Gate`): Basis to measure. Can be a qibo gate
+            or a callable that accepts a qubit, for example: ``lambda q: gates.RX(q, 0.2)``
+            Default is Z.
         p0 (dict): Optional bitflip probability map. Can be:
             A dictionary that maps each measured qubit to the probability
             that it is flipped, a list or tuple that has the same length
@@ -37,6 +41,7 @@ def __init__(
         *q,
         register_name: Optional[str] = None,
         collapse: bool = False,
+        basis: Gate = Z,
         p0: Optional["ProbsType"] = None,
         p1: Optional["ProbsType"] = None,
     ):
@@ -70,6 +75,14 @@ def __init__(
             p0 = p1
         self.bitflip_map = (self._get_bitflip_map(p0), self._get_bitflip_map(p1))
 
+        # list of gates that will be added to the circuit before the
+        # measurement, in order to rotate to the given basis
+        self.basis = []
+        for q in self.target_qubits:
+            gate = basis(q).basis_rotation()
+            if gate is not None:
+                self.basis.append(gate)
+
     @staticmethod
     def _get_bitflip_tuple(qubits: Tuple[int], probs: "ProbsType") -> Tuple[float]:
         if isinstance(probs, float):
diff --git a/src/qibo/tests/test_gates_gates.py b/src/qibo/tests/test_gates_gates.py
index 82f12bb8c7..73064b68b6 100644
--- a/src/qibo/tests/test_gates_gates.py
+++ b/src/qibo/tests/test_gates_gates.py
@@ -774,6 +774,22 @@ def test_generalizedfsim_dagger(backend):
     backend.assert_allclose(final_state, initial_state)
 
 
+###############################################################################
+
+############################# Test basis rotation #############################
+
+
+def test_gate_basis_rotation(backend):
+    gate = gates.X(0).basis_rotation()
+    assert isinstance(gate, gates.H)
+    gate = gates.Y(0).basis_rotation()
+    assert isinstance(gate, gates.Unitary)
+    target_matrix = np.array([[1, -1j], [1j, -1]]) / np.sqrt(2)
+    backend.assert_allclose(gate.asmatrix(backend), target_matrix)
+    with pytest.raises(NotImplementedError):
+        gates.RX(0, np.pi / 2).basis_rotation()
+
+
 ###############################################################################
 
 ########################### Test gate decomposition ###########################

From a72f6f730fdc1fa65a9711f196957dfdb02788a3 Mon Sep 17 00:00:00 2001
From: Stavros <35475381+stavros11@users.noreply.github.com>
Date: Fri, 17 Feb 2023 20:47:55 +0400
Subject: [PATCH 02/10] Add basis rotation to circuit before measurement

---
 src/qibo/gates/gates.py             |  3 +++
 src/qibo/models/circuit.py          |  1 +
 src/qibo/tests/test_measurements.py | 12 ++++++++++++
 3 files changed, 16 insertions(+)

diff --git a/src/qibo/gates/gates.py b/src/qibo/gates/gates.py
index d296083d41..548d560277 100644
--- a/src/qibo/gates/gates.py
+++ b/src/qibo/gates/gates.py
@@ -261,6 +261,9 @@ def controlled_by(self, *q):
             gate = super().controlled_by(*q)
         return gate
 
+    def basis_rotation(self):
+        return None
+
 
 class S(Gate):
     """The S gate.
diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 1cfcc80464..8f0c77010f 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -561,6 +561,7 @@ def add(self, gate):
 
             self.queue.append(gate)
             if isinstance(gate, gates.M):
+                self.add(gate.basis)
                 if gate.register_name is None:
                     # add default register name
                     nreg = self.queue.nmeasurements - 1
diff --git a/src/qibo/tests/test_measurements.py b/src/qibo/tests/test_measurements.py
index a21a2bbf71..1f77d12046 100644
--- a/src/qibo/tests/test_measurements.py
+++ b/src/qibo/tests/test_measurements.py
@@ -435,3 +435,15 @@ def test_measurement_result_vs_circuit_result(backend, accelerators):
     frequencies = result.frequencies(registers=True)
     assert ma_freq == frequencies.get("a")
     assert mb_freq == frequencies.get("b")
+
+
+@pytest.mark.parametrize("nqubits", [1, 4])
+@pytest.mark.parametrize("outcome", [0, 1])
+def test_measurement_basis(backend, nqubits, outcome):
+    c = models.Circuit(nqubits)
+    if outcome:
+        c.add(gates.X(q) for q in range(nqubits))
+    c.add(gates.H(q) for q in range(nqubits))
+    c.add(gates.M(*range(nqubits), basis=gates.X))
+    result = c(nshots=100)
+    assert result.frequencies() == {nqubits * str(outcome): 100}

From fc48f93bb3b01e4093a8828742ab293e161b05cb Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Mon, 6 Mar 2023 10:23:59 +0400
Subject: [PATCH 03/10] Add circuit as initial state

---
 src/qibo/backends/numpy.py                      |  3 +++
 src/qibo/tests/test_models_circuit_execution.py | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 639cc68475..ea3d5520c9 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -331,6 +331,9 @@ def thermal_error_density_matrix(self, gate, state, nqubits):
     def execute_circuit(
         self, circuit, initial_state=None, nshots=None, return_array=False
     ):
+        if isinstance(initial_state, type(circuit)):
+            return self.execute_circuit(initial_state, None, nshots)
+
         if circuit.repeated_execution:
             return self.execute_circuit_repeated(circuit, initial_state, nshots)
 
diff --git a/src/qibo/tests/test_models_circuit_execution.py b/src/qibo/tests/test_models_circuit_execution.py
index b330057e17..d0c9966117 100644
--- a/src/qibo/tests/test_models_circuit_execution.py
+++ b/src/qibo/tests/test_models_circuit_execution.py
@@ -106,3 +106,18 @@ def test_density_matrix_circuit(backend):
     target_rho = m2.dot(target_rho).dot(m2.T.conj())
     target_rho = m3.dot(target_rho).dot(m3.T.conj())
     backend.assert_allclose(final_rho, target_rho)
+
+
+def test_circuit_initial_state(backend):
+    c = Circuit(2)
+    c.add(gates.H(0))
+    c.add(gates.X(1))
+
+    c1 = Circuit(2)
+    c1.add(gates.H(0))
+    c1.add(gates.H(1))
+
+    c_final_state = backend.execute_circuit(c, c1)
+    c1_final_state = backend.execute_circuit(c1)
+
+    backend.assert_allclose(c_final_state, c1_final_state)

From e8ccd7ea404f9cc7cbdb51c86772cb05508e2d59 Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Mon, 6 Mar 2023 12:14:40 +0400
Subject: [PATCH 04/10] Fix behavior

---
 src/qibo/backends/numpy.py                    |  9 ++++-
 .../tests/test_models_circuit_execution.py    | 34 +++++++++++++------
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index ea3d5520c9..37e05e2c1b 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -332,7 +332,14 @@ def execute_circuit(
         self, circuit, initial_state=None, nshots=None, return_array=False
     ):
         if isinstance(initial_state, type(circuit)):
-            return self.execute_circuit(initial_state, None, nshots)
+            if not initial_state.density_matrix == circuit.density_matrix:
+                raise_error(
+                    ValueError,
+                    f"""Cannot set circuit with density_matrix {initial_state.density_matrix} a
+                      initial state for circuit with density_matrix {circuit.density_matrix}.""",
+                )
+            else:
+                return self.execute_circuit(initial_state + circuit, None, nshots)
 
         if circuit.repeated_execution:
             return self.execute_circuit_repeated(circuit, initial_state, nshots)
diff --git a/src/qibo/tests/test_models_circuit_execution.py b/src/qibo/tests/test_models_circuit_execution.py
index d0c9966117..33f5b78ae8 100644
--- a/src/qibo/tests/test_models_circuit_execution.py
+++ b/src/qibo/tests/test_models_circuit_execution.py
@@ -108,16 +108,30 @@ def test_density_matrix_circuit(backend):
     backend.assert_allclose(final_rho, target_rho)
 
 
-def test_circuit_initial_state(backend):
-    c = Circuit(2)
-    c.add(gates.H(0))
-    c.add(gates.X(1))
+@pytest.mark.parametrize("density_matrix", [True, False])
+def test_circuit_as_initial_state(backend, density_matrix):
+    nqubits = 10
+    c = Circuit(nqubits, density_matrix=density_matrix)
+    c.add(gates.X(i) for i in range(nqubits))
+
+    c1 = Circuit(nqubits, density_matrix=density_matrix)
+    c1.add(gates.H(i) for i in range(nqubits))
+
+    actual_circuit = c1 + c
+
+    output = backend.execute_circuit(c, c1)
+    target = backend.execute_circuit(actual_circuit)
+
+    backend.assert_allclose(target, output)
+
 
-    c1 = Circuit(2)
-    c1.add(gates.H(0))
-    c1.add(gates.H(1))
+def test_initial_state_error(backend):
+    nqubits = 10
+    c = Circuit(nqubits)
+    c.add(gates.X(i) for i in range(nqubits))
 
-    c_final_state = backend.execute_circuit(c, c1)
-    c1_final_state = backend.execute_circuit(c1)
+    c1 = Circuit(nqubits, density_matrix=True)
+    c1.add(gates.H(i) for i in range(nqubits))
 
-    backend.assert_allclose(c_final_state, c1_final_state)
+    with pytest.raises(ValueError):
+        backend.execute_circuit(c, c1)

From 235a77bda6d89da6bdd8434d7c054f958138ed83 Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Mon, 6 Mar 2023 15:09:37 +0400
Subject: [PATCH 05/10] Add docstring

---
 src/qibo/models/circuit.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 8f0c77010f..e7ffdb4a83 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -967,8 +967,11 @@ def compile(self, backend=None):
     def execute(self, initial_state=None, nshots=None):
         """Executes the circuit. Exact implementation depends on the backend.
 
-        See :meth:`qibo.core.circuit.Circuit.execute` for more
-        details.
+        Args:
+            initial_state (`np.ndarray` or :class:`qibo.models.circuit.Circuit`): Initial configuration.
+                Can be specified by the setting the state vector using an array or a circuit. If ``None``
+                the initial state is ``|000..00>``.
+            nshots (int): Number of shots.
         """
         if self.compiled:
             # pylint: disable=E1101

From 188a5517319b3b8efdc98b2fc31d2c65ff14f31c Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Mon, 6 Mar 2023 21:23:23 +0400
Subject: [PATCH 06/10] Add check for accelerators

---
 src/qibo/backends/numpy.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 37e05e2c1b..81e3019344 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -335,9 +335,15 @@ def execute_circuit(
             if not initial_state.density_matrix == circuit.density_matrix:
                 raise_error(
                     ValueError,
-                    f"""Cannot set circuit with density_matrix {initial_state.density_matrix} a
+                    f"""Cannot set circuit with density_matrix {initial_state.density_matrix} as
                       initial state for circuit with density_matrix {circuit.density_matrix}.""",
                 )
+            elif not initial_state.accelerators == circuit.accelerators:
+                raise_error(
+                    ValueError,
+                    f"""Cannot set circuit with accelerators {initial_state.density_matrix} as
+                      initial state for circuit with accelerators {circuit.density_matrix}.""",
+                )
             else:
                 return self.execute_circuit(initial_state + circuit, None, nshots)
 

From 3f6264c2b93d01c92e87bec4d43d8c6cfa79f269 Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Tue, 7 Mar 2023 09:05:51 +0400
Subject: [PATCH 07/10] Fix coverage

---
 src/qibo/backends/numpy.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 81e3019344..25e1762b02 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -338,7 +338,9 @@ def execute_circuit(
                     f"""Cannot set circuit with density_matrix {initial_state.density_matrix} as
                       initial state for circuit with density_matrix {circuit.density_matrix}.""",
                 )
-            elif not initial_state.accelerators == circuit.accelerators:
+            elif (
+                not initial_state.accelerators == circuit.accelerators
+            ):  # pragma: no cover
                 raise_error(
                     ValueError,
                     f"""Cannot set circuit with accelerators {initial_state.density_matrix} as

From 429cc6b13207b851b695dec3c988e3ed6d8acae3 Mon Sep 17 00:00:00 2001
From: Stavros <35475381+stavros11@users.noreply.github.com>
Date: Wed, 8 Mar 2023 15:52:30 +0400
Subject: [PATCH 08/10] Implement list for basis

---
 src/qibo/gates/measurements.py | 13 +++++++++----
 tests/test_measurements.py     | 11 +++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/qibo/gates/measurements.py b/src/qibo/gates/measurements.py
index b2da8e19a6..d0153f8b44 100644
--- a/src/qibo/gates/measurements.py
+++ b/src/qibo/gates/measurements.py
@@ -21,8 +21,11 @@ class M(Gate):
             performed. Can be used only for single shot measurements.
             If ``True`` the collapsed state vector is returned. If ``False``
             the measurement result is returned.
-        basis (:class:`qibo.gates.Gate`): Basis to measure. Can be a qibo gate
-            or a callable that accepts a qubit, for example: ``lambda q: gates.RX(q, 0.2)``
+        basis (:class:`qibo.gates.Gate`, list): Basis to measure.
+            Can be a qibo gate or a callable that accepts a qubit,
+            for example: ``lambda q: gates.RX(q, 0.2)``
+            or a list of these, if a different basis will be used for each
+            measurement qubit.
             Default is Z.
         p0 (dict): Optional bitflip probability map. Can be:
             A dictionary that maps each measured qubit to the probability
@@ -77,9 +80,11 @@ def __init__(
 
         # list of gates that will be added to the circuit before the
         # measurement, in order to rotate to the given basis
+        if not isinstance(basis, list):
+            basis = len(self.target_qubits) * [basis]
         self.basis = []
-        for q in self.target_qubits:
-            gate = basis(q).basis_rotation()
+        for qubit, basis_cls in zip(self.target_qubits, basis):
+            gate = basis_cls(qubit).basis_rotation()
             if gate is not None:
                 self.basis.append(gate)
 
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index 1f77d12046..feb79dbb08 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -447,3 +447,14 @@ def test_measurement_basis(backend, nqubits, outcome):
     c.add(gates.M(*range(nqubits), basis=gates.X))
     result = c(nshots=100)
     assert result.frequencies() == {nqubits * str(outcome): 100}
+
+
+def test_measurement_basis_list(backend):
+    c = models.Circuit(4)
+    c.add(gates.H(0))
+    c.add(gates.X(2))
+    c.add(gates.H(2))
+    c.add(gates.X(3))
+    c.add(gates.M(0, 1, 2, 3, basis=[gates.X, gates.Z, gates.X, gates.Z]))
+    result = c(nshots=100)
+    assert result.frequencies() == {"0011": 100}

From 4d33ee82a0d415cca850d9b61a9ddb602c43f7c9 Mon Sep 17 00:00:00 2001
From: Stavros <35475381+stavros11@users.noreply.github.com>
Date: Wed, 8 Mar 2023 15:58:27 +0400
Subject: [PATCH 09/10] Check length and raise error

---
 src/qibo/gates/measurements.py | 6 ++++++
 tests/test_measurements.py     | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/src/qibo/gates/measurements.py b/src/qibo/gates/measurements.py
index d0153f8b44..a567a66ee1 100644
--- a/src/qibo/gates/measurements.py
+++ b/src/qibo/gates/measurements.py
@@ -82,6 +82,12 @@ def __init__(
         # measurement, in order to rotate to the given basis
         if not isinstance(basis, list):
             basis = len(self.target_qubits) * [basis]
+        elif len(basis) != len(self.target_qubits):
+            raise_error(
+                ValueError,
+                f"Given basis list has length {len(basis)} while "
+                f"we are measuring {len(self.target_qubits)} qubits.",
+            )
         self.basis = []
         for qubit, basis_cls in zip(self.target_qubits, basis):
             gate = basis_cls(qubit).basis_rotation()
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index feb79dbb08..3d0a7555ab 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -458,3 +458,9 @@ def test_measurement_basis_list(backend):
     c.add(gates.M(0, 1, 2, 3, basis=[gates.X, gates.Z, gates.X, gates.Z]))
     result = c(nshots=100)
     assert result.frequencies() == {"0011": 100}
+
+
+def test_measurement_basis_list_error(backend):
+    c = models.Circuit(4)
+    with pytest.raises(ValueError):
+        c.add(gates.M(0, 1, 2, 3, basis=[gates.X, gates.Z, gates.X]))

From 379886f075f207d81b85399ed55e11e839c87fba Mon Sep 17 00:00:00 2001
From: Stavros <35475381+stavros11@users.noreply.github.com>
Date: Wed, 8 Mar 2023 16:22:29 +0400
Subject: [PATCH 10/10] Fix basis gate order

---
 src/qibo/models/circuit.py | 3 ++-
 tests/test_measurements.py | 8 ++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 182f0e9940..0e6f983321 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -575,9 +575,9 @@ def add(self, gate):
                         "".format(gate.target_qubits, self.nqubits),
                     )
 
-            self.queue.append(gate)
             if isinstance(gate, gates.M):
                 self.add(gate.basis)
+                self.queue.append(gate)
                 if gate.register_name is None:
                     # add default register name
                     nreg = self.queue.nmeasurements - 1
@@ -598,6 +598,7 @@ def add(self, gate):
                 return gate.result
 
             else:
+                self.queue.append(gate)
                 for measurement in list(self.measurements):
                     if set(measurement.qubits) & set(gate.qubits):
                         measurement.collapse = True
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index 3d0a7555ab..276b6e67b1 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -458,6 +458,14 @@ def test_measurement_basis_list(backend):
     c.add(gates.M(0, 1, 2, 3, basis=[gates.X, gates.Z, gates.X, gates.Z]))
     result = c(nshots=100)
     assert result.frequencies() == {"0011": 100}
+    print(c.draw())
+    assert (
+        c.draw()
+        == """q0: ─H─H───M─
+q1: ───────M─
+q2: ─X─H─H─M─
+q3: ─X─────M─"""
+    )
 
 
 def test_measurement_basis_list_error(backend):