Measurement in Pauli basis

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

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.
@@ -252,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.

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):

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

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 ###########################

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}