improve intial parameters guess Lorentzian fit

src/qibocal/protocols/utils.py

@@ -108,26 +108,35 @@ def lorentzian_fit(data, resonator_type=None, fit=None):
     voltages = data.signal
 
     # Guess parameters for Lorentzian max or min
-    # TODO: probably this is not working on HW
-    guess_offset = np.mean(
-        voltages[np.abs(voltages - np.mean(voltages)) < np.std(voltages)]
-    )
-    guess_slope = 0.0
+    guess_slope = (voltages[-1] - voltages[0]) / (frequencies[-1] - frequencies[0])
+    guess_offset = voltages[0] - guess_slope * frequencies[0]
+    guess_background = guess_offset + guess_slope * frequencies
+    voltages_no_background = voltages - guess_background
+
     if (resonator_type == "3D" and fit == "resonator") or (
         resonator_type == "2D" and fit == "qubit"
     ):
-        guess_center = frequencies[
-            np.argmax(voltages)
-        ]  # Argmax = Returns the indices of the maximum values along an axis.
-        guess_sigma = abs(frequencies[np.argmin(voltages)] - guess_center)
-        guess_amp = (np.max(voltages) - guess_offset) * guess_sigma * np.pi
-
+        guess_center = frequencies[np.argmax(voltages_no_background)]
+        guess_peak_height = voltages_no_background.max()
+        indices_beyond_half = np.where(voltages_no_background > guess_peak_height / 2)[
+            0
+        ]
     else:
-        guess_center = frequencies[
-            np.argmin(voltages)
-        ]  # Argmin = Returns the indices of the minimum values along an axis.
-        guess_sigma = abs(frequencies[np.argmax(voltages)] - guess_center)
-        guess_amp = (np.min(voltages) - guess_offset) * guess_sigma * np.pi
+        guess_center = frequencies[np.argmin(voltages_no_background)]
+        guess_peak_height = voltages_no_background.min()
+        indices_beyond_half = np.where(voltages_no_background < guess_peak_height / 2)[
+            0
+        ]
+
+    if len(indices_beyond_half) >= 1:
+        guess_sigma = (
+            frequencies[indices_beyond_half[-1]] - frequencies[indices_beyond_half[0]]
+        ) / 2
+    else:
+        # if there is no clear peak, we give a high flexibility
+        guess_sigma = frequencies[-1] - frequencies[0]
+
+    guess_amp = guess_peak_height * guess_sigma * np.pi
 
     initial_parameters = [
         guess_amp,
@@ -136,30 +145,26 @@ def lorentzian_fit(data, resonator_type=None, fit=None):
         guess_offset,
         guess_slope,
     ]
+    freq_domain_size = frequencies[-1] - frequencies[0]
+    bounds = (
+        [-np.inf, frequencies[0], 0.0, -np.inf, -np.inf],
+        [np.inf, frequencies[-1], freq_domain_size, np.inf, np.inf],
+    )
+
     # fit the model with the data and guessed parameters
     try:
-        sigma = (
-            data.error_signal
-            if (hasattr(data, "error_signal") and not np.isnan(data.error_signal).any())
-            else None
-        )
-
-        fit_parameters, perr = curve_fit(
+        fit_parameters, parameters_cov = curve_fit(
             lorentzian,
             frequencies,
             voltages,
             p0=initial_parameters,
-            sigma=sigma,
+            bounds=bounds,
         )
         # The output results are stored in a json, but ndarray is not JSON serializable,
         # so the parameters are converted to list.
-        perr = (
-            np.sqrt(np.diag(perr)).tolist()
-            if sigma is not None
-            else [0] * len(initial_parameters)
-        )
+        parameter_errors = np.sqrt(np.diag(parameters_cov)).tolist()
         model_parameters = fit_parameters.tolist()
-        return model_parameters[1] * GHZ_TO_HZ, model_parameters, perr
+        return model_parameters[1] * GHZ_TO_HZ, model_parameters, parameter_errors
     except RuntimeError as e:
         log.warning(f"Lorentzian fit not successful due to {e}")
 

tests/conftest.py

@@ -10,6 +10,9 @@
     "mock",
 ]
 
+TEST_FILE_DIR = Path(__file__).resolve().parent
+PATH_TESTING_DATA = TEST_FILE_DIR / "tests_data"
+
 
 @pytest.fixture(autouse=True)
 def cd(tmp_path_factory: pytest.TempdirFactory):

tests/test_fit_functions.py

@@ -1,28 +1,17 @@
 import json
 import math
-from pathlib import Path
 
 import numpy as np
-import pytest
+from conftest import TEST_FILE_DIR
 
 from qibocal.protocols.rabi.utils import (
     fit_amplitude_function as rabi_fit_amplitude_function,
 )
 from qibocal.protocols.rabi.utils import fit_length_function as rabi_fit_length_function
 from qibocal.protocols.ramsey.processing import fitting as ramsey_fitting
 from qibocal.protocols.ramsey.processing import process_fit as ramsey_process_fit
-from qibocal.protocols.resonator_spectroscopies.resonator_spectroscopy import (
-    ResonatorSpectroscopyData,
-    ResonatorSpectroscopyResults,
-)
-from qibocal.protocols.resonator_spectroscopies.resonator_spectroscopy import (
-    _fit as resonator_spectroscopy_fit,
-)
 from qibocal.protocols.utils import fallback_period, guess_period
 
-TEST_FILE_DIR = Path(__file__).resolve().parent
-PATH_TESTING_DATA = TEST_FILE_DIR / "tests_data"
-
 
 def test_ramsey_fit():
     test_folder = TEST_FILE_DIR / "ramsey_fit_data"
@@ -158,13 +147,3 @@ def test_rabi_fit():
                     true_duration = results['"duration"'][f]
 
                 assert math.isclose(true_duration, new_duration, rel_tol=2.5e-2)
-
-
-def test_resonator_spectroscopy_fit():
-    results_folder = PATH_TESTING_DATA / "resonator_spectroscopy-0"
-    data = ResonatorSpectroscopyData.load(results_folder)
-    expected = ResonatorSpectroscopyResults.load(results_folder)
-    assert data is not None and expected is not None
-    fitted = resonator_spectroscopy_fit(data)
-    qubit = 0  # the data contains only qubit 0
-    assert fitted.frequency[qubit] == pytest.approx(expected.frequency[qubit])

tests/test_resonator_spectroscopy.py

@@ -0,0 +1,20 @@
+import pytest
+from conftest import PATH_TESTING_DATA
+
+from qibocal.protocols.resonator_spectroscopies.resonator_spectroscopy import (
+    ResonatorSpectroscopyData,
+    ResonatorSpectroscopyResults,
+)
+from qibocal.protocols.resonator_spectroscopies.resonator_spectroscopy import (
+    _fit as resonator_spectroscopy_fit,
+)
+
+
+def test_resonator_spectroscopy_fit():
+    results_folder = PATH_TESTING_DATA / "resonator_spectroscopy-0"
+    data = ResonatorSpectroscopyData.load(results_folder)
+    expected = ResonatorSpectroscopyResults.load(results_folder)
+    assert data is not None and expected is not None
+    fitted = resonator_spectroscopy_fit(data)
+    qubit = 0  # the data contains only qubit 0
+    assert fitted.frequency[qubit] == pytest.approx(expected.frequency[qubit])