Merge pull request #602 from qiboteam/fix_twpa_update

Fix update in readout optimization protocols

src/qibocal/protocols/characterization/readout_optimization/resonator_amplitude.py

@@ -34,6 +34,8 @@ class ResonatorAmplitudeParameters(Parameters):
     [
         ("error", np.float64),
         ("amp", np.float64),
+        ("angle", np.float64),
+        ("threshold", np.float64),
     ]
 )
 """Custom dtype for Optimization RO amplitude."""
@@ -54,6 +56,10 @@ class ResonatorAmplitudeResults(Results):
     """Lowest probability errors"""
     best_amp: dict[QubitId, list]
     """Amplitude with lowest error"""
+    best_angle: dict[QubitId, float]
+    """IQ angle that gives lower error."""
+    best_threshold: dict[QubitId, float]
+    """Thershold that gives lower error."""
 
 
 def _acquisition(
@@ -128,6 +134,8 @@ def _acquisition(
                 dict(
                     amp=np.array([new_amp]),
                     error=np.array([error]),
+                    angle=np.array([model.angle]),
+                    threshold=np.array([model.threshold]),
                 ),
             )
             platform.qubits[qubit].native_gates.MZ.amplitude = old_amp
@@ -138,14 +146,18 @@ def _acquisition(
 def _fit(data: ResonatorAmplitudeData) -> ResonatorAmplitudeResults:
     qubits = data.qubits
     best_amps = {}
+    best_angle = {}
+    best_threshold = {}
     lowest_err = {}
     for qubit in qubits:
         data_qubit = data[qubit]
         index_best_err = np.argmin(data_qubit["error"])
         lowest_err[qubit] = data_qubit["error"][index_best_err]
         best_amps[qubit] = data_qubit["amp"][index_best_err]
+        best_angle[qubit] = data_qubit["angle"][index_best_err]
+        best_threshold[qubit] = data_qubit["threshold"][index_best_err]
 
-    return ResonatorAmplitudeResults(lowest_err, best_amps)
+    return ResonatorAmplitudeResults(lowest_err, best_amps, best_angle, best_threshold)
 
 
 def _plot(data: ResonatorAmplitudeData, fit: ResonatorAmplitudeResults, qubit):
@@ -192,6 +204,8 @@ def _plot(data: ResonatorAmplitudeData, fit: ResonatorAmplitudeResults, qubit):
 
 def _update(results: ResonatorAmplitudeResults, platform: Platform, qubit: QubitId):
     update.readout_amplitude(results.best_amp[qubit], platform, qubit)
+    update.iq_angle(results.best_angle[qubit], platform, qubit)
+    update.threshold(results.best_threshold[qubit], platform, qubit)
 
 
 resonator_amplitude = Routine(_acquisition, _fit, _plot, _update)

src/qibocal/protocols/characterization/readout_optimization/resonator_frequency.py

@@ -34,10 +34,19 @@ class ResonatorFrequencyResults(Results):
     """Assignment fidelities."""
     best_freq: dict[QubitId, float]
     """Resonator Frequency with the highest assignment fidelity."""
+    best_angle: dict[QubitId, float]
+    """IQ angle that maximes assignment fidelity"""
+    best_threshold: dict[QubitId, float]
+    """Threshold that maximes assignment fidelity"""
 
 
 ResonatorFrequencyType = np.dtype(
-    [("freq", np.float64), ("assignment_fidelity", np.float64)]
+    [
+        ("freq", np.float64),
+        ("assignment_fidelity", np.float64),
+        ("angle", np.float64),
+        ("threshold", np.float64),
+    ]
 )
 """Custom dtype for Optimization RO frequency."""
 
@@ -144,6 +153,8 @@ def _acquisition(
                 dict(
                     freq=np.array([(ro_pulses[qubit].frequency + freq) * HZ_TO_GHZ]),
                     assignment_fidelity=np.array([model.assignment_fidelity]),
+                    angle=np.array([model.angle]),
+                    threshold=np.array([model.threshold]),
                 ),
             )
     return data
@@ -153,16 +164,22 @@ def _fit(data: ResonatorFrequencyData) -> ResonatorFrequencyResults:
     """Post-Processing for Optimization RO frequency"""
     qubits = data.qubits
     best_freq = {}
+    best_angle = {}
+    best_threshold = {}
     highest_fidelity = {}
     for qubit in qubits:
         data_qubit = data[qubit]
         index_best_fid = np.argmax(data_qubit["assignment_fidelity"])
         highest_fidelity[qubit] = data_qubit["assignment_fidelity"][index_best_fid]
         best_freq[qubit] = data_qubit["freq"][index_best_fid]
+        best_angle[qubit] = data_qubit["angle"][index_best_fid]
+        best_threshold[qubit] = data_qubit["threshold"][index_best_fid]
 
     return ResonatorFrequencyResults(
         fidelities=highest_fidelity,
         best_freq=best_freq,
+        best_angle=best_angle,
+        best_threshold=best_threshold,
     )
 
 
@@ -210,6 +227,8 @@ def _plot(data: ResonatorFrequencyData, fit: ResonatorFrequencyResults, qubit):
 
 def _update(results: ResonatorFrequencyResults, platform: Platform, qubit: QubitId):
     update.readout_frequency(results.best_freq[qubit], platform, qubit)
+    update.threshold(results.best_threshold[qubit], platform, qubit)
+    update.iq_angle(results.best_angle[qubit], platform, qubit)
 
 
 resonator_frequency = Routine(_acquisition, _fit, _plot, _update)

src/qibocal/protocols/characterization/readout_optimization/twpa_calibration/frequency.py

@@ -9,6 +9,9 @@
 from qibocal import update
 from qibocal.auto.operation import Data, Parameters, Qubits, Results, Routine
 from qibocal.protocols.characterization import classification
+from qibocal.protocols.characterization.readout_optimization.resonator_frequency import (
+    ResonatorFrequencyType,
+)
 from qibocal.protocols.characterization.utils import HZ_TO_GHZ, table_dict, table_html
 
 
@@ -22,14 +25,6 @@ class TwpaFrequencyParameters(Parameters):
     """Frequency step [Hz]"""
 
 
-TwpaFrequencyType = np.dtype(
-    [
-        ("freq", np.float64),
-        ("assignment_fidelity", np.float64),
-    ]
-)
-
-
 @dataclass
 class TwpaFrequencyData(Data):
     """TwpaFrequency acquisition outputs."""
@@ -48,6 +43,8 @@ class TwpaFrequencyResults(Results):
 
     best_freqs: dict[QubitId, float] = field(default_factory=dict)
     best_fidelities: dict[QubitId, float] = field(default_factory=dict)
+    best_angles: dict[QubitId, float] = field(default_factory=dict)
+    best_thresholds: dict[QubitId, float] = field(default_factory=dict)
 
 
 def _acquisition(
@@ -101,7 +98,7 @@ def _acquisition(
         classification_result = classification._fit(classification_data)
         for qubit in qubits:
             data.register_qubit(
-                TwpaFrequencyType,
+                ResonatorFrequencyType,
                 (qubit),
                 dict(
                     freq=np.array(
@@ -111,9 +108,10 @@ def _acquisition(
                     assignment_fidelity=np.array(
                         [classification_result.assignment_fidelity[qubit]],
                     ),
+                    angle=np.array([classification_result.rotation_angle[qubit]]),
+                    threshold=np.array([classification_result.threshold[qubit]]),
                 ),
             )
-    print(data)
     return data
 
 
@@ -124,12 +122,19 @@ def _fit(data: TwpaFrequencyData) -> TwpaFrequencyResults:
     qubits = data.qubits
     best_freq = {}
     best_fidelity = {}
+    best_angle = {}
+    best_threshold = {}
     for qubit in qubits:
         data_qubit = data[qubit]
         index_best_err = np.argmax(data_qubit["assignment_fidelity"])
         best_fidelity[qubit] = data_qubit["assignment_fidelity"][index_best_err]
         best_freq[qubit] = data_qubit["freq"][index_best_err]
-    return TwpaFrequencyResults(best_freq, best_fidelity)
+        best_angle[qubit] = data_qubit["angle"][index_best_err]
+        best_threshold[qubit] = data_qubit["threshold"][index_best_err]
+
+    return TwpaFrequencyResults(
+        best_freq, best_fidelity, best_thresholds=best_threshold, best_angles=best_angle
+    )
 
 
 def _plot(data: TwpaFrequencyData, fit: TwpaFrequencyResults, qubit):
@@ -169,6 +174,8 @@ def _plot(data: TwpaFrequencyData, fit: TwpaFrequencyResults, qubit):
 
 def _update(results: TwpaFrequencyResults, platform: Platform, qubit: QubitId):
     update.twpa_frequency(results.best_freqs[qubit], platform, qubit)
+    update.iq_angle(results.best_angles[qubit], platform, qubit)
+    update.threshold(results.best_thresholds[qubit], platform, qubit)
 
 
 twpa_frequency = Routine(_acquisition, _fit, _plot, _update)

src/qibocal/protocols/characterization/readout_optimization/twpa_calibration/frequency_power.py

@@ -31,6 +31,8 @@ class TwpaFrequencyPowerParameters(Parameters):
         ("freq", np.float64),
         ("power", np.float64),
         ("assignment_fidelity", np.float64),
+        ("angle", np.float64),
+        ("threshold", np.float64),
     ]
 )
 
@@ -56,6 +58,8 @@ class TwpaFrequencyPowerResults(Results):
     best_freqs: dict[QubitId, float] = field(default_factory=dict)
     best_powers: dict[QubitId, float] = field(default_factory=dict)
     best_fidelities: dict[QubitId, float] = field(default_factory=dict)
+    best_angles: dict[QubitId, float] = field(default_factory=dict)
+    best_thresholds: dict[QubitId, float] = field(default_factory=dict)
 
 
 def _acquisition(
@@ -132,6 +136,8 @@ def _acquisition(
                         assignment_fidelity=np.array(
                             [classification_result.assignment_fidelity[qubit]],
                         ),
+                        angle=np.array([classification_result.rotation_angle[qubit]]),
+                        threshold=np.array([classification_result.threshold[qubit]]),
                     ),
                 )
     return data
@@ -144,6 +150,8 @@ def _fit(data: TwpaFrequencyPowerData) -> TwpaFrequencyPowerResults:
     best_freq = {}
     best_power = {}
     best_fidelity = {}
+    best_angle = {}
+    best_threshold = {}
     qubits = data.qubits
 
     for qubit in qubits:
@@ -152,8 +160,16 @@ def _fit(data: TwpaFrequencyPowerData) -> TwpaFrequencyPowerResults:
         best_fidelity[qubit] = data_qubit["assignment_fidelity"][index_best_err]
         best_freq[qubit] = data_qubit["freq"][index_best_err]
         best_power[qubit] = data_qubit["power"][index_best_err]
-
-    return TwpaFrequencyPowerResults(best_freq, best_power, best_fidelity)
+        best_angle[qubit] = data_qubit["angle"][index_best_err]
+        best_threshold[qubit] = data_qubit["threshold"][index_best_err]
+
+    return TwpaFrequencyPowerResults(
+        best_freq,
+        best_power,
+        best_fidelity,
+        best_angles=best_angle,
+        best_thresholds=best_threshold,
+    )
 
 
 def _plot(data: TwpaFrequencyPowerData, fit: TwpaFrequencyPowerResults, qubit):
@@ -201,6 +217,8 @@ def _plot(data: TwpaFrequencyPowerData, fit: TwpaFrequencyPowerResults, qubit):
 def _update(results: TwpaFrequencyPowerResults, platform: Platform, qubit: QubitId):
     update.twpa_frequency(results.best_freqs[qubit], platform, qubit)
     update.twpa_power(results.best_powers[qubit], platform, qubit)
+    update.iq_angle(results.best_angles[qubit], platform, qubit)
+    update.threshold(results.best_thresholds[qubit], platform, qubit)
 
 
 twpa_frequency_power = Routine(_acquisition, _fit, _plot, _update)

src/qibocal/protocols/characterization/readout_optimization/twpa_calibration/power.py

@@ -27,6 +27,8 @@ class TwpaPowerParameters(Parameters):
     [
         ("power", np.float64),
         ("assignment_fidelity", np.float64),
+        ("angle", np.float64),
+        ("threshold", np.float64),
     ]
 )
 
@@ -45,6 +47,8 @@ class TwpaPowerResults(Results):
 
     best_powers: dict[QubitId, float] = field(default_factory=dict)
     best_fidelities: dict[QubitId, float] = field(default_factory=dict)
+    best_angles: dict[QubitId, float] = field(default_factory=dict)
+    best_thresholds: dict[QubitId, float] = field(default_factory=dict)
 
 
 def _acquisition(
@@ -106,6 +110,8 @@ def _acquisition(
                     assignment_fidelity=np.array(
                         [classification_result.assignment_fidelity[qubit]]
                     ),
+                    angle=np.array([classification_result.rotation_angle[qubit]]),
+                    threshold=np.array([classification_result.threshold[qubit]]),
                 ),
             )
     return data
@@ -117,13 +123,22 @@ def _fit(data: TwpaPowerData) -> TwpaPowerResults:
     qubits = data.qubits
     best_power = {}
     best_fidelity = {}
+    best_angle = {}
+    best_threshold = {}
     for qubit in qubits:
         data_qubit = data[qubit]
         index_best_err = np.argmax(data_qubit["assignment_fidelity"])
         best_fidelity[qubit] = data_qubit["assignment_fidelity"][index_best_err]
         best_power[qubit] = data_qubit["power"][index_best_err]
-
-    return TwpaPowerResults(best_power, best_fidelity)
+        best_angle[qubit] = data_qubit["angle"][index_best_err]
+        best_threshold[qubit] = data_qubit["threshold"][index_best_err]
+
+    return TwpaPowerResults(
+        best_power,
+        best_fidelity,
+        best_angles=best_angle,
+        best_thresholds=best_threshold,
+    )
 
 
 def _plot(data: TwpaPowerData, fit: TwpaPowerResults, qubit):
@@ -162,6 +177,8 @@ def _plot(data: TwpaPowerData, fit: TwpaPowerResults, qubit):
 
 def _update(results: TwpaPowerResults, platform: Platform, qubit: QubitId):
     update.twpa_power(results.best_powers[qubit], platform, qubit)
+    update.iq_angle(results.best_angles[qubit], platform, qubit)
+    update.threshold(results.best_thresholds[qubit], platform, qubit)
 
 
 twpa_power = Routine(_acquisition, _fit, _plot, _update)