Merge pull request #1045 from qiboteam/readout_mitigation

Refactor readout mitigation matrix routine

Author: Edoardo-Pedicillo

src/qibocal/protocols/readout_mitigation_matrix.py

@@ -7,24 +7,18 @@
 from qibo import gates
 from qibo.backends import GlobalBackend
 from qibo.models import Circuit
-from qibolab import ExecutionParameters
 from qibolab.platform import Platform
-from qibolab.pulses import PulseSequence
 from qibolab.qubits import QubitId
 
 from qibocal.auto.operation import Data, Parameters, Results, Routine
 from qibocal.auto.transpile import dummy_transpiler, execute_transpiled_circuit
 from qibocal.config import log
 
-from .utils import calculate_frequencies
-
 
 @dataclass
 class ReadoutMitigationMatrixParameters(Parameters):
     """ReadoutMitigationMatrix matrix inputs."""
 
-    pulses: Optional[bool] = True
-    """Get readout mitigation matrix using pulses. If False gates will be used."""
     nshots: Optional[int] = None
     """Number of shots."""
     relaxation_time: Optional[int] = None
@@ -37,10 +31,14 @@ class ReadoutMitigationMatrixResults(Results):
         field(default_factory=dict)
     )
     """Readout mitigation matrices (inverse of measurement matrix)."""
-    measurement_matrix: dict[tuple[QubitId, ...], npt.NDArray[np.float64]] = field(
-        default_factory=dict
-    )
-    """Matrix containing measurement matrices for each state."""
+
+
+ReadoutMitigationMatrixType = np.dtype(
+    [
+        ("state", int),
+        ("frequency", np.float64),
+    ]
+)
 
 
 @dataclass
@@ -54,40 +52,6 @@ class ReadoutMitigationMatrixData(Data):
     data: dict = field(default_factory=dict)
     """Raw data acquited."""
 
-    def add(self, qubits, state, freqs):
-        for result_state, freq in freqs.items():
-            self.data[
-                qubits
-                + (
-                    state,
-                    result_state,
-                )
-            ] = freq
-
-        for basis in [format(i, f"0{len(qubits)}b") for i in range(2 ** len(qubits))]:
-            if (
-                qubits
-                + (
-                    state,
-                    basis,
-                )
-                not in self.data
-            ):
-                self.data[
-                    qubits
-                    + (
-                        state,
-                        basis,
-                    )
-                ] = 0
-
-    def __getitem__(self, qubits):
-        return {
-            index: value
-            for index, value in self.data.items()
-            if qubits == list(index[: len(index) - 2])
-        }
-
 
 def _acquisition(
     params: ReadoutMitigationMatrixParameters,
@@ -105,75 +69,52 @@ def _acquisition(
         nqubits = len(qubits)
         for i in range(2**nqubits):
             state = format(i, f"0{nqubits}b")
-            if params.pulses:
-                sequence = PulseSequence()
-                for q, bit in enumerate(state):
-                    if bit == "1":
-                        sequence.add(
-                            platform.create_RX_pulse(
-                                qubits[q], start=0, relative_phase=0
-                            )
-                        )
-                measurement_start = sequence.finish
-                for q in range(len(state)):
-                    MZ_pulse = platform.create_MZ_pulse(
-                        qubits[q], start=measurement_start
-                    )
-                    sequence.add(MZ_pulse)
-                results = platform.execute_pulse_sequence(
-                    sequence, ExecutionParameters(nshots=params.nshots)
-                )
-                data.add(
-                    tuple(qubits), state, calculate_frequencies(results, tuple(qubits))
+            c = Circuit(
+                nqubits,
+            )
+            for q, bit in enumerate(state):
+                if bit == "1":
+                    c.add(gates.X(q))
+            c.add(gates.M(*range(nqubits)))
+            _, results = execute_transpiled_circuit(
+                c, qubits, backend, nshots=params.nshots, transpiler=transpiler
+            )
+            frequencies = np.zeros(2 ** len(qubits))
+            for i, freq in results.frequencies().items():
+                frequencies[int(i, 2)] = freq
+            for freq in frequencies:
+                data.register_qubit(
+                    ReadoutMitigationMatrixType,
+                    (qubits),
+                    dict(
+                        state=np.array([int(state, 2)]),
+                        frequency=freq,
+                    ),
                 )
-            else:
-                c = Circuit(
-                    platform.nqubits,
-                    wire_names=[str(i) for i in range(platform.nqubits)],
-                )
-                for q, bit in enumerate(state):
-                    if bit == "1":
-                        c.add(gates.X(qubits[q]))
-                c.add(gates.M(*[qubits[i] for i in range(len(state))]))
-                _, results = execute_transpiled_circuit(
-                    c, qubit_map, backend, nshots=params.nshots, transpiler=transpiler
-                )
-                data.add(tuple(qubits), state, dict(results.frequencies()))
     return data
 
 
 def _fit(data: ReadoutMitigationMatrixData) -> ReadoutMitigationMatrixResults:
     """Post processing for readout mitigation matrix protocol."""
     readout_mitigation_matrix = {}
-    measurement_matrix = {}
-    for qubit in data.qubit_list:
-        qubit_data = data[qubit]
-        matrix = np.zeros((2 ** len(qubit), 2 ** len(qubit)))
-        computational_basis = [
-            format(i, f"0{len(qubit)}b") for i in range(2 ** len(qubit))
-        ]
-        for state in computational_basis:
-            column = np.zeros(2 ** len(qubit))
-            qubit_state_data = {
-                index: value
-                for index, value in qubit_data.items()
-                if index[-2] == state
-            }
-            for index, value in qubit_state_data.items():
-                column[(int(index[-1], 2))] = value / data.nshots
-            matrix[:, int(state, 2)] = np.flip(column)
-
-        measurement_matrix[tuple(qubit)] = matrix.tolist()
+    for qubits in data.qubit_list:
+        qubit_data = data.data[tuple(qubits)]
+        mitigation_matrix = []
+        for state in range(2 ** len(qubits)):
+            mitigation_matrix.append(qubit_data[qubit_data.state == state].frequency)
+        mitigation_matrix = np.vstack(mitigation_matrix) / data.nshots
         try:
-            readout_mitigation_matrix[tuple(qubit)] = np.linalg.inv(matrix).tolist()
+            readout_mitigation_matrix[tuple(qubits)] = np.linalg.inv(
+                mitigation_matrix
+            ).tolist()
         except np.linalg.LinAlgError as e:
             log.warning(f"ReadoutMitigationMatrix: the fitting was not succesful. {e}")
-
-    return ReadoutMitigationMatrixResults(
+    res = ReadoutMitigationMatrixResults(
         readout_mitigation_matrix=readout_mitigation_matrix,
-        measurement_matrix=measurement_matrix,
     )
 
+    return res
+
 
 def _plot(
     data: ReadoutMitigationMatrixData,
@@ -187,12 +128,12 @@ def _plot(
         computational_basis = [
             format(i, f"0{len(target)}b") for i in range(2 ** len(target))
         ]
-        z = fit.measurement_matrix[tuple(target)]
-
+        measurement_matrix = np.linalg.inv(fit.readout_mitigation_matrix[tuple(target)])
+        z = measurement_matrix
         fig = px.imshow(
             z,
             x=computational_basis,
-            y=computational_basis[::-1],
+            y=computational_basis,
             text_auto=True,
             labels={
                 "x": "Prepeared States",

src/qibocal/protocols/two_qubit_interaction/chsh/circuits.py

@@ -7,7 +7,7 @@
 from .utils import READOUT_BASIS
 
 
-def create_bell_circuit(nqubits, qubits, theta=np.pi / 4, bell_state=0):
+def create_bell_circuit(theta=np.pi / 4, bell_state=0):
     """Creates the circuit to generate the bell states and with a theta-measurement
     bell_state chooses the initial bell state for the test:
     0 -> |00>+|11>
@@ -17,24 +17,24 @@ def create_bell_circuit(nqubits, qubits, theta=np.pi / 4, bell_state=0):
     Native defaults to only using GPI2 and GPI gates.
     """
     p = [0, 0]
-    c = Circuit(nqubits)
-    c.add(gates.H(qubits[0]))
-    c.add(gates.H(qubits[1]))
-    c.add(gates.CZ(qubits[0], qubits[1]))
-    c.add(gates.H(qubits[1]))
+    c = Circuit(2)
+    c.add(gates.H(0))
+    c.add(gates.H(1))
+    c.add(gates.CZ(0, 1))
+    c.add(gates.H(1))
     if bell_state == 1:
-        c.add(gates.Z(qubits[0]))
+        c.add(gates.Z(0))
     elif bell_state == 2:
-        c.add(gates.Z(qubits[0]))
-        c.add(gates.X(qubits[0]))
+        c.add(gates.Z(0))
+        c.add(gates.X(0))
     elif bell_state == 3:
-        c.add(gates.X(qubits[0]))
+        c.add(gates.X(0))
 
-    c.add(gates.RY(qubits[0], theta))
+    c.add(gates.RY(0, theta))
     return c, p
 
 
-def create_bell_circuit_native(nqubits, qubits, theta=np.pi / 4, bell_state=0):
+def create_bell_circuit_native(theta=np.pi / 4, bell_state=0):
     """Creates the circuit to generate the bell states and with a theta-measurement
     bell_state chooses the initial bell state for the test:
     0 -> |00>+|11>
@@ -44,35 +44,33 @@ def create_bell_circuit_native(nqubits, qubits, theta=np.pi / 4, bell_state=0):
     Native defaults to only using GPI2 and GPI gates.
     """
 
-    c = Circuit(nqubits)
+    c = Circuit(2)
     p = [0, 0]
-    c.add(gates.GPI2(qubits[0], np.pi / 2))
-    c.add(gates.GPI2(qubits[1], np.pi / 2))
-    c.add(gates.CZ(qubits[0], qubits[1]))
-    c.add(gates.GPI2(qubits[1], -np.pi / 2))
+    c.add(gates.GPI2(0, np.pi / 2))
+    c.add(gates.GPI2(1, np.pi / 2))
+    c.add(gates.CZ(0, 1))
+    c.add(gates.GPI2(1, -np.pi / 2))
     if bell_state == 0:
         p[0] += np.pi
     elif bell_state == 1:
         p[0] += 0
     elif bell_state == 2:
         p[0] += 0
-        c.add(gates.GPI2(qubits[0], p[0]))
-        c.add(gates.GPI2(qubits[0], p[0]))
+        c.add(gates.GPI2(0, p[0]))
+        c.add(gates.GPI2(0, p[0]))
     elif bell_state == 3:
         p[0] += np.pi
-        c.add(gates.GPI2(qubits[0], p[0]))
-        c.add(gates.GPI2(qubits[0], p[0]))
+        c.add(gates.GPI2(0, p[0]))
+        c.add(gates.GPI2(0, p[0]))
 
-    c.add(gates.GPI2(qubits[0], p[0]))
+    c.add(gates.GPI2(0, p[0]))
     p[0] += theta
-    c.add(gates.GPI2(qubits[0], p[0] + np.pi))
+    c.add(gates.GPI2(0, p[0] + np.pi))
 
     return c, p
 
 
 def create_chsh_circuits(
-    platform,
-    qubits,
     theta=np.pi / 4,
     bell_state=0,
     native=True,
@@ -84,15 +82,14 @@ def create_chsh_circuits(
     """
     create_bell = create_bell_circuit_native if native else create_bell_circuit
     chsh_circuits = {}
-    nqubits = platform.nqubits if platform else max(qubits) + 1
     for basis in readout_basis:
-        c, p = create_bell(nqubits, qubits, theta, bell_state)
+        c, p = create_bell(theta, bell_state)
         for i, base in enumerate(basis):
             if base == "X":
                 if native:
-                    c.add(gates.GPI2(qubits[i], p[i] + np.pi / 2))
+                    c.add(gates.GPI2(i, p[i] + np.pi / 2))
                 else:
-                    c.add(gates.H(qubits[i]))
-        c.add(gates.M(*qubits))
+                    c.add(gates.H(i))
+        c.add(gates.M(0, 1))
         chsh_circuits[basis] = c
     return chsh_circuits

src/qibocal/protocols/two_qubit_interaction/chsh/protocol.py

@@ -76,14 +76,16 @@ class CHSHData(Data):
 
     def save(self, path: Path):
         """Saving data including mitigation matrix."""
-
-        np.savez(
-            path / f"{MITIGATION_MATRIX_FILE}.npz",
-            **{
-                json.dumps((control, target)): self.mitigation_matrix[control, target]
-                for control, target, _, _, _ in self.data
-            },
-        )
+        if self.mitigation_matrix:
+            np.savez(
+                path / f"{MITIGATION_MATRIX_FILE}.npz",
+                **{
+                    json.dumps((control, target)): self.mitigation_matrix[
+                        control, target
+                    ]
+                    for control, target, _, _, _ in self.data
+                },
+            )
         super().save(path=path)
 
     @classmethod
@@ -174,7 +176,7 @@ def _acquisition_pulses(
 
     if params.apply_error_mitigation:
         mitigation_data = mitigation_acquisition(
-            mitigation_params(pulses=True, nshots=params.nshots), platform, targets
+            mitigation_params(nshots=params.nshots), platform, targets
         )
         mitigation_results = mitigation_fit(mitigation_data)
 
@@ -223,10 +225,9 @@ def _acquisition_circuits(
     backend = GlobalBackend()
     backend.platform = platform
     transpiler = dummy_transpiler(backend)
-    qubit_map = [i for i in range(platform.nqubits)]
     if params.apply_error_mitigation:
         mitigation_data = mitigation_acquisition(
-            mitigation_params(pulses=False, nshots=params.nshots), platform, targets
+            mitigation_params(nshots=params.nshots), platform, targets
         )
         mitigation_results = mitigation_fit(mitigation_data)
     for pair in targets:
@@ -242,8 +243,6 @@ def _acquisition_circuits(
         for bell_state in params.bell_states:
             for theta in thetas:
                 chsh_circuits = create_chsh_circuits(
-                    platform,
-                    qubits=pair,
                     bell_state=bell_state,
                     theta=theta,
                     native=params.native,
@@ -254,7 +253,7 @@ def _acquisition_circuits(
                         nshots=params.nshots,
                         transpiler=transpiler,
                         backend=backend,
-                        qubit_map=qubit_map,
+                        qubit_map=pair,
                     )
                     frequencies = result.frequencies()
                     data.register_basis(pair, bell_state, basis, frequencies)