From 47804e167725970c0ffb5d3b6ac86cd94504c906 Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Tue, 2 Apr 2024 16:09:11 +0400
Subject: [PATCH] feat: Improve data selection using confidence interval

---
 .../flux_dependence/qubit_flux_dependence.py  | 13 +-----
 .../resonator_flux_dependence.py              | 13 +-----
 .../characterization/flux_dependence/utils.py | 44 +++++++++++++------
 3 files changed, 32 insertions(+), 38 deletions(-)

diff --git a/src/qibocal/protocols/characterization/flux_dependence/qubit_flux_dependence.py b/src/qibocal/protocols/characterization/flux_dependence/qubit_flux_dependence.py
index 73186b6f7..9f523b73a 100644
--- a/src/qibocal/protocols/characterization/flux_dependence/qubit_flux_dependence.py
+++ b/src/qibocal/protocols/characterization/flux_dependence/qubit_flux_dependence.py
@@ -196,18 +196,7 @@ def _fit(data: QubitFluxData) -> QubitFluxResults:
         frequencies = qubit_data.freq
         signal = qubit_data.signal
 
-        if data.resonator_type == "3D":
-            frequencies, biases = utils.extract_min_feature(
-                frequencies,
-                biases,
-                signal,
-            )
-        else:
-            frequencies, biases = utils.extract_max_feature(
-                frequencies,
-                biases,
-                signal,
-            )
+        frequencies, biases = utils.extract_feature(frequencies, biases, signal, "max")
 
         try:
             popt = curve_fit(
diff --git a/src/qibocal/protocols/characterization/flux_dependence/resonator_flux_dependence.py b/src/qibocal/protocols/characterization/flux_dependence/resonator_flux_dependence.py
index d3050baf2..3f262674e 100644
--- a/src/qibocal/protocols/characterization/flux_dependence/resonator_flux_dependence.py
+++ b/src/qibocal/protocols/characterization/flux_dependence/resonator_flux_dependence.py
@@ -299,18 +299,7 @@ def _fit(data: ResonatorFluxData) -> ResonatorFluxResults:
         frequencies = qubit_data.freq
         signal = qubit_data.signal
 
-        if data.resonator_type == "3D":
-            frequencies, biases = utils.extract_max_feature(
-                frequencies,
-                biases,
-                signal,
-            )
-        else:
-            frequencies, biases = utils.extract_min_feature(
-                frequencies,
-                biases,
-                signal,
-            )
+        frequencies, biases = utils.extract_feature(frequencies, biases, signal, "min")
 
         try:
             popt = curve_fit(
diff --git a/src/qibocal/protocols/characterization/flux_dependence/utils.py b/src/qibocal/protocols/characterization/flux_dependence/utils.py
index 3884d2cb4..3127eaee9 100644
--- a/src/qibocal/protocols/characterization/flux_dependence/utils.py
+++ b/src/qibocal/protocols/characterization/flux_dependence/utils.py
@@ -4,6 +4,11 @@
 
 from ..utils import HZ_TO_GHZ
 
+CONFIDENCE_INTERVAL_FIRST_MASK = 99
+"""Confidence interval used to mask flux data."""
+CONFIDENCE_INTERVAL_SECOND_MASK = 70
+"""Confidence interval used to clean outliers."""
+
 
 def is_crosstalk(data):
     """Check if keys are tuple which corresponds to crosstalk data structure."""
@@ -231,22 +236,33 @@ def transmon_readout_frequency(x, w_max, d, element, offset, resonator_freq, g):
     )
 
 
-def extract_min_feature(freq, bias, signal, threshold=1.5):
-    """Extract min feature using SNR."""
-    mean_signal = np.mean(signal)
-    std_signal = np.std(signal)
-    snr_map = (signal - mean_signal) / std_signal
-    binary_mask = snr_map < -threshold
-    return freq[binary_mask], bias[binary_mask]
+def extract_feature(freq: np.ndarray, bias: np.ndarray, signal: np.ndarray, feat: str):
+    """Extract feature using confidence intervals.
 
+    A first mask is construct by looking at 99% confidence interval for each bias bin.
+    A second mask is applied by looking at 70% confidence interval to remove outliers.
+    """
 
-def extract_max_feature(freq, bias, signal, threshold=1.5):
-    """Extract max feature using SNR."""
-    mean_signal = np.mean(signal)
-    std_signal = np.std(signal)
-    snr_map = (signal - mean_signal) / std_signal
-    binary_mask = snr_map > threshold
-    return freq[binary_mask], bias[binary_mask]
+    masks = []
+    for bias_bin in np.unique(bias):
+        signal_fixed_bias = signal[bias == bias_bin]
+        min, max = np.percentile(
+            signal_fixed_bias,
+            [100 - CONFIDENCE_INTERVAL_FIRST_MASK, CONFIDENCE_INTERVAL_FIRST_MASK],
+        )
+        masks.append(
+            signal_fixed_bias < min if feat == "min" else signal_fixed_bias > max
+        )
+
+    first_mask = np.vstack(masks).ravel()
+    min, max = np.percentile(
+        signal[first_mask],
+        [100 - CONFIDENCE_INTERVAL_SECOND_MASK, CONFIDENCE_INTERVAL_SECOND_MASK],
+    )
+    second_mask = (
+        signal[first_mask] < max if feat == "min" else signal[first_mask] > min
+    )
+    return freq[first_mask][second_mask], bias[first_mask][second_mask]
 
 
 def qubit_flux_dependence_fit_bounds(qubit_frequency: float, bias: np.array):