refactor: Improve plot for flux spectroscopies

Author: andrea-pasquale

src/qibocal/protocols/flux_dependence/resonator_flux_dependence.py

@@ -79,7 +79,6 @@ class ResonatorFluxData(Data):
     """Qubit bare resonator frequency power provided by the user."""
     charging_energy: dict[QubitId, float] = field(default_factory=dict)
     """Qubit charging energy in Hz."""
-
     data: dict[QubitId, npt.NDArray[ResFluxType]] = field(default_factory=dict)
     """Raw data acquired."""
 
@@ -89,6 +88,20 @@ def register_qubit(self, qubit, freq, bias, signal, phase):
             freq, bias, signal, phase, dtype=ResFluxType
         )
 
+    @property
+    def feature(self) -> str:
+        """Get feature (either min or max) depending on resonator type."""
+        return "min" if self.resonator_type == "2D" else "max"
+
+    def filtered_data(self, qubit: QubitId) -> np.ndarray:
+        """Apply mask to specific qubit data."""
+        return extract_feature(
+            self.data[qubit].freq,
+            self.data[qubit].bias,
+            self.data[qubit].signal,
+            self.feature,
+        )
+
 
 def _acquisition(
     params: ResonatorFluxParameters,
@@ -191,15 +204,8 @@ def _fit(data: ResonatorFluxData) -> ResonatorFluxResults:
     matrix_element = {}
 
     for qubit in data.qubits:
-        qubit_data = data[qubit]
-        biases = qubit_data.bias
-        frequencies = qubit_data.freq
-        signal = qubit_data.signal
-
         # extract signal from 2D plot based on SNR mask
-        frequencies, biases = extract_feature(
-            frequencies, biases, signal, "min" if data.resonator_type == "2D" else "max"
-        )
+        frequencies, biases = data.filtered_data(qubit)
 
         # define fit function
         def fit_function(

src/qibocal/protocols/flux_dependence/utils.py

@@ -27,16 +27,29 @@ def flux_dependence_plot(data, fit, qubit, fit_function=None):
     qubit_data = data[qubit]
     frequencies = qubit_data.freq * HZ_TO_GHZ
 
-    fig = go.Figure()
+    filtered_freq, filtered_bias = data.filtered_data(qubit)
 
+    fig = go.Figure()
     fig.add_trace(
         go.Heatmap(
             x=qubit_data.freq * HZ_TO_GHZ,
             y=qubit_data.bias,
             z=qubit_data.signal,
+            colorbar=dict(title="Signal [a.u.]"),
+            colorscale="Viridis",
         ),
     )
 
+    fig.add_trace(
+        go.Scatter(
+            x=filtered_freq * HZ_TO_GHZ,
+            y=filtered_bias,
+            name="Estimated points",
+            mode="markers",
+            marker=dict(color="orange"),
+        )
+    )
+
     # TODO: This fit is for frequency, can it be reused here, do we even want the fit ?
     if (
         fit is not None
@@ -52,7 +65,7 @@ def flux_dependence_plot(data, fit, qubit, fit_function=None):
                 y=bias,
                 showlegend=True,
                 name="Fit",
-                marker=dict(color="green"),
+                marker=dict(color="orange"),
             ),
         )
 
@@ -67,8 +80,7 @@ def flux_dependence_plot(data, fit, qubit, fit_function=None):
                 mode="markers",
                 marker=dict(
                     size=8,
-                    color="black",
-                    symbol="cross",
+                    color="red",
                 ),
                 name="Sweetspot",
                 showlegend=True,