From a8763d32d5840177451e92a1a3e299c3b8dff0b1 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Fri, 1 Dec 2023 15:47:42 +0400
Subject: [PATCH 01/13] fix fit

---
 .../protocols/characterization/rabi/length_signal.py   | 10 +++++-----
 src/qibocal/protocols/characterization/rabi/utils.py   | 10 ++++++++++
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/length_signal.py b/src/qibocal/protocols/characterization/rabi/length_signal.py
index d771920ae..22ff10f1e 100644
--- a/src/qibocal/protocols/characterization/rabi/length_signal.py
+++ b/src/qibocal/protocols/characterization/rabi/length_signal.py
@@ -131,7 +131,7 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
         x_min = np.min(rabi_parameter)
         x_max = np.max(rabi_parameter)
         x = (rabi_parameter - x_min) / (x_max - x_min)
-        y = (voltages - y_min) / (y_max - y_min)
+        y = (voltages - y_min) / (y_max - y_min) - 1 / 2
 
         # Guessing period using fourier transform
         ft = np.fft.rfft(y)
@@ -141,21 +141,21 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
         # 0.5 hardcoded guess for less than one oscillation
         f = x[index] / (x[1] - x[0]) if index is not None else 0.5
 
-        pguess = [0.5, 0.5, 1 / f, np.pi / 2, 0]
+        pguess = [0, np.sign(y[0]) * 0.5, 1 / f, 0, 0]
         try:
             popt, _ = curve_fit(
-                utils.rabi_length_fit,
+                utils.rabi_length_signal,
                 x,
                 y,
                 p0=pguess,
                 maxfev=100000,
                 bounds=(
-                    [0, 0, 0, -np.pi, 0],
+                    [0, -1, 0, -np.pi, 0],
                     [1, 1, np.inf, np.pi, np.inf],
                 ),
             )
             translated_popt = [  # change it according to the fit function
-                (y_max - y_min) * popt[0] + y_min,
+                (y_max - y_min) * (popt[0] + 1 / 2) + y_min,
                 (y_max - y_min) * popt[1] * np.exp(x_min * popt[4] / (x_max - x_min)),
                 popt[2] * (x_max - x_min),
                 popt[3] - 2 * np.pi * x_min / popt[2] / (x_max - x_min),
diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 4caa2c293..792b9a33d 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -25,6 +25,16 @@ def rabi_length_fit(x, p0, p1, p2, p3, p4):
     return p0 + p1 * np.sin(2 * np.pi * x / p2 + p3) * np.exp(-x * p4)
 
 
+def rabi_length_signal(x, p0, p1, p2, p3, p4):
+    # A fit to Superconducting Qubit Rabi Oscillation
+    #   Offset                       : p[0]
+    #   Oscillation amplitude        : p[1]
+    #   Period    T                  : 1/p[2]
+    #   Phase                        : p[3]
+    #   Arbitrary parameter T_2      : 1/p[4]
+    return p0 + p1 * np.cos(2 * np.pi * x / p2 + p3) * np.exp(-x * p4)
+
+
 def plot(data, qubit, fit):
     if "RabiAmplitude" in data.__class__.__name__:
         quantity = "amp"

From 4fac79a3f0c28e257ae6c7cad32e0285c5e706e5 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Fri, 1 Dec 2023 16:08:21 +0400
Subject: [PATCH 02/13] fix plot

---
 src/qibocal/protocols/characterization/rabi/length_signal.py | 1 -
 src/qibocal/protocols/characterization/rabi/utils.py         | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/length_signal.py b/src/qibocal/protocols/characterization/rabi/length_signal.py
index 22ff10f1e..61296f66f 100644
--- a/src/qibocal/protocols/characterization/rabi/length_signal.py
+++ b/src/qibocal/protocols/characterization/rabi/length_signal.py
@@ -166,7 +166,6 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
             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
 
diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 792b9a33d..d90f95ce1 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -43,7 +43,7 @@ def plot(data, qubit, fit):
     elif data.__class__.__name__ == "RabiLengthVoltData":
         quantity = "length"
         title = "Time (ns)"
-        fitting = rabi_length_fit
+        fitting = rabi_length_signal
 
     figures = []
     fitting_report = ""

From 47b9e1b16cd1592a534b5582ec4b6605c6b5aa33 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Mon, 4 Dec 2023 11:42:39 +0400
Subject: [PATCH 03/13] propagate fit in rabi with probabilities

---
 .../protocols/characterization/rabi/length.py |  4 +--
 .../characterization/rabi/length_signal.py    |  2 +-
 .../protocols/characterization/rabi/utils.py  | 29 ++++++++-----------
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/length.py b/src/qibocal/protocols/characterization/rabi/length.py
index 83bb332be..f8b09c700 100644
--- a/src/qibocal/protocols/characterization/rabi/length.py
+++ b/src/qibocal/protocols/characterization/rabi/length.py
@@ -154,10 +154,10 @@ def _fit(data: RabiLengthData) -> RabiLengthResults:
         index = local_maxima[0] if len(local_maxima) > 0 else None
         # 0.5 hardcoded guess for less than one oscillation
         f = x[index] / (x[1] - x[0]) if index is not None else 0.5
-        pguess = [0.5, 0.5, np.max(x) / f, np.pi / 2, 0]
+        pguess = [0.5, 0.5, np.max(x) / f, 0, 0]
         try:
             popt, perr = curve_fit(
-                utils.rabi_length_fit,
+                utils.rabi_function,
                 x,
                 y,
                 p0=pguess,
diff --git a/src/qibocal/protocols/characterization/rabi/length_signal.py b/src/qibocal/protocols/characterization/rabi/length_signal.py
index 22ff10f1e..0663380ed 100644
--- a/src/qibocal/protocols/characterization/rabi/length_signal.py
+++ b/src/qibocal/protocols/characterization/rabi/length_signal.py
@@ -144,7 +144,7 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
         pguess = [0, np.sign(y[0]) * 0.5, 1 / f, 0, 0]
         try:
             popt, _ = curve_fit(
-                utils.rabi_length_signal,
+                utils.rabi_function,
                 x,
                 y,
                 p0=pguess,
diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 792b9a33d..d4e460ec3 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -15,23 +15,18 @@ def rabi_amplitude_fit(x, p0, p1, p2, p3):
     return p0 + p1 * np.sin(2 * np.pi * x / p2 + p3)
 
 
-def rabi_length_fit(x, p0, p1, p2, p3, p4):
-    # A fit to Superconducting Qubit Rabi Oscillation
-    #   Offset                       : p[0]
-    #   Oscillation amplitude        : p[1]
-    #   Period    T                  : 1/p[2]
-    #   Phase                        : p[3]
-    #   Arbitrary parameter T_2      : 1/p[4]
-    return p0 + p1 * np.sin(2 * np.pi * x / p2 + p3) * np.exp(-x * p4)
-
-
-def rabi_length_signal(x, p0, p1, p2, p3, p4):
-    # A fit to Superconducting Qubit Rabi Oscillation
-    #   Offset                       : p[0]
-    #   Oscillation amplitude        : p[1]
-    #   Period    T                  : 1/p[2]
-    #   Phase                        : p[3]
-    #   Arbitrary parameter T_2      : 1/p[4]
+def rabi_function(x, p0, p1, p2, p3, p4):
+    """
+    Fit function of rabi length signal experiment.
+
+    Args:
+        x: Input data.
+        p0: Offset.
+        p1: Oscillation Amplitude.
+        p2: Period.
+        p3: Phase.
+        p4: T2.
+    """
     return p0 + p1 * np.cos(2 * np.pi * x / p2 + p3) * np.exp(-x * p4)
 
 

From ac8d71b9bb165d564bba0dcdc9d7cac5d11d84d8 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Mon, 4 Dec 2023 12:05:54 +0400
Subject: [PATCH 04/13] rafactor rabi utils

---
 .../characterization/rabi/amplitude.py        |  4 +-
 .../characterization/rabi/amplitude_signal.py |  2 +-
 .../protocols/characterization/rabi/length.py |  4 +-
 .../characterization/rabi/length_signal.py    |  2 +-
 .../protocols/characterization/rabi/utils.py  | 49 +++++++++----------
 5 files changed, 30 insertions(+), 31 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/amplitude.py b/src/qibocal/protocols/characterization/rabi/amplitude.py
index 15a7ebac6..4328a7f5a 100644
--- a/src/qibocal/protocols/characterization/rabi/amplitude.py
+++ b/src/qibocal/protocols/characterization/rabi/amplitude.py
@@ -153,7 +153,7 @@ def _fit(data: RabiAmplitudeData) -> RabiAmplitudeResults:
         pguess = [0.5, 0.5, 1 / f, np.pi / 2]
         try:
             popt, perr = curve_fit(
-                utils.rabi_amplitude_fit,
+                utils.rabi_amplitude_function,
                 x,
                 y,
                 p0=pguess,
@@ -179,7 +179,7 @@ def _fit(data: RabiAmplitudeData) -> RabiAmplitudeResults:
         chi2[qubit] = (
             chi2_reduced(
                 y,
-                utils.rabi_amplitude_fit(x, *popt),
+                utils.rabi_amplitude_function(x, *popt),
                 qubit_data.error,
             ),
             np.sqrt(2 / len(y)),
diff --git a/src/qibocal/protocols/characterization/rabi/amplitude_signal.py b/src/qibocal/protocols/characterization/rabi/amplitude_signal.py
index 37d7c47ec..5447cd65f 100644
--- a/src/qibocal/protocols/characterization/rabi/amplitude_signal.py
+++ b/src/qibocal/protocols/characterization/rabi/amplitude_signal.py
@@ -139,7 +139,7 @@ def _fit(data: RabiAmplitudeVoltData) -> RabiAmplitudeVoltResults:
         pguess = [0.5, 1, 1 / f, np.pi / 2]
         try:
             popt, _ = curve_fit(
-                utils.rabi_amplitude_fit,
+                utils.rabi_amplitude_function,
                 x,
                 y,
                 p0=pguess,
diff --git a/src/qibocal/protocols/characterization/rabi/length.py b/src/qibocal/protocols/characterization/rabi/length.py
index f8b09c700..2e069a7d5 100644
--- a/src/qibocal/protocols/characterization/rabi/length.py
+++ b/src/qibocal/protocols/characterization/rabi/length.py
@@ -157,7 +157,7 @@ def _fit(data: RabiLengthData) -> RabiLengthResults:
         pguess = [0.5, 0.5, np.max(x) / f, 0, 0]
         try:
             popt, perr = curve_fit(
-                utils.rabi_function,
+                utils.rabi_length_function,
                 x,
                 y,
                 p0=pguess,
@@ -180,7 +180,7 @@ def _fit(data: RabiLengthData) -> RabiLengthResults:
         chi2[qubit] = (
             chi2_reduced(
                 y,
-                utils.rabi_length_fit(x, *popt),
+                utils.rabi_length_function(x, *popt),
                 qubit_data.error,
             ),
             np.sqrt(2 / len(y)),
diff --git a/src/qibocal/protocols/characterization/rabi/length_signal.py b/src/qibocal/protocols/characterization/rabi/length_signal.py
index 0663380ed..16c35ed76 100644
--- a/src/qibocal/protocols/characterization/rabi/length_signal.py
+++ b/src/qibocal/protocols/characterization/rabi/length_signal.py
@@ -144,7 +144,7 @@ def _fit(data: RabiLengthVoltData) -> RabiLengthVoltResults:
         pguess = [0, np.sign(y[0]) * 0.5, 1 / f, 0, 0]
         try:
             popt, _ = curve_fit(
-                utils.rabi_function,
+                utils.rabi_amplitude_function,
                 x,
                 y,
                 p0=pguess,
diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index d4e460ec3..386e03e8f 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -5,17 +5,21 @@
 from ..utils import COLORBAND, COLORBAND_LINE, table_dict, table_html
 
 
-def rabi_amplitude_fit(x, p0, p1, p2, p3):
-    # A fit to Superconducting Qubit Rabi Oscillation
-    #   Offset                       : p[0]
-    #   Oscillation amplitude        : p[1]
-    #   Period    T                  : 1/p[2]
-    #   Phase                        : p[3]
-    #   Arbitrary parameter T_2      : 1/p[4]
+def rabi_amplitude_function(x, p0, p1, p2, p3):
+    """
+    Fit function of rabi amplitude signal experiment.
+
+    Args:
+        x: Input data.
+        p0: Offset.
+        p1: Oscillation Amplitude.
+        p2: Period.
+        p3: Phase.
+    """
     return p0 + p1 * np.sin(2 * np.pi * x / p2 + p3)
 
 
-def rabi_function(x, p0, p1, p2, p3, p4):
+def rabi_length_function(x, p0, p1, p2, p3, p4):
     """
     Fit function of rabi length signal experiment.
 
@@ -31,14 +35,7 @@ def rabi_function(x, p0, p1, p2, p3, p4):
 
 
 def plot(data, qubit, fit):
-    if "RabiAmplitude" in data.__class__.__name__:
-        quantity = "amp"
-        title = "Amplitude (dimensionless)"
-        fitting = rabi_amplitude_fit
-    elif data.__class__.__name__ == "RabiLengthVoltData":
-        quantity = "length"
-        title = "Time (ns)"
-        fitting = rabi_length_fit
+    quantity, title, fitting = extract_rabi(data)
 
     figures = []
     fitting_report = ""
@@ -124,15 +121,7 @@ def plot(data, qubit, fit):
 
 
 def plot_probabilities(data, qubit, fit):
-    if data.__class__.__name__ == "RabiAmplitudeData":
-        quantity = "amp"
-        title = "Amplitude (dimensionless)"
-        fitting = rabi_amplitude_fit
-    elif data.__class__.__name__ == "RabiLengthData":
-        quantity = "length"
-        title = "Time (ns)"
-        fitting = rabi_length_fit
-
+    quantity, title, fitting = extract_rabi(data)
     figures = []
     fitting_report = ""
 
@@ -200,3 +189,13 @@ def plot_probabilities(data, qubit, fit):
     figures.append(fig)
 
     return figures, fitting_report
+
+
+def extract_rabi(data):
+    """
+    Extract Rabi fit info.
+    """
+    if data.__class__.__name__ == "RabiAmplitudeData":
+        return "amp", "Amplitude (dimensionless)", rabi_amplitude_function
+    elif data.__class__.__name__ == "RabiLengthData":
+        return "length", "Time (ns)", rabi_length_function

From db6f4dab4a2f6d07b6f386a7574f2ba166149132 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Mon, 4 Dec 2023 12:12:24 +0400
Subject: [PATCH 05/13] update docs

---
 src/qibocal/protocols/characterization/rabi/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 5f1b0de4d..7dbeba772 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -7,7 +7,7 @@
 
 def rabi_amplitude_function(x, p0, p1, p2, p3):
     """
-    Fit function of rabi amplitude signal experiment.
+    Fit function of Rabi amplitude signal experiment.
 
     Args:
         x: Input data.
@@ -21,7 +21,7 @@ def rabi_amplitude_function(x, p0, p1, p2, p3):
 
 def rabi_length_function(x, p0, p1, p2, p3, p4):
     """
-    Fit function of rabi length signal experiment.
+    Fit function of Rabi length signal experiment.
 
     Args:
         x: Input data.

From a72fd9f89edda1ddb8a3d9d8edd5243832bde140 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Mon, 4 Dec 2023 12:50:37 +0400
Subject: [PATCH 06/13] fix tests

---
 src/qibocal/protocols/characterization/rabi/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 7dbeba772..4062ee5bb 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -194,7 +194,7 @@ def extract_rabi(data):
     """
     Extract Rabi fit info.
     """
-    if data.__class__.__name__ == "RabiAmplitudeData":
+    if "RabiAmplitude" in data.__class__.__name__:
         return "amp", "Amplitude (dimensionless)", rabi_amplitude_function
-    elif data.__class__.__name__ == "RabiLengthData":
+    elif "RabiLength" in data.__class__.__name__:
         return "length", "Time (ns)", rabi_length_function

From bba59dd894599e8121e79b5a2ef9c3072afdd1a7 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 11:01:40 +0400
Subject: [PATCH 07/13] change fit parameters names

---
 .../protocols/characterization/rabi/utils.py  | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 4062ee5bb..b07edcc6f 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -5,33 +5,26 @@
 from ..utils import COLORBAND, COLORBAND_LINE, table_dict, table_html
 
 
-def rabi_amplitude_function(x, p0, p1, p2, p3):
+def rabi_amplitude_function(x, offset, amplitude, period, phase):
     """
     Fit function of Rabi amplitude signal experiment.
 
     Args:
         x: Input data.
-        p0: Offset.
-        p1: Oscillation Amplitude.
-        p2: Period.
-        p3: Phase.
     """
-    return p0 + p1 * np.sin(2 * np.pi * x / p2 + p3)
+    return offset + amplitude * np.sin(2 * np.pi * x / period + phase)
 
 
-def rabi_length_function(x, p0, p1, p2, p3, p4):
+def rabi_length_function(x, offset, amplitude, period, phase, t2_inv):
     """
     Fit function of Rabi length signal experiment.
 
     Args:
         x: Input data.
-        p0: Offset.
-        p1: Oscillation Amplitude.
-        p2: Period.
-        p3: Phase.
-        p4: T2.
     """
-    return p0 + p1 * np.cos(2 * np.pi * x / p2 + p3) * np.exp(-x * p4)
+    return offset + amplitude * np.cos(2 * np.pi * x / period + phase) * np.exp(
+        -x * t2_inv
+    )
 
 
 def plot(data, qubit, fit):

From ef690485b34588d8ac3595f08e5fc360c7c1ebb8 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 11:15:04 +0400
Subject: [PATCH 08/13] raise error in extract rabi

---
 src/qibocal/protocols/characterization/rabi/utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index b07edcc6f..a0c084701 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -191,3 +191,4 @@ def extract_rabi(data):
         return "amp", "Amplitude (dimensionless)", rabi_amplitude_function
     elif "RabiLength" in data.__class__.__name__:
         return "length", "Time (ns)", rabi_length_function
+    raise RuntimeError("Data has to be a data structure of the Rabi routines.")

From 747551f4d8c121fb2b49991211bfb980d6f3c449 Mon Sep 17 00:00:00 2001
From: Edoardo Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 16:10:00 +0400
Subject: [PATCH 09/13] Update
 src/qibocal/protocols/characterization/rabi/utils.py

Co-authored-by: Gabriele Palazzo <73099233+GabrielePalazzo@users.noreply.github.com>
---
 src/qibocal/protocols/characterization/rabi/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 44f84203e..04c35941f 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -188,5 +188,5 @@ def extract_rabi(data):
     if "RabiAmplitude" in data.__class__.__name__:
         return "amp", "Amplitude (dimensionless)", rabi_amplitude_function
     elif "RabiLength" in data.__class__.__name__:
-        return "length", "Time (ns)", rabi_length_function
+        return "length", "Time [ns]", rabi_length_function
     raise RuntimeError("Data has to be a data structure of the Rabi routines.")

From 03ecf92f5a21c762a649a4c42dcc47ac55c5c8d8 Mon Sep 17 00:00:00 2001
From: Edoardo Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 16:10:15 +0400
Subject: [PATCH 10/13] Update
 src/qibocal/protocols/characterization/rabi/utils.py

Co-authored-by: Gabriele Palazzo <73099233+GabrielePalazzo@users.noreply.github.com>
---
 src/qibocal/protocols/characterization/rabi/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 04c35941f..65dcbfdc1 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -186,7 +186,7 @@ def extract_rabi(data):
     Extract Rabi fit info.
     """
     if "RabiAmplitude" in data.__class__.__name__:
-        return "amp", "Amplitude (dimensionless)", rabi_amplitude_function
+        return "amp", "Amplitude [dimensionless]", rabi_amplitude_function
     elif "RabiLength" in data.__class__.__name__:
         return "length", "Time [ns]", rabi_length_function
     raise RuntimeError("Data has to be a data structure of the Rabi routines.")

From 2af932ba5b375512036a3c66f79f24c328763f80 Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 17:54:22 +0400
Subject: [PATCH 11/13] test extract_rabi

---
 tests/test_protocols.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/test_protocols.py b/tests/test_protocols.py
index 9881ec46f..7acf591ba 100644
--- a/tests/test_protocols.py
+++ b/tests/test_protocols.py
@@ -7,6 +7,14 @@
 from qibolab import create_platform
 
 from qibocal.cli._base import command
+from qibocal.protocols.characterization.rabi.amplitude import RabiAmplitudeData
+from qibocal.protocols.characterization.rabi.ef import RabiAmplitudeEFData
+from qibocal.protocols.characterization.rabi.length import RabiLengthData
+from qibocal.protocols.characterization.rabi.utils import (
+    extract_rabi,
+    rabi_amplitude_function,
+    rabi_length_function,
+)
 
 PATH_TO_RUNCARD = pathlib.Path(__file__).parent / "runcards/protocols.yml"
 PLATFORM = create_platform("dummy")
@@ -119,4 +127,19 @@ def test_fit_command(runcard, tmp_path):
     assert results_plot.exit_code == 0
 
 
+def test_extract_rabi():
+    assert extract_rabi(RabiAmplitudeData()) == (
+        "amp",
+        "Amplitude (dimensionless)",
+        rabi_amplitude_function,
+    )
+    assert extract_rabi(RabiLengthData()) == (
+        "length",
+        "Time (ns)",
+        rabi_length_function,
+    )
+    with pytest.raises(RuntimeError):
+        extract_rabi(RabiAmplitudeEFData)
+
+
 # TODO: compare report by calling qq report

From e6c26b9b7108cff005d24d1a374b0bc8be899e9f Mon Sep 17 00:00:00 2001
From: Edoardo-Pedicillo <edoardo.pedicillo@gmail.com>
Date: Tue, 5 Dec 2023 18:10:13 +0400
Subject: [PATCH 12/13] fix tests

---
 tests/test_protocols.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_protocols.py b/tests/test_protocols.py
index 7acf591ba..d94fd7f3f 100644
--- a/tests/test_protocols.py
+++ b/tests/test_protocols.py
@@ -130,12 +130,12 @@ def test_fit_command(runcard, tmp_path):
 def test_extract_rabi():
     assert extract_rabi(RabiAmplitudeData()) == (
         "amp",
-        "Amplitude (dimensionless)",
+        "Amplitude [dimensionless]",
         rabi_amplitude_function,
     )
     assert extract_rabi(RabiLengthData()) == (
         "length",
-        "Time (ns)",
+        "Time [ns]",
         rabi_length_function,
     )
     with pytest.raises(RuntimeError):

From 4e2863f49319974601c96abbb92935d07eeb7ef7 Mon Sep 17 00:00:00 2001
From: Edoardo Pedicillo <edoardo.pedicillo@gmail.com>
Date: Wed, 6 Dec 2023 10:26:13 +0400
Subject: [PATCH 13/13] Update
 src/qibocal/protocols/characterization/rabi/utils.py

Co-authored-by: Alessandro Candido <candido.ale@gmail.com>
---
 src/qibocal/protocols/characterization/rabi/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibocal/protocols/characterization/rabi/utils.py b/src/qibocal/protocols/characterization/rabi/utils.py
index 65dcbfdc1..00b1e8aeb 100644
--- a/src/qibocal/protocols/characterization/rabi/utils.py
+++ b/src/qibocal/protocols/characterization/rabi/utils.py
@@ -187,6 +187,6 @@ def extract_rabi(data):
     """
     if "RabiAmplitude" in data.__class__.__name__:
         return "amp", "Amplitude [dimensionless]", rabi_amplitude_function
-    elif "RabiLength" in data.__class__.__name__:
+    if "RabiLength" in data.__class__.__name__:
         return "length", "Time [ns]", rabi_length_function
     raise RuntimeError("Data has to be a data structure of the Rabi routines.")