add run action on 18a_coupler_zero_point

qualibration_graphs/superconducting/calibration_utils/coupler_zero_point/__init__.py

@@ -0,0 +1,12 @@
+from .parameters import Parameters
+from .analysis import process_raw_dataset, fit_raw_data, log_fitted_results
+from .plotting import  plot_raw_data_with_fit, plot_individual_data_with
+
+__all__ = [
+    "Parameters",
+    "process_raw_dataset",
+    "fit_raw_data",
+    "log_fitted_results",
+    "plot_raw_data_with_fit",
+    "plot_individual_data_with",
+]

qualibration_graphs/superconducting/calibration_utils/coupler_zero_point/analysis.py

@@ -0,0 +1,113 @@
+import logging
+from dataclasses import dataclass
+from typing import Dict, Tuple
+
+import numpy as np
+import xarray as xr
+from qualibrate import QualibrationNode
+
+from quam.components.quantum_components import qubit
+from qualibration_libs.analysis import oscillation
+
+
+
+@dataclass
+class FitParameters:
+    """Stores the relevant qubit spectroscopy experiment fit parameters for a single qubit"""
+
+    success: bool
+    coupler_flux_min: float
+    qubit_flux_max: float
+
+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.
+    """
+    if log_callable is None:
+        log_callable = logging.getLogger(__name__).info
+    pass
+
+
+def process_raw_dataset(ds: xr.Dataset, node: QualibrationNode):
+    """
+    Parameters:
+    -----------
+    ds : xr.Dataset
+        Dataset containing the processed data.
+    node : QualibrationNode
+        The calibration node containing parameters and qubit pairs.
+
+    Returns:
+    --------
+    Tuple[xr.Dataset, Dict[str, FitResults]]
+        Dataset with fit results and dictionary of fit results for each qubit pair.
+    """
+    detuning_mode = "quadratic" # "cosine" or "quadratic"
+    qubit_pairs = [node.machine.qubit_pairs[pair] for pair in node.parameters.qubit_pairs]
+    fluxes_qp = node.namespace["fluxes_qp"]
+    fluxes_coupler = ds.coupler_flux.values
+    qubit_flux_full = np.array([fluxes_qp[qp.name] for qp in qubit_pairs])
+    ds = ds.assign_coords({"qubit_flux_full": (["qubit_pair", "qubit_flux"], qubit_flux_full)})
+
+    coupler_flux_full = np.array([fluxes_coupler + qp.coupler.decouple_offset for qp in qubit_pairs])
+    if detuning_mode == "quadratic":
+        detuning = np.array([-fluxes_qp[qp.name] ** 2 * qp.qubit_control.freq_vs_flux_01_quad_term  for qp in qubit_pairs])
+    elif detuning_mode == "cosine":
+        detuning = np.array([oscillation(fluxes_qp, qp.qubit_control.extras['a'], qp.qubit_control.extras['f'], qp.qubit_control.extras['phi'], qp.qubit_control.extras['offset']) for qp in qubit_pairs])
+    ds = ds.assign_coords({"coupler_flux_full": (["qubit_pair", "coupler_flux"], coupler_flux_full)})
+    ds = ds.assign_coords({"detuning": (["qubit_pair", "qubit_flux"], detuning)})
+
+    return ds
+
+
+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.
+
+    Parameters:
+    -----------
+    ds : xr.Dataset
+        Dataset containing the raw data.
+    node_parameters : Parameters
+        Parameters related to the node, including whether state discrimination is used.
+
+    Returns:
+    --------
+    xr.Dataset
+        Dataset containing the fit results.
+    """
+
+    # Extract the relevant fitted parameters
+    fit_data, fit_results = _extract_relevant_fit_parameters(ds, node)
+    return fit_data, fit_results
+
+
+def _extract_relevant_fit_parameters(fit: xr.Dataset, node: QualibrationNode):
+    """Add metadata to the dataset and fit results."""
+
+    # Populate the FitParameters class with fitted values
+    fit_results = {}
+    for qp in fit.qubit_pair.values:
+        if node.parameters.use_state_discrimination:
+            res_sum = -fit.state_control + fit.state_target
+        else:
+            res_sum = -fit.I_control + fit.I_target
+        fluxes_qp = node.namespace["fluxes_qp"]
+        coupler_min_arg = res_sum.sel(qubit_pair=qp).mean(dim='qubit_flux').argmin()
+        coupler_flux_min = fit.coupler_flux_full.sel(qubit_pair=qp)[coupler_min_arg]
+        qubit_max_arg = res_sum.sel(qubit_pair=qp).mean(dim="coupler_flux").argmax()
+        qubit_flux_max = fluxes_qp[qp][qubit_max_arg]
+        #fit_results[qp] = {"flux_coupler_min": float(flux_coupler_min.values), "flux_qubit_max": float(flux_qubit_max)}
+        fit_results[qp] = FitParameters(
+            success=True,
+            coupler_flux_min=float(coupler_flux_min.values),
+            qubit_flux_max=float(qubit_flux_max)
+        )
+    return fit, fit_results

qualibration_graphs/superconducting/calibration_utils/coupler_zero_point/parameters.py

@@ -0,0 +1,27 @@
+from typing import Literal, Optional, List
+
+from qualibrate import NodeParameters
+from qualibrate.parameters import RunnableParameters
+from qualibration_libs.parameters import QubitsExperimentNodeParameters, CommonNodeParameters
+
+
+class NodeSpecificParameters(RunnableParameters):
+    qubit_pairs: Optional[List[str]] = ["qA1-qA2"]
+    num_averages: int = 500
+    coupler_flux_min : float = -0.05 #relative to the coupler set point
+    coupler_flux_max : float = 0.05 #relative to the coupler set point
+    coupler_flux_step : float = 0.0004
+    qubit_flux_span : float = 0.026 # relative to the known/calculated detuning between the qubits
+    qubit_flux_step : float = 0.0002
+    pulse_duration_ns: int = 232
+    cz_or_iswap: Literal["cz", "iswap"] = "cz"
+    use_saved_detuning: bool = False
+    flux_point_joint_or_independent_or_pairwise: Literal["joint", "independent", "pairwise"] = "joint"
+
+class Parameters(
+    NodeParameters,
+    CommonNodeParameters,
+    NodeSpecificParameters,
+    QubitsExperimentNodeParameters,
+):
+    pass

qualibration_graphs/superconducting/calibration_utils/coupler_zero_point/plotting.py

@@ -0,0 +1,180 @@
+from typing import List
+import xarray as xr
+from matplotlib.axes import Axes
+import matplotlib.pyplot as plt
+import numpy as np
+
+from qualang_tools.units import unit
+from quam_builder.architecture.superconducting.qubit import AnyTransmon
+u = unit(coerce_to_integer=True)
+
+def plot_raw_data_with_fit(ds: xr.Dataset, qubit_pairs: List[AnyTransmon], fits: xr.Dataset):
+    """
+    Plots the resonator spectroscopy amplitude IQ_abs with fitted curves for the given qubits.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset containing the quadrature data.
+    qubits : list of AnyTransmon
+        A list of qubits to plot.
+    fits : xr.Dataset
+        The dataset containing the fit parameters.
+
+    Returns
+    -------
+    Figure
+        The matplotlib figure object containing the plots.
+
+    Notes
+    -----
+    - The function creates a grid of subplots, one for each qubit.
+    - Each subplot contains the raw data and the fitted curve.
+    """
+    fig, axs = plt.subplots(nrows=len(qubit_pairs), ncols=1, figsize=(15, 9))
+    for ii, qp in enumerate(qubit_pairs):
+        ax = axs[ii] if len(qubit_pairs) > 1 else axs
+
+        # Try to get fit data for this qubit pair, handle if missing
+        try:
+            fit_data = fits[qp.id] if fits is not None else None
+        except (KeyError, ValueError):
+            # If this qubit pair is not in the fit results, set fit_data to None
+            fit_data = None
+
+        plot_individual_data_with(ax, ds, qp.id, fit_data)
+
+    fig.suptitle("coupler zero point")
+    fig.set_size_inches(15, 9)
+    fig.tight_layout()
+
+    return fig
+
+
+
+
+def plot_individual_data_with(ax: Axes, ds: xr.Dataset, qubit_pair: str, fit: xr.Dataset = None):
+    """
+    Plots individual qubit data on a given axis with optional fit.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        The axis on which to plot the data.
+    ds : xr.Dataset
+        The dataset containing the quadrature data.
+    qubit : dict[str, str]
+        mapping to the qubit to plot.
+    fit : xr.Dataset, optional
+        The dataset containing the fit parameters (default is None).
+
+    Notes
+    -----
+    - If the fit dataset is provided, the fitted curve is plotted along with the raw data.
+    """
+
+    if hasattr(ds, "state_target"):
+        # If the dataset has 'state_target', use it for plotting
+        data = ds.state_target.sel(qubit_pair=qubit_pair)
+    else:
+        data = ds.I_target.sel(qubit_pair=qubit_pair)
+
+    data.assign_coords({"qubit_flux_mV": 1e3*data.qubit_flux_full, "coupler_flux_mV": 1e3*data.coupler_flux_full}).plot(x='qubit_flux_mV',y='coupler_flux_mV')
+
+    # Only plot fit results if they exist and are valid
+    if fit is not None:
+        try:
+            # Check if fit values exist and are valid (not NaN)
+            qubit_flux_max = float(fit['qubit_flux_max'])
+            coupler_flux_min = float(fit['coupler_flux_min'])
+            if not (np.isnan(qubit_flux_max) or np.isnan(coupler_flux_min)):
+                ax.axhline(1e3*coupler_flux_min, color = 'red', lw = 0.5, ls = '--')
+                #ax.axhline(1e3*machine.qubit_pairs[qp['qubit']].coupler.decouple_offset, color = 'blue', lw =0.5, ls = '--')
+                ax.axvline(1e3*qubit_flux_max, color = 'red', lw =0.5, ls = '--')
+            else:
+                ax.set_title(f"{qubit_pair} - Fit Failed (Invalid Parameters)")
+        except (ValueError, TypeError, AttributeError):
+            ax.set_title(f"{qubit_pair} - Fit Failed")
+    else:
+        ax.set_title(f"{qubit_pair} - No Fit Data")
+    detuning_data = ds.sel(qubit_pair=qubit_pair).detuning.values * 1e-6
+    def flux_to_detuning(x):
+        return np.interp(x, 1e3*data.qubit_flux_full, detuning_data)
+
+    def detuning_to_flux(y):
+        return np.interp(y, detuning_data, 1e3*data.qubit_flux_full)
+
+    sec_ax = ax.secondary_xaxis('top', functions=(flux_to_detuning, detuning_to_flux))
+    sec_ax.set_xlabel('Detuning [MHz]')
+
+    ax.set_xlabel("qubit flux [mV]")
+    ax.set_ylabel("coupler flux [mV]")
+
+"""
+# %% {Plotting}
+if not node.parameters.simulate:
+    grid_names, qubit_pair_names = grid_pair_names(qubit_pairs)
+    grid = QubitPairGrid(grid_names, qubit_pair_names)
+    for ax, qp in grid_iter(grid):
+        if node.parameters.use_state_discrimination:
+            values_to_plot = ds.state_control.sel(qubit=qp['qubit'])
+        else:
+            values_to_plot = ds.I_control.sel(qubit=qp['qubit'])
+
+        values_to_plot.assign_coords({"flux_qubit_mV": 1e3*values_to_plot.flux_qubit_full, "flux_coupler_mV": 1e3*values_to_plot.flux_coupler_full}).plot(ax = ax, cmap = 'viridis', x = 'flux_qubit_mV', y = 'flux_coupler_mV')
+        qubit_pair = machine.qubit_pairs[qp['qubit']]
+        ax.set_title(f"{qp['qubit']}, idle: {qubit_pair.coupler.decouple_offset}V, g=0: {node.results['results'][qp['qubit']]['flux_coupler_min']:.4f}V", fontsize = 10)
+        ax.axhline(1e3*node.results["results"][qp["qubit"]]["flux_coupler_min"], color = 'red', lw = 0.5, ls = '--')
+        ax.axhline(1e3*machine.qubit_pairs[qp['qubit']].coupler.decouple_offset, color = 'blue', lw =0.5, ls = '--')
+        ax.axvline(1e3*node.results["results"][qp["qubit"]]["flux_qubit_max"], color = 'red', lw =0.5, ls = '--')
+        # Create a secondary x-axis for detuning
+        flux_qubit_data = ds.sel(qubit=qp['qubit']).flux_qubit_full.values*1e3
+        detuning_data = ds.sel(qubit=qp['qubit']).detuning.values * 1e-6
+
+        def flux_to_detuning(x):
+            return np.interp(x, flux_qubit_data, detuning_data)
+
+        def detuning_to_flux(y):
+            return np.interp(y, detuning_data, flux_qubit_data)
+
+        sec_ax = ax.secondary_xaxis('top', functions=(flux_to_detuning, detuning_to_flux))
+        sec_ax.set_xlabel('Detuning [MHz]')
+        ax.set_xlabel('Qubit flux pulse [mV]')
+        ax.set_ylabel('Coupler flux pulse [mV]')
+    grid.fig.suptitle('Control')
+    plt.tight_layout()
+    plt.show()
+    node.results['figure_control'] = grid.fig
+
+    grid = QubitPairGrid(grid_names, qubit_pair_names)
+    for ax, qp in grid_iter(grid):
+        if node.parameters.use_state_discrimination:
+            values_to_plot = ds.state_target.sel(qubit=qp['qubit'])
+        else:
+            values_to_plot = ds.I_target.sel(qubit=qp['qubit'])
+
+        values_to_plot.assign_coords({"flux_qubit_mV": 1e3*values_to_plot.flux_qubit_full, "flux_coupler_mV": 1e3*values_to_plot.flux_coupler_full}).plot(ax = ax, cmap = 'viridis', x = 'flux_qubit_mV', y = 'flux_coupler_mV')
+        qubit_pair = machine.qubit_pairs[qp['qubit']]
+        ax.set_title(f"{qp['qubit']}, idle: {qubit_pair.coupler.decouple_offset}V, g=0: {node.results['results'][qp['qubit']]['flux_coupler_min']:.4f}V", fontsize = 10)
+        ax.axhline(1e3*node.results["results"][qp["qubit"]]["flux_coupler_min"], color = 'red', lw = 0.5, ls = '--')
+        ax.axhline(1e3*machine.qubit_pairs[qp['qubit']].coupler.decouple_offset, color = 'blue', lw =0.5, ls = '--')
+        ax.axvline(1e3*node.results["results"][qp["qubit"]]["flux_qubit_max"], color = 'red', lw =0.5, ls = '--')
+        # Create a secondary x-axis for detuning
+        flux_qubit_data = ds.sel(qubit=qp['qubit']).flux_qubit_full.values*1e3
+        detuning_data = ds.sel(qubit=qp['qubit']).detuning.values * 1e-6
+
+        def flux_to_detuning(x):
+            return np.interp(x, flux_qubit_data, detuning_data)
+
+        def detuning_to_flux(y):
+            return np.interp(y, detuning_data, flux_qubit_data)
+
+        sec_ax = ax.secondary_xaxis('top', functions=(flux_to_detuning, detuning_to_flux))
+        sec_ax.set_xlabel('Detuning [MHz]')
+        ax.set_xlabel('Qubit flux shift [mV]')
+        ax.set_ylabel('Coupler flux [mV]')
+    grid.fig.suptitle('Target')
+    plt.tight_layout()
+    plt.show()
+    node.results['figure_target'] = grid.fig
+"""