Adding PCA fit for rabi signal

src/qibocal/protocols/drag/drag.py

@@ -213,7 +213,9 @@ def _fit(data: DragTuningData) -> DragTuningResults:
         beta_params = qubit_data["beta"]
 
         # Guessing period using fourier transform
-        period_guess = fallback_period(guess_period(beta_params, qubit_data["prob"]))
+        period_guess = fallback_period(
+            guess_period(beta_params, qubit_data["prob"])
+        ).tolist()
         pguess = [0.5, 0.5, period_guess, 0]
         try:
             popt, _ = curve_fit(

src/qibocal/protocols/rabi/amplitude.py

@@ -115,7 +115,6 @@ def _fit(data: RabiAmplitudeData) -> RabiAmplitudeResults:
                 y,
                 pguess,
                 sigma=qubit_data.error,
-                signal=False,
             )
             pi_pulse_amplitudes[qubit] = [pi_pulse_parameter, perr[2] / 2]
             fitted_parameters[qubit] = popt

src/qibocal/protocols/rabi/amplitude_frequency.py

@@ -5,7 +5,6 @@
 import numpy as np
 import numpy.typing as npt
 import plotly.graph_objects as go
-from plotly.subplots import make_subplots
 from qibolab import AcquisitionType, AveragingMode, Parameter, Sweeper
 
 from qibocal.auto.operation import Protocol, QubitId
@@ -17,8 +16,8 @@
     table_dict,
     table_html,
 )
+from qibocal.result import probability
 
-from ...result import probability
 from .amplitude_frequency_signal import (
     RabiAmplitudeFreqSignalData,
     RabiAmplitudeFrequencySignalParameters,
@@ -27,6 +26,7 @@
 )
 from .utils import (
     fit_amplitude_function,
+    plot_probabilities,
     rabi_amplitude_function,
     rabi_initial_guess,
     sequence_amplitude,
@@ -144,39 +144,41 @@ def _fit(data: RabiAmplitudeFreqData) -> RabiAmplitudeFrequencyResults:
         probability_matrix = probability.reshape(len(amps), len(freqs)).T
 
         # guess optimal frequency maximizing oscillation amplitude
-        index = np.argmax([max(x) - min(x) for x in probability_matrix])
-        frequency = freqs[index]
-
-        y = probability_matrix[index]
-        error = data[qubit].error[data[qubit].freq == frequency]
+        # here prob_matrix has dimensions (n_freqs, n_amps), so we
+        # need to compute initial guesses over axis==1
+        full_pguesses = rabi_initial_guess(
+            amps, probability_matrix, "amp", signal=False, axis=1
+        )
 
-        y_min = np.min(y)
-        y_max = np.max(y)
-        x_min = np.min(amps)
-        x_max = np.max(amps)
-        x = (amps - x_min) / (x_max - x_min)
-        y = (y - y_min) / (y_max - y_min)
+        # guess has the following elements:
+        # 0. median guess
+        # 1. amplitude guess
+        # 2. period guess
+        # 3. phase guess
+        # 4. decaying constant guess
+        # we estimate the best frequency by maximizing the amplitude estimation
+        index = np.argmax(full_pguesses[1])
 
-        pguess = rabi_initial_guess(x, y, "amp", signal=False)
+        frequency = freqs[index]
+        y = probability_matrix[index, :].ravel()
+        error = data[qubit].error[data[qubit].freq == frequency]
 
+        # initial guesses for the best frequency row
+        pguess = [p[index] for p in full_pguesses]
         try:
             popt, perr, pi_pulse_parameter = fit_amplitude_function(
-                x,
+                amps,
                 y,
                 pguess,
                 sigma=error,
-                signal=False,
-                x_limits=(x_min, x_max),
-                y_limits=(y_min, y_max),
             )
             fitted_frequencies[qubit] = frequency
             fitted_amplitudes[qubit] = [pi_pulse_parameter, perr[2] / 2]
             fitted_parameters[qubit] = popt if isinstance(popt, list) else popt.tolist()
-
             chi2[qubit] = (
                 chi2_reduced(
                     y,
-                    rabi_amplitude_function(x, *popt),
+                    rabi_amplitude_function(amps, *popt),
                     error,
                 ),
                 np.sqrt(2 / len(y)),
@@ -186,7 +188,7 @@ def _fit(data: RabiAmplitudeFreqData) -> RabiAmplitudeFrequencyResults:
 
     return RabiAmplitudeFrequencyResults(
         amplitude=fitted_amplitudes,
-        length=data.durations,
+        length={key: [value, 0] for key, value in data.durations.items()},
         fitted_parameters=fitted_parameters,
         frequency=fitted_frequencies,
         chi2=chi2,
@@ -197,47 +199,39 @@ def _fit(data: RabiAmplitudeFreqData) -> RabiAmplitudeFrequencyResults:
 def _plot(
     data: RabiAmplitudeFreqData,
     target: QubitId,
-    fit: RabiAmplitudeFrequencyResults = None,
+    fit: RabiAmplitudeFrequencyResults | None = None,
 ):
     """Plotting function for RabiAmplitudeFrequency."""
     figures = []
     fitting_report = ""
-    fig = make_subplots(
-        rows=1,
-        cols=1,
-        horizontal_spacing=0.1,
-        vertical_spacing=0.2,
-        subplot_titles=("Probability",),
-    )
+    fig = go.Figure()
+    frequencies = data.frequencies(target) * HZ_TO_GHZ
+    amplitudes = data.amplitudes(target)
     qubit_data = data[target]
-    frequencies = qubit_data.freq * HZ_TO_GHZ
-    amplitudes = qubit_data.amp
-
-    fig.update_xaxes(title_text="Amplitude [a.u.]", row=1, col=1)
-    fig.update_yaxes(title_text="Frequency [GHz]", row=1, col=1)
-
-    figures.append(fig)
 
     fig.add_trace(
         go.Heatmap(
             x=amplitudes,
             y=frequencies,
-            z=qubit_data.prob,
+            z=qubit_data.prob.reshape(len(amplitudes), len(frequencies)).T,
         ),
-        row=1,
-        col=1,
+    )
+    fig.update_layout(
+        title="Probability",
+        xaxis_title="Amplitude [a.u.]",
+        yaxis_title="Frequency [GHz]",
     )
 
     if fit is not None:
+        selected_frequency = fit.frequency[target]
+
         fig.add_trace(
             go.Scatter(
                 x=[min(amplitudes), max(amplitudes)],
-                y=[fit.frequency[target] * HZ_TO_GHZ] * 2,
+                y=[selected_frequency * HZ_TO_GHZ] * 2,
                 mode="lines",
                 line={"color": "white", "width": 4, "dash": "dash"},
             ),
-            row=1,
-            col=1,
         )
         pulse_name = "Pi-half pulse" if data.rx90 else "Pi pulse"
 
@@ -252,10 +246,15 @@ def _plot(
             )
         )
 
-    fig.update_layout(
-        showlegend=False,
-        legend={"orientation": "h"},
-    )
+        fitted_data = data.return_row_data(selected_frequency, target)
+        rabi1d_figure, rabi1d_report = plot_probabilities(
+            fitted_data, target, fit, data.rx90
+        )
+        fitting_report += rabi1d_report
+        figures.extend(rabi1d_figure)
+
+    figures.insert(0, fig)
+
     return figures, fitting_report