Merge pull request #707 from qiboteam/refactor_results

Partial refactor in `fit` functions

doc/source/tutorials/advanced.rst

@@ -275,8 +275,8 @@ In the acquisition function we are going to perform the experiment.
 Result class
 ^^^^^^^^^^^^
 
-Here we decided to code a generic `Result` that contains the fitted
-parameters for each quibt.
+Here we decided to code a generic `Results` that contains the fitted
+parameters for each qubit.
 
 .. code-block:: python
 
@@ -287,6 +287,13 @@ parameters for each quibt.
         """Results object for data"""
         fitted_parameters: dict[QubitId, list] = field(default_factory=dict)
 
+.. note::
+
+    To check whether fitted parameters for a specific ``Qubit`` it might
+    be necessary to re-write the ``__contains__`` method if the ``Results``
+    inheritance include non-dictionary attributes.
+
+
 Fit function
 ^^^^^^^^^^^^
 

src/qibocal/auto/operation.py

@@ -2,7 +2,7 @@
 import json
 import time
 from copy import deepcopy
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, fields
 from functools import wraps
 from pathlib import Path
 from typing import Callable, Generic, NewType, Optional, TypeVar, Union
@@ -222,6 +222,18 @@ def load(cls, path: Path):
 class Results(AbstractData):
     """Generic runcard update."""
 
+    def __contains__(self, key: Union[QubitId, QubitPairId, tuple[QubitId, ...]]):
+        """Checking if qubit is in Results.
+
+        If key is not present means that fitting failed or
+        was not performed.
+        """
+        return all(
+            key in getattr(self, field.name)
+            for field in fields(self)
+            if isinstance(getattr(self, field.name), dict)
+        )
+
     @classmethod
     def load(cls, path: Path):
         """Load results."""
@@ -276,7 +288,7 @@ def data_type(self):
 
     @property
     def results_type(self):
-        """ "Results object type return by data acquisition."""
+        """Results object type returned by data acquisition."""
         return inspect.signature(self.fit).return_annotation
 
     # TODO: I don't like these properties but it seems to work

src/qibocal/auto/task.py

@@ -118,6 +118,7 @@ def run(
         except (RuntimeError, AttributeError):
             operation = dummy_operation
             parameters = DummyPars()
+
         if mode.name in ["autocalibration", "acquire"]:
             if operation.platform_dependent and operation.targets_dependent:
                 completed.data, completed.data_time = operation.acquisition(

src/qibocal/cli/report.py

@@ -74,8 +74,11 @@ def routine_targets(self, task_id: TaskId):
     def single_qubit_plot(self, task_id: TaskId, qubit: QubitId):
         """Generate single qubit plot."""
         node = self.history[task_id]
+        fit = node.results if node.results and qubit in node.results else None
+        # the fit is shown only if fitted parameters for the corresponding
+        # key are in Results.
         figures, fitting_report = node.task.operation.report(
-            data=node.data, fit=node.results, target=qubit
+            data=node.data, fit=fit, target=qubit
         )
         with tempfile.NamedTemporaryFile(delete=False) as temp:
             html_list = []

src/qibocal/protocols/characterization/allxy/drag_pulse_tuning.py

@@ -178,17 +178,15 @@ def _fit(data: DragPulseTuningData) -> DragPulseTuningResults:
         beta_params = qubit_data.beta
 
         try:
-            popt, pcov = curve_fit(drag_fit, beta_params, voltages)
+            popt, _ = curve_fit(drag_fit, beta_params, voltages)
             smooth_dataset = drag_fit(beta_params, popt[0], popt[1], popt[2], popt[3])
             min_abs_index = np.abs(smooth_dataset).argmin()
             beta_optimal = beta_params[min_abs_index]
-        except:
-            log.warning("drag_tuning_fit: the fitting was not succesful")
-            popt = np.array([0, 0, 1, 0])
-            beta_optimal = 0
+            fitted_parameters[qubit] = popt.tolist()
+            betas_optimal[qubit] = beta_optimal
+        except Exception as e:
+            log.warning(f"Drag fit failed on qubit {qubit} due to {e}.")
 
-        fitted_parameters[qubit] = popt.tolist()
-        betas_optimal[qubit] = beta_optimal
     return DragPulseTuningResults(betas_optimal, fitted_parameters)
 
 

src/qibocal/protocols/characterization/coherence/spin_echo.py

@@ -150,7 +150,7 @@ def _fit(data: SpinEchoData) -> SpinEchoResults:
             ),
             np.sqrt(2 / len(data[qubit].prob)),
         )
-        for qubit in data.qubits
+        for qubit in fitted_parameters
     }
 
     return SpinEchoResults(t2Echos, fitted_parameters, chi2)

src/qibocal/protocols/characterization/coherence/t2.py

@@ -125,7 +125,7 @@ def _fit(data: T2Data) -> T2Results:
             ),
             np.sqrt(2 / len(data[qubit].prob)),
         )
-        for qubit in data.qubits
+        for qubit in fitted_parameters
     }
     return T2Results(t2s, fitted_parameters, chi2)
 

src/qibocal/protocols/characterization/coherence/utils.py

@@ -51,14 +51,11 @@ def exponential_fit(data, zeno=None):
                 popt[2] * (x_max - x_min),
             ]
             t2 = popt[2]
+            fitted_parameters[qubit] = popt
+            decay[qubit] = t2
 
         except Exception as e:
-            log.warning(f"Exp decay fitting was not succesful. {e}")
-            popt = [0] * 3
-            t2 = 5.0
-
-        fitted_parameters[qubit] = popt
-        decay[qubit] = t2
+            log.warning(f"Exponential decay fit failed for qubit {qubit} due to {e}")
 
     return decay, fitted_parameters
 
@@ -100,14 +97,10 @@ def exponential_fit_probability(data):
                 popt[2] * (x_max - x_min),
             ]
             perr = np.sqrt(np.diag(perr))
+            fitted_parameters[qubit] = popt
+            dec = popt[2]
+            decay[qubit] = (dec, perr[2])
         except Exception as e:
-            log.warning(f"Exp decay fitting was not succesful. {e}")
-            popt = [0] * 3
-            dec = 5
-            perr = [1] * 3
-
-        fitted_parameters[qubit] = popt
-        dec = popt[2]
-        decay[qubit] = (dec, perr[2])
+            log.warning(f"Exponential decay fit failed for qubit {qubit} due to {e}")
 
     return decay, fitted_parameters

src/qibocal/protocols/characterization/coherence/zeno.py

@@ -122,7 +122,7 @@ def _fit(data: ZenoData) -> ZenoResults:
             ),
             np.sqrt(2 / len(data[qubit].prob)),
         )
-        for qubit in data.qubits
+        for qubit in fitted_parameters
     }
     return ZenoResults(t1s, fitted_parameters, chi2)
 

src/qibocal/protocols/characterization/dispersive_shift.py

@@ -173,35 +173,37 @@ def _fit(data: DispersiveShiftData) -> DispersiveShiftResults:
 
     frequency_0 = {}
     frequency_1 = {}
+    best_freqs = {}
     fitted_parameters_0 = {}
     fitted_parameters_1 = {}
 
     for i in range(2):
         for qubit in qubits:
             data_i = data[qubit, i]
-            freq, fitted_params, _ = lorentzian_fit(
+            fit_result = lorentzian_fit(
                 data_i, resonator_type=data.resonator_type, fit="resonator"
             )
-            if i == 0:
-                frequency_0[qubit] = freq
-                fitted_parameters_0[qubit] = fitted_params
-            else:
-                frequency_1[qubit] = freq
-                fitted_parameters_1[qubit] = fitted_params
+            if fit_result is not None:
+                if i == 0:
+                    frequency_0[qubit], fitted_parameters_0[qubit], _ = fit_result
+                else:
+                    frequency_1[qubit], fitted_parameters_1[qubit], _ = fit_result
+
             i_measures = data_i.i
             q_measures = data_i.q
 
             iq_couples[i].append(np.stack((i_measures, q_measures), axis=-1))
-    # for each qubit find the iq couple of 0-1 states that maximize the distance
+        # for each qubit find the iq couple of 0-1 states that maximize the distance
     iq_couples = np.array(iq_couples)
-    best_freqs = {}
+
     for idx, qubit in enumerate(qubits):
         frequencies = data[qubit, 0].freq
 
         max_index = np.argmax(
             np.linalg.norm(iq_couples[0][idx] - iq_couples[1][idx], axis=-1)
         )
         best_freqs[qubit] = frequencies[max_index]
+
     return DispersiveShiftResults(
         frequency_state_zero=frequency_0,
         frequency_state_one=frequency_1,

src/qibocal/protocols/characterization/dispersive_shift_qutrit.py

@@ -160,18 +160,16 @@ def _fit(data: DispersiveShiftQutritData) -> DispersiveShiftQutritResults:
     for i in range(3):
         for qubit in qubits:
             data_i = data[qubit, i]
-            freq, fitted_params, _ = lorentzian_fit(
+            fit_result = lorentzian_fit(
                 data_i, resonator_type=data.resonator_type, fit="resonator"
             )
-            if i == 0:
-                frequency_0[qubit] = freq
-                fitted_parameters_0[qubit] = fitted_params
-            elif i == 1:
-                frequency_1[qubit] = freq
-                fitted_parameters_1[qubit] = fitted_params
-            else:
-                frequency_2[qubit] = freq
-                fitted_parameters_2[qubit] = fitted_params
+            if fit_result is not None:
+                if i == 0:
+                    frequency_0[qubit], fitted_parameters_0[qubit], _ = fit_result
+                elif i == 1:
+                    frequency_1[qubit], fitted_parameters_1[qubit], _ = fit_result
+                else:
+                    frequency_2[qubit], fitted_parameters_2[qubit], _ = fit_result
 
     return DispersiveShiftQutritResults(
         frequency_state_zero=frequency_0,

src/qibocal/protocols/characterization/qubit_spectroscopy.py

@@ -131,20 +131,21 @@ def _fit(data: QubitSpectroscopyData) -> QubitSpectroscopyResults:
     error_fit_pars = {}
     chi2 = {}
     for qubit in qubits:
-        freq, fitted_params, perr = lorentzian_fit(
+        fit_result = lorentzian_fit(
             data[qubit], resonator_type=data.resonator_type, fit="qubit"
         )
-        frequency[qubit] = freq
-        fitted_parameters[qubit] = fitted_params
-        error_fit_pars[qubit] = perr
-        chi2[qubit] = (
-            chi2_reduced(
-                data[qubit].freq,
-                lorentzian(data[qubit].freq, *fitted_params),
-                data[qubit].error_signal,
-            ),
-            np.sqrt(2 / len(data[qubit].freq)),
-        )
+        if fit_result is not None:
+            frequency[qubit], fitted_parameters[qubit], error_fit_pars[qubit] = (
+                fit_result
+            )
+            chi2[qubit] = (
+                chi2_reduced(
+                    data[qubit].freq,
+                    lorentzian(data[qubit].freq, *fitted_parameters[qubit]),
+                    data[qubit].error_signal,
+                ),
+                np.sqrt(2 / len(data[qubit].freq)),
+            )
     return QubitSpectroscopyResults(
         frequency=frequency,
         fitted_parameters=fitted_parameters,

src/qibocal/protocols/characterization/qubit_spectroscopy_ef.py

@@ -47,6 +47,7 @@ def _fit_ef(data: QubitSpectroscopyEFData) -> QubitSpectroscopyEFResults:
     anharmoncities = {
         qubit: data.drive_frequencies[qubit] - results.frequency[qubit]
         for qubit in data.qubits
+        if qubit in results
     }
     params = asdict(results)
     params.update({"anharmonicity": anharmoncities})

src/qibocal/protocols/characterization/rabi/amplitude.py

@@ -168,24 +168,21 @@ def _fit(data: RabiAmplitudeData) -> RabiAmplitudeResults:
             pi_pulse_parameter = (
                 popt[2] / 2 * utils.period_correction_factor(phase=popt[3])
             )
+            pi_pulse_amplitudes[qubit] = (pi_pulse_parameter, perr[2] / 2)
+            fitted_parameters[qubit] = popt.tolist()
+            durations = {key: (value, 0) for key, value in data.durations.items()}
+            chi2[qubit] = (
+                chi2_reduced(
+                    y,
+                    utils.rabi_amplitude_function(x, *popt),
+                    qubit_data.error,
+                ),
+                np.sqrt(2 / len(y)),
+            )
+
+        except Exception as e:
+            log.warning(f"Rabi fit failed for qubit {qubit} due to {e}.")
 
-        except:
-            log.warning("rabi_fit: the fitting was not succesful")
-            pi_pulse_parameter = 0
-            popt = [0] * 4
-            perr = [1] * 4
-
-        pi_pulse_amplitudes[qubit] = (pi_pulse_parameter, perr[2] / 2)
-        fitted_parameters[qubit] = popt.tolist()
-        durations = {key: (value, 0) for key, value in data.durations.items()}
-        chi2[qubit] = (
-            chi2_reduced(
-                y,
-                utils.rabi_amplitude_function(x, *popt),
-                qubit_data.error,
-            ),
-            np.sqrt(2 / len(y)),
-        )
     return RabiAmplitudeResults(pi_pulse_amplitudes, durations, fitted_parameters, chi2)
 
 

src/qibocal/protocols/characterization/rabi/amplitude_signal.py

@@ -160,14 +160,11 @@ def _fit(data: RabiAmplitudeVoltData) -> RabiAmplitudeVoltResults:
                 / 2
                 * utils.period_correction_factor(phase=translated_popt[3])
             )
+            pi_pulse_amplitudes[qubit] = pi_pulse_parameter
+            fitted_parameters[qubit] = translated_popt
 
-        except:
-            log.warning("rabi_fit: the fitting was not succesful")
-            pi_pulse_parameter = 0
-            fitted_parameters = [0] * 4
-
-        pi_pulse_amplitudes[qubit] = pi_pulse_parameter
-        fitted_parameters[qubit] = translated_popt
+        except Exception as e:
+            log.warning(f"Rabi fit failed for qubit {qubit} due to {e}.")
 
     return RabiAmplitudeVoltResults(
         pi_pulse_amplitudes, data.durations, fitted_parameters

src/qibocal/protocols/characterization/rabi/length.py

@@ -172,21 +172,20 @@ def _fit(data: RabiLengthData) -> RabiLengthResults:
             pi_pulse_parameter = (
                 popt[2] / 2 * utils.period_correction_factor(phase=popt[3])
             )
-        except:
-            log.warning("rabi_fit: the fitting was not succesful")
-            pi_pulse_parameter = 0
-            popt = [0] * 4 + [1]
-        durations[qubit] = (pi_pulse_parameter, perr[2] / 2)
-        fitted_parameters[qubit] = popt.tolist()
-        amplitudes = {key: (value, 0) for key, value in data.amplitudes.items()}
-        chi2[qubit] = (
-            chi2_reduced(
-                y,
-                utils.rabi_length_function(x, *popt),
-                qubit_data.error,
-            ),
-            np.sqrt(2 / len(y)),
-        )
+            durations[qubit] = (pi_pulse_parameter, perr[2] / 2)
+            fitted_parameters[qubit] = popt.tolist()
+            amplitudes = {key: (value, 0) for key, value in data.amplitudes.items()}
+            chi2[qubit] = (
+                chi2_reduced(
+                    y,
+                    utils.rabi_length_function(x, *popt),
+                    qubit_data.error,
+                ),
+                np.sqrt(2 / len(y)),
+            )
+        except Exception as e:
+            log.warning(f"Rabi fit failed for qubit {qubit} due to {e}.")
+
     return RabiLengthResults(durations, amplitudes, fitted_parameters, chi2)
 
 

src/qibocal/protocols/characterization/rabi/length_signal.py

@@ -166,13 +166,11 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
                 / 2
                 * utils.period_correction_factor(phase=translated_popt[3])
             )
+            durations[qubit] = pi_pulse_parameter
+            fitted_parameters[qubit] = translated_popt
 
-        except:
-            log.warning("rabi_fit: the fitting was not succesful")
-            pi_pulse_parameter = 0
-            translated_popt = [0, 0, 1, 0, 0]
-        durations[qubit] = pi_pulse_parameter
-        fitted_parameters[qubit] = translated_popt
+        except Exception as e:
+            log.warning(f"Rabi fit failed for qubit {qubit} due to {e}.")
 
     return RabiLengthVoltResults(durations, data.amplitudes, fitted_parameters)