Repository: qiboteam/qibo

Issue:
Mesurement and initialization in Pauli basis
I would like to suggest including a feature for measuring in a different Pauli basis besides the computational basis, i.e. include an argument in the Measurement class `qibo.gates.M(..., basis,...)` such that

> If basis = 'Z', pass
> If basis = 'X', add a Hadamard gate automatically to rotate the state to the X basis
> If basis = 'Y', add a Haddamard + S gate automatically to rotate the state to the Y basis

That would be very useful for several applications: computing expectation values "manually" (useful for educational purposes, closer to the experimental treatment of the quantum computing output) or quantum circuit cutting (which requires measuring and preparing the states on a particular basis)

Following this idea, I would also suggest automatically initializing all qubits in the |+>, |->, |1>, |+i>, |-i>, etc states. e.g. including some function `initialize(basis=X, Y, Z, eigenstate = plus, minus)` or similar.  


Discussion:

Patch:
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index b08361993..646ef5e02 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -331,6 +331,24 @@ class NumpyBackend(Backend):
     def execute_circuit(
         self, circuit, initial_state=None, nshots=None, return_array=False
     ):
+        if isinstance(initial_state, type(circuit)):
+            if not initial_state.density_matrix == circuit.density_matrix:
+                raise_error(
+                    ValueError,
+                    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
+            ):  # pragma: no cover
+                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)
+
         if circuit.repeated_execution:
             return self.execute_circuit_repeated(circuit, initial_state, nshots)
 
diff --git a/src/qibo/gates/abstract.py b/src/qibo/gates/abstract.py
index 42fb80445..135605e16 100644
--- a/src/qibo/gates/abstract.py
+++ b/src/qibo/gates/abstract.py
@@ -270,6 +270,12 @@ class Gate:
             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 a455988cf..548d56027 100644
--- a/src/qibo/gates/gates.py
+++ b/src/qibo/gates/gates.py
@@ -215,6 +215,9 @@ class X(Gate):
         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 @@ class Y(Gate):
         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.
@@ -252,6 +261,9 @@ class Z(Gate):
             gate = super().controlled_by(*q)
         return gate
 
+    def basis_rotation(self):
+        return None
+
 
 class S(Gate):
     """The S gate.
diff --git a/src/qibo/gates/measurements.py b/src/qibo/gates/measurements.py
index 834fde706..a567a66ee 100644
--- a/src/qibo/gates/measurements.py
+++ b/src/qibo/gates/measurements.py
@@ -2,11 +2,12 @@ from typing import Dict, Optional, Tuple
 
 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,12 @@ 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`, 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
             that it is flipped, a list or tuple that has the same length
@@ -37,6 +44,7 @@ class M(Gate):
         *q,
         register_name: Optional[str] = None,
         collapse: bool = False,
+        basis: Gate = Z,
         p0: Optional["ProbsType"] = None,
         p1: Optional["ProbsType"] = None,
     ):
@@ -70,6 +78,22 @@ class M(Gate):
             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
+        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()
+            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/models/circuit.py b/src/qibo/models/circuit.py
index df3aad4e0..0e6f98332 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -575,8 +575,9 @@ class Circuit:
                         "".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
@@ -597,6 +598,7 @@ class Circuit:
                 return gate.result
 
             else:
+                self.queue.append(gate)
                 for measurement in list(self.measurements):
                     if set(measurement.qubits) & set(gate.qubits):
                         measurement.collapse = True
@@ -982,8 +984,11 @@ class Circuit:
     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


Test patch:
diff --git a/tests/test_gates_gates.py b/tests/test_gates_gates.py
index 3f9ec3476..09d752e23 100644
--- a/tests/test_gates_gates.py
+++ b/tests/test_gates_gates.py
@@ -785,6 +785,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 ###########################
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index a21a2bbf7..276b6e67b 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -435,3 +435,40 @@ 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}
+
+
+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}
+    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):
+    c = models.Circuit(4)
+    with pytest.raises(ValueError):
+        c.add(gates.M(0, 1, 2, 3, basis=[gates.X, gates.Z, gates.X]))
diff --git a/tests/test_models_circuit_execution.py b/tests/test_models_circuit_execution.py
index 5056011f7..0e426ecbb 100644
--- a/tests/test_models_circuit_execution.py
+++ b/tests/test_models_circuit_execution.py
@@ -105,3 +105,32 @@ 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)
+
+
+@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)
+
+
+def test_initial_state_error(backend):
+    nqubits = 10
+    c = Circuit(nqubits)
+    c.add(gates.X(i) for i in range(nqubits))
+
+    c1 = Circuit(nqubits, density_matrix=True)
+    c1.add(gates.H(i) for i in range(nqubits))
+
+    with pytest.raises(ValueError):
+        backend.execute_circuit(c, c1)