From 7e40aafb5816ca9f79cf68adf90f710d21a48e44 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 20 May 2024 19:17:51 +0400
Subject: [PATCH] refactor: reduce code repetition further

---
 src/qibocal/protocols/state_tomography.py     | 149 ++++++++++--------
 .../protocols/two_qubit_state_tomography.py   |  71 +--------
 2 files changed, 83 insertions(+), 137 deletions(-)

diff --git a/src/qibocal/protocols/state_tomography.py b/src/qibocal/protocols/state_tomography.py
index 17608994e..956b7baa0 100644
--- a/src/qibocal/protocols/state_tomography.py
+++ b/src/qibocal/protocols/state_tomography.py
@@ -213,8 +213,8 @@ def plot_rho(fig, zz, trace_options, figure_options, showlegend=None):
         )
 
 
-def _plot(data: StateTomographyData, fit: StateTomographyResults, target: QubitId):
-    """Plotting for state tomography"""
+def plot_reconstruction(ideal, measured):
+    """Plot 3D plot with reconstruction of ideal and measured density matrix."""
     fig = make_subplots(
         rows=1,
         cols=2,
@@ -226,82 +226,91 @@ def _plot(data: StateTomographyData, fit: StateTomographyResults, target: QubitI
         ),
     )
 
-    if fit is not None:
-        # computing limits for colorscale
-        min_re, max_re = np.min(fit.target_density_matrix_real[target]), np.max(
-            fit.target_density_matrix_real[target]
-        )
-        min_im, max_im = np.min(fit.target_density_matrix_imag[target]), np.max(
-            fit.target_density_matrix_imag[target]
-        )
+    # computing limits for colorscale
+    min_re, max_re = np.min(ideal.real), np.max(ideal.real)
+    min_im, max_im = np.min(ideal.imag), np.max(ideal.imag)
+
+    # add offset
+    if np.abs(min_re - max_re) < 1e-5:
+        min_re = min_re - 0.1
+        max_re = max_re + 0.1
+    if np.abs(min_im - max_im) < 1e-5:
+        min_im = min_im - 0.1
+        max_im = max_im + 0.1
+
+    plot_rho(
+        fig,
+        measured.real,
+        trace_options=dict(
+            color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
+        ),
+        figure_options=dict(row=1, col=1),
+    )
 
-        # add offset
-        if np.abs(min_re - max_re) < 1e-5:
-            min_re = min_re - 0.1
-            max_re = max_re + 0.1
-
-        if np.abs(min_im - max_im) < 1e-5:
-            min_im = min_im - 0.1
-            max_im = max_im + 0.1
-
-        plot_rho(
-            fig,
-            fit.measured_density_matrix_real[target],
-            trace_options=dict(
-                color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
-            ),
-            figure_options=dict(row=1, col=1),
-        )
+    plot_rho(
+        fig,
+        ideal.real,
+        trace_options=dict(
+            color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
+        ),
+        figure_options=dict(row=1, col=1),
+    )
 
-        plot_rho(
-            fig,
-            fit.target_density_matrix_real[target],
-            trace_options=dict(
-                color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
-            ),
-            figure_options=dict(row=1, col=1),
-        )
+    plot_rho(
+        fig,
+        measured.imag,
+        trace_options=dict(
+            color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
+        ),
+        figure_options=dict(row=1, col=2),
+        showlegend=False,
+    )
 
-        plot_rho(
-            fig,
-            fit.measured_density_matrix_imag[target],
-            trace_options=dict(
-                color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
-            ),
-            figure_options=dict(row=1, col=2),
-            showlegend=False,
-        )
+    plot_rho(
+        fig,
+        ideal.imag,
+        trace_options=dict(
+            color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
+        ),
+        figure_options=dict(row=1, col=2),
+        showlegend=False,
+    )
 
-        plot_rho(
-            fig,
-            fit.target_density_matrix_imag[target],
-            trace_options=dict(
-                color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
-            ),
-            figure_options=dict(row=1, col=2),
-            showlegend=False,
-        )
+    fig.update_scenes(
+        xaxis=dict(tickvals=list(range(len(ideal)))),
+        yaxis=dict(tickvals=list(range(len(ideal)))),
+        zaxis=dict(range=[-1, 1]),
+    )
 
-        fig.update_scenes(
-            xaxis=dict(tickvals=[0, 1]),
-            yaxis=dict(tickvals=[0, 1]),
-            zaxis=dict(range=[-1, 1]),
-        )
+    return fig
 
-        fitting_report = table_html(
-            table_dict(
-                target,
-                [
-                    "Fidelity",
-                ],
-                [
-                    np.round(fit.fidelity[target], 4),
-                ],
-            )
+
+def _plot(data: StateTomographyData, fit: StateTomographyResults, target: QubitId):
+    """Plotting for state tomography"""
+    if fit is None:
+        return [], ""
+
+    ideal = np.array(fit.target_density_matrix_real[target]) + 1j * np.array(
+        fit.target_density_matrix_imag[target]
+    )
+    measured = np.array(fit.measured_density_matrix_real[target]) + 1j * np.array(
+        fit.measured_density_matrix_imag[target]
+    )
+    fig = plot_reconstruction(ideal, measured)
+
+    fitting_report = table_html(
+        table_dict(
+            target,
+            [
+                "Fidelity",
+            ],
+            [
+                np.round(fit.fidelity[target], 4),
+            ],
         )
+    )
 
-        return [fig], fitting_report
-    return [], ""
+    return [fig], fitting_report
 
 
 state_tomography = Routine(_acquisition, _fit, _plot)
diff --git a/src/qibocal/protocols/two_qubit_state_tomography.py b/src/qibocal/protocols/two_qubit_state_tomography.py
index 83a4945d4..361e22ed3 100644
--- a/src/qibocal/protocols/two_qubit_state_tomography.py
+++ b/src/qibocal/protocols/two_qubit_state_tomography.py
@@ -18,7 +18,7 @@
 from qibocal.auto.operation import DATAFILE, Data, Results, Routine
 from qibocal.auto.transpile import dummy_transpiler, execute_transpiled_circuit
 
-from .state_tomography import StateTomographyParameters, plot_rho
+from .state_tomography import StateTomographyParameters, plot_reconstruction
 from .utils import table_dict, table_html
 
 SINGLE_QUBIT_BASIS = ["X", "Y", "Z"]
@@ -222,7 +222,6 @@ def plot_measurements(data: StateTomographyData, target: QubitPairId):
     fig = make_subplots(
         rows=3,
         cols=3,
-        # "$\text{Plot 1}$"
         subplot_titles=tuple(
             f"{basis1}<sub>1</sub>{basis2}<sub>2</sub>"
             for basis1, basis2 in TWO_QUBIT_BASIS
@@ -282,72 +281,10 @@ def _plot(data: StateTomographyData, fit: StateTomographyResults, target: QubitP
         )
         return [fig_measurements], fitting_report
 
-    fig = make_subplots(
-        rows=1,
-        cols=2,
-        start_cell="top-left",
-        specs=[[{"type": "scatter3d"}, {"type": "scatter3d"}]],
-        subplot_titles=(
-            "Re(ρ)",
-            "Im(ρ)",
-        ),
-    )
-    # computing limits for colorscale
-    min_re, max_re = np.min(data.ideal[target].real), np.max(data.ideal[target].real)
-    min_im, max_im = np.min(data.ideal[target].imag), np.max(data.ideal[target].imag)
-
-    # add offset
-    if np.abs(min_re - max_re) < 1e-5:
-        min_re = min_re - 0.1
-        max_re = max_re + 0.1
-
-    if np.abs(min_im - max_im) < 1e-5:
-        min_im = min_im - 0.1
-        max_im = max_im + 0.1
-
-    plot_rho(
-        fig,
-        fit.measured_density_matrix_real[target],
-        trace_options=dict(
-            color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
-        ),
-        figure_options=dict(row=1, col=1),
-    )
-
-    plot_rho(
-        fig,
-        data.ideal[target].real,
-        trace_options=dict(
-            color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
-        ),
-        figure_options=dict(row=1, col=1),
-    )
-
-    plot_rho(
-        fig,
-        fit.measured_density_matrix_imag[target],
-        trace_options=dict(
-            color="rgba(255,100,0,0.1)", name="experiment", legendgroup="experiment"
-        ),
-        figure_options=dict(row=1, col=2),
-        showlegend=False,
-    )
-
-    plot_rho(
-        fig,
-        data.ideal[target].imag,
-        trace_options=dict(
-            color="rgba(100,0,100,0.1)", name="simulation", legendgroup="simulation"
-        ),
-        figure_options=dict(row=1, col=2),
-        showlegend=False,
-    )
-
-    fig.update_scenes(
-        xaxis=dict(tickvals=list(range(4))),
-        yaxis=dict(tickvals=list(range(4))),
-        zaxis=dict(range=[-1, 1]),
+    measured = np.array(fit.measured_density_matrix_real[target]) + 1j * np.array(
+        fit.measured_raw_density_matrix_imag[target]
     )
+    fig = plot_reconstruction(data.ideal[target], measured)
 
     fitting_report = table_html(
         table_dict(