Chore: add iq blobs with precommit compliant

qualibration_graphs/superconducting/calibration_utils/iq_blobs/__init__.py

@@ -1,3 +1,5 @@
+"""IQ blob calibration utilities."""
+
 from .parameters import Parameters
 from .analysis import process_raw_dataset, fit_raw_data, log_fitted_results
 from .plotting import plot_iq_blobs, plot_confusion_matrices

qualibration_graphs/superconducting/calibration_utils/iq_blobs/analysis.py

@@ -1,6 +1,9 @@
+"""Analysis helpers for IQ blob readout calibration."""
+
 import logging
 from dataclasses import dataclass
-from typing import Tuple, Dict
+from typing import Dict, Tuple
+
 import numpy as np
 import xarray as xr
 
@@ -11,7 +14,7 @@
 
 @dataclass
 class FitParameters:
-    """Stores the relevant qubit spectroscopy experiment fit parameters for a single qubit"""
+    """Stores fitted IQ blob discrimination parameters for a single qubit."""
 
     iw_angle: float
     ge_threshold: float
@@ -22,17 +25,7 @@ class FitParameters:
 
 
 def log_fitted_results(fit_results: Dict, log_callable=None):
-    """
-    Logs the node-specific fitted results for all qubits from the fit xarray Dataset.
-
-    Parameters:
-    -----------
-    ds : xr.Dataset
-        Dataset containing the fitted results for all qubits.
-    log_callable : callable, optional
-        Callable for logging the fitted results. If None, a default logger is used.
-
-    """
+    """Log fitted IQ blob thresholds and readout fidelity per qubit."""
     if log_callable is None:
         log_callable = logging.getLogger(__name__).info
     for q in fit_results.keys():
@@ -49,6 +42,8 @@ def log_fitted_results(fit_results: Dict, log_callable=None):
 
 
 def process_raw_dataset(ds: xr.Dataset, node: QualibrationNode):
+    """Normalize raw IQ data to volts and remove tuple nesting."""
+
     # Fix the structure of ds to avoid tuples
     def extract_value(element):
         if isinstance(element, tuple):
@@ -68,27 +63,29 @@ def extract_value(element):
 
 def fit_raw_data(ds: xr.Dataset, node: QualibrationNode) -> Tuple[xr.Dataset, dict[str, FitParameters]]:
     """
-    Fit the qubit frequency and FWHM for each qubit in the dataset.
+    Compute rotation angle, discrimination thresholds, and readout fidelity for each qubit.
 
     Parameters:
     -----------
     ds : xr.Dataset
         Dataset containing the raw data.
-    node_parameters : Parameters
-        Parameters related to the node, including whether state discrimination is used.
+    node : QualibrationNode
+        Node providing the active qubits and metadata.
 
     Returns:
     --------
     xr.Dataset
-        Dataset containing the fit results.
+        Dataset containing the fitted parameters.
+    dict[str, FitParameters]
+        Per-qubit fit parameters for logging and state updates.
     """
     ds_fit = ds
     # Condition to have the Q equal for both states:
     angle = np.arctan2(
         ds_fit.Qe.mean(dim="n_runs") - ds_fit.Qg.mean(dim="n_runs"),
         ds_fit.Ig.mean(dim="n_runs") - ds_fit.Ie.mean(dim="n_runs"),
     )
-    ds_fit = ds_fit.assign({"iw_angle": xr.DataArray(angle, coords=dict(qubit=ds_fit.qubit.data))})
+    ds_fit = ds_fit.assign({"iw_angle": xr.DataArray(angle, coords={"qubit": ds_fit.qubit.data})})
 
     C = np.cos(angle)
     S = np.sin(angle)
@@ -110,7 +107,7 @@ def fit_raw_data(ds: xr.Dataset, node: QualibrationNode) -> Tuple[xr.Dataset, di
         hist[ii][1][1:][np.argmax(np.histogram(ds_fit.Ig_rot.sel(qubit=q.name), bins=100)[0])]
         for ii, q in enumerate(node.namespace["qubits"])
     ]
-    ds_fit = ds_fit.assign({"rus_threshold": xr.DataArray(rus_threshold, coords=dict(qubit=ds_fit.qubit.data))})
+    ds_fit = ds_fit.assign({"rus_threshold": xr.DataArray(rus_threshold, coords={"qubit": ds_fit.qubit.data})})
 
     threshold = []
     gg, ge, eg, ee = [], [], [], []
@@ -126,13 +123,13 @@ def fit_raw_data(ds: xr.Dataset, node: QualibrationNode) -> Tuple[xr.Dataset, di
         ge.append(np.sum(ds_fit.Ig_rot.sel(qubit=q.name) > fit.x[0]) / len(ds_fit.Ig_rot.sel(qubit=q.name)))
         eg.append(np.sum(ds_fit.Ie_rot.sel(qubit=q.name) < fit.x[0]) / len(ds_fit.Ie_rot.sel(qubit=q.name)))
         ee.append(np.sum(ds_fit.Ie_rot.sel(qubit=q.name) > fit.x[0]) / len(ds_fit.Ie_rot.sel(qubit=q.name)))
-    ds_fit = ds_fit.assign({"ge_threshold": xr.DataArray(threshold, coords=dict(qubit=ds_fit.qubit.data))})
-    ds_fit = ds_fit.assign({"gg": xr.DataArray(gg, coords=dict(qubit=ds_fit.qubit.data))})
-    ds_fit = ds_fit.assign({"ge": xr.DataArray(ge, coords=dict(qubit=ds_fit.qubit.data))})
-    ds_fit = ds_fit.assign({"eg": xr.DataArray(eg, coords=dict(qubit=ds_fit.qubit.data))})
-    ds_fit = ds_fit.assign({"ee": xr.DataArray(ee, coords=dict(qubit=ds_fit.qubit.data))})
+    ds_fit = ds_fit.assign({"ge_threshold": xr.DataArray(threshold, coords={"qubit": ds_fit.qubit.data})})
+    ds_fit = ds_fit.assign({"gg": xr.DataArray(gg, coords={"qubit": ds_fit.qubit.data})})
+    ds_fit = ds_fit.assign({"ge": xr.DataArray(ge, coords={"qubit": ds_fit.qubit.data})})
+    ds_fit = ds_fit.assign({"eg": xr.DataArray(eg, coords={"qubit": ds_fit.qubit.data})})
+    ds_fit = ds_fit.assign({"ee": xr.DataArray(ee, coords={"qubit": ds_fit.qubit.data})})
     ds_fit = ds_fit.assign(
-        {"readout_fidelity": xr.DataArray(100 * (ds_fit.gg + ds_fit.ee) / 2, coords=dict(qubit=ds_fit.qubit.data))}
+        {"readout_fidelity": xr.DataArray(100 * (ds_fit.gg + ds_fit.ee) / 2, coords={"qubit": ds_fit.qubit.data})}
     )
 
     # Extract the relevant fitted parameters
@@ -171,7 +168,7 @@ def _extract_relevant_fit_parameters(fit: xr.Dataset, node: QualibrationNode):
                 [float(fit.sel(qubit=q).gg), float(fit.sel(qubit=q).ge)],
                 [float(fit.sel(qubit=q).eg), float(fit.sel(qubit=q).ee)],
             ],
-            success=fit.sel(qubit=q).success.values.__bool__(),
+            success=bool(fit.sel(qubit=q).success.values),
         )
         for q in fit.qubit.values
     }

qualibration_graphs/superconducting/calibration_utils/iq_blobs/parameters.py

@@ -1,10 +1,14 @@
+"""Parameter definitions for the IQ blob readout calibration."""
+
 from typing import Literal
 from qualibrate import NodeParameters
 from qualibrate.parameters import RunnableParameters
 from qualibration_libs.parameters import QubitsExperimentNodeParameters, CommonNodeParameters
 
 
 class NodeSpecificParameters(RunnableParameters):
+    """Run-time settings for IQ blob acquisition."""
+
     num_shots: int = 2000
     """Number of runs to perform. Default is 2000."""
     operation: Literal["readout", "readout_QND"] = "readout"
@@ -17,4 +21,6 @@ class Parameters(
     NodeSpecificParameters,
     QubitsExperimentNodeParameters,
 ):
+    """Combined parameters for the IQ blob calibration node."""
+
     pass

qualibration_graphs/superconducting/calibration_utils/iq_blobs/plotting.py

@@ -1,6 +1,9 @@
+"""Plotting helpers for IQ blob readout analysis."""
+
 from typing import List
-import xarray as xr
+
 import numpy as np
+import xarray as xr
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 
@@ -13,7 +16,7 @@
 
 def plot_iq_blobs(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
     """
-    Plots the IQ blobs with the derived thresholds for the given qubits.
+    Plot IQ blobs and discrimination thresholds for the given qubits.
 
     Parameters
     ----------
@@ -32,7 +35,7 @@ def plot_iq_blobs(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
     Notes
     -----
     - The function creates a grid of subplots, one for each qubit.
-    - Each subplot contains the raw data and the fitted curve.
+    - Each subplot shows rotated IQ points and threshold lines.
     """
     grid = QubitGrid(ds, [q.grid_location for q in qubits])
     for ax, qubit in grid_iter(grid):
@@ -50,7 +53,7 @@ def plot_iq_blobs(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
 
 def plot_individual_iq_blobs(ax: Axes, ds: xr.Dataset, qubit: dict[str, str], fit: xr.Dataset = None):
     """
-    Plots individual qubit data on a given axis with optional fit.
+    Plot a single qubit's rotated IQ points and thresholds.
 
     Parameters
     ----------
@@ -65,7 +68,7 @@ def plot_individual_iq_blobs(ax: Axes, ds: xr.Dataset, qubit: dict[str, str], fi
 
     Notes
     -----
-    - If the fit dataset is provided, the fitted curve is plotted along with the raw data.
+    - If the fit dataset is provided, thresholds are drawn on top of the points.
     """
 
     ax.plot(1e3 * fit.Ig_rot, 1e3 * fit.Qg_rot, ".", alpha=0.2, label="Ground", markersize=1)
@@ -93,7 +96,7 @@ def plot_individual_iq_blobs(ax: Axes, ds: xr.Dataset, qubit: dict[str, str], fi
 
 def plot_historams(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
     """
-    Plots the IQ blobs with the derived thresholds for the given qubits.
+    Plot rotated-I histograms and thresholds for the given qubits.
 
     Parameters
     ----------
@@ -112,7 +115,7 @@ def plot_historams(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
     Notes
     -----
     - The function creates a grid of subplots, one for each qubit.
-    - Each subplot contains the raw data and the fitted curve.
+    - Each subplot shows ground/excited histograms and threshold lines.
     """
     grid = QubitGrid(ds, [q.grid_location for q in qubits])
     for ax, qubit in grid_iter(grid):
@@ -128,7 +131,7 @@ def plot_historams(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
 
 def plot_individual_histograms(ax: Axes, ds: xr.Dataset, qubit: dict[str, str], fit: xr.Dataset = None):
     """
-    Plots individual qubit data on a given axis with optional fit.
+    Plot rotated-I histograms and thresholds for a single qubit.
 
     Parameters
     ----------
@@ -143,7 +146,7 @@ def plot_individual_histograms(ax: Axes, ds: xr.Dataset, qubit: dict[str, str],
 
     Notes
     -----
-    - If the fit dataset is provided, the fitted curve is plotted along with the raw data.
+    - If the fit dataset is provided, thresholds are drawn on the histograms.
     """
 
     ax.hist(1e3 * fit.Ig_rot, bins=100, alpha=0.5, label="Ground")
@@ -163,7 +166,7 @@ def plot_individual_histograms(ax: Axes, ds: xr.Dataset, qubit: dict[str, str],
 
 def plot_confusion_matrices(ds: xr.Dataset, qubits: List[AnyTransmon], fits: xr.Dataset):
     """
-    Plots the confusion matrix for the given qubits.
+    Plot readout confusion matrices for the given qubits.
 
     Parameters
     ----------

qualibration_graphs/superconducting/calibrations/1Q_calibrations/07_iq_blobs.py

@@ -1,9 +1,11 @@
+"""IQ blob readout calibration node for setting discrimination thresholds."""
+
 # %% {Imports}
+from dataclasses import asdict
+
 import matplotlib.pyplot as plt
 import numpy as np
-from calibration_utils.iq_blobs.plotting import plot_historams
 import xarray as xr
-from dataclasses import asdict
 
 from qm.qua import *
 
@@ -12,18 +14,20 @@
 from qualang_tools.units import unit
 
 from qualibrate import QualibrationNode
+from qualibration_libs.data import XarrayDataFetcher
+from qualibration_libs.parameters import get_qubits
+from qualibration_libs.runtime import simulate_and_plot
 from quam_config import Quam
+
 from calibration_utils.iq_blobs import (
     Parameters,
-    process_raw_dataset,
     fit_raw_data,
     log_fitted_results,
-    plot_iq_blobs,
     plot_confusion_matrices,
+    plot_iq_blobs,
+    process_raw_dataset,
 )
-from qualibration_libs.parameters import get_qubits
-from qualibration_libs.runtime import simulate_and_plot
-from qualibration_libs.data import XarrayDataFetcher
+from calibration_utils.iq_blobs.plotting import plot_historams
 
 
 # %% {Description}
@@ -267,4 +271,5 @@ def update_state(node: QualibrationNode[Parameters, Quam]):
 # %% {Save_results}
 @node.run_action()
 def save_results(node: QualibrationNode[Parameters, Quam]):
+    """Persist the node results and state updates."""
     node.save()