Chore/styx fix typos

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

@@ -3,6 +3,7 @@
 # %% {Imports}
 import warnings
 from dataclasses import asdict
+from typing import List, Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -15,19 +16,49 @@
 from qualang_tools.results import progress_counter
 from qualang_tools.units import unit
 
-from qualibrate import QualibrationNode
+from qualibrate import NodeParameters, QualibrationNode
 from quam_config import Quam
 from calibration_utils.resonator_spectroscopy_vs_flux import (
-    Parameters,
     process_raw_dataset,
     fit_raw_data,
     log_fitted_results,
     plot_raw_data_with_fit,
 )
-from qualibration_libs.parameters import get_qubit_pairs
+from qualibration_libs.parameters import CommonNodeParameters
 from qualibration_libs.runtime import simulate_and_plot
 from qualibration_libs.data import XarrayDataFetcher
 
+
+class Parameters(
+    NodeParameters,
+    CommonNodeParameters,
+):
+    """Parameters for resonator spectroscopy versus coupler flux calibration."""
+
+    qubits: Optional[List[str]] = None
+    """List of qubit names to calibrate. If None or empty, all active qubits are used."""
+    coupler: Optional[str] = None
+    """Name of the coupler to use for flux pulsing (e.g., 'coupler_q1_q2')."""
+    num_shots: int = 100
+    """Number of averages to perform. Default is 100."""
+    min_flux_offset_in_v: float = -0.5
+    """Minimum flux bias offset in volts. Default is -0.5 V."""
+    max_flux_offset_in_v: float = 0.5
+    """Maximum flux bias offset in volts. Default is 0.5 V."""
+    num_flux_points: int = 101
+    """Number of flux points. Default is 101."""
+    frequency_span_in_mhz: float = 15
+    """Frequency span in MHz. Default is 15 MHz."""
+    frequency_step_in_mhz: float = 0.1
+    """Frequency step in MHz. Default is 0.1 MHz."""
+    input_line_impedance_in_ohm: float = 50
+    """Input line impedance in ohms. Default is 50 Ohm."""
+    line_attenuation_in_db: float = 0
+    """Line attenuation in dB. Default is 0 dB."""
+    update_flux_min: bool = False
+    """Flag to update flux minimum frequency point. Default is False."""
+
+
 # %% {Node initialisation}
 description = """
         RESONATOR SPECTROSCOPY VERSUS COUPLER FLUX
@@ -62,7 +93,8 @@
 def custom_param(node: QualibrationNode[Parameters, Quam]):
     """Allow the user to locally set the node parameters for debugging purposes, or execution in the Python IDE."""
     # You can get type hinting in your IDE by typing node.parameters.
-    # node.parameters.qubit_pairs = ["q1-q2"]
+    # node.parameters.qubits = ["q1", "q2"]
+    # node.parameters.coupler = "coupler_q1_q2"
     pass
 
 
@@ -76,28 +108,25 @@ def create_qua_program(node: QualibrationNode[Parameters, Quam]):  # pylint: dis
     """Create the sweep axes and generate the QUA program from the pulse sequence and the node parameters."""
     # Class containing tools to help handle units and conversions.
     u = unit(coerce_to_integer=True)
-    # Get the active qubit pairs from the node and organize them by batches
-    node.namespace["qubit_pairs"] = qubit_pairs = get_qubit_pairs(node)
-    num_qubit_pairs = len(qubit_pairs)
-
-    # Extract all qubits from qubit pairs for measurement
-    # Create a list of all pairs with their global indices for QUA variable indexing
-    measured_qubits = []
-    all_pairs_list = []  # List of all pairs in order
-    qubit_idx = 0
-    for qp in qubit_pairs:
-        measured_qubits.append(qp.qubit_control)
-        measured_qubits.append(qp.qubit_target)
-        all_pairs_list.append(qp)
-    node.namespace["measured_qubits"] = measured_qubits
-    node.namespace["all_pairs_list"] = all_pairs_list
-    # Also set in qubits for compatibility with utility functions
-    node.namespace["qubits"] = measured_qubits
-    num_qubits = len(measured_qubits)
-
-    # Check if the qubit pairs have couplers attached
-    if any(qp.coupler is None for qp in qubit_pairs):
-        warnings.warn("Found qubit pairs without a coupler. Skipping")
+    machine = node.machine
+
+    # Get the qubits from parameters or use active qubits
+    if not node.parameters.qubits:
+        qubits = machine.active_qubits
+    else:
+        qubits = [machine.qubits[q] for q in node.parameters.qubits]
+
+    # Get the coupler from parameters
+    assert node.parameters.coupler is not None, "Coupler parameter must be specified."
+    coupler = machine.qubit_pairs[node.parameters.coupler].coupler
+
+    node.namespace["qubits"] = qubits
+    node.namespace["coupler"] = coupler
+    num_qubits = len(qubits)
+
+    # Check if the coupler is attached
+    if coupler is None:
+        warnings.warn("Coupler not found. Skipping")
     # Extract the sweep parameters and axes from the node parameters
     n_avg = node.parameters.num_shots
     # Coupler flux bias sweep in V
@@ -113,7 +142,7 @@ def create_qua_program(node: QualibrationNode[Parameters, Quam]):  # pylint: dis
 
     # Register the sweep axes to be added to the dataset when fetching data
     node.namespace["sweep_axes"] = {
-        "qubit": xr.DataArray([q.name for q in measured_qubits]),
+        "qubit": xr.DataArray([q.name for q in qubits]),
         "flux_bias": xr.DataArray(dcs, attrs={"long_name": "coupler flux bias", "units": "V"}),
         "detuning": xr.DataArray(dfs, attrs={"long_name": "readout frequency", "units": "Hz"}),
     }
@@ -124,54 +153,32 @@ def create_qua_program(node: QualibrationNode[Parameters, Quam]):  # pylint: dis
         dc = declare(fixed)  # QUA variable for the coupler flux bias
         df = declare(int)  # QUA variable for the readout frequency detuning
 
-        # Initialize all qubits from all pairs
-        for qp in all_pairs_list:
-            node.machine.initialize_qpu(target=qp.qubit_control)
-            node.machine.initialize_qpu(target=qp.qubit_target)
+        # Initialize all qubits
+        for qubit in qubits:
+            node.machine.initialize_qpu(target=qubit)
         align()
 
         with for_(n, 0, n < n_avg, n + 1):
             save(n, n_st)
             with for_(*from_array(dc, dcs)):
-                # Apply coupler flux bias for all pairs via DC offset
-                for qp in all_pairs_list:
-                    # Coupler flux sweeping by setting the OPX DC offset associated with the coupler element
-                    qp.coupler.set_dc_offset(dc)
-                # Wait for couplers to settle
+                # Apply coupler flux bias via DC offset
+                coupler.set_dc_offset(dc)
+                # Wait for coupler to settle
                 wait(100, "resonator*")
-                # Measure all qubits from all pairs in order
-                # Use compile-time constant indices: 0, 1 for first pair, 2, 3 for second, etc.
-                qubit_idx = 0
-                for qp in all_pairs_list:
-                    control_idx = qubit_idx
-                    target_idx = qubit_idx + 1
-                    # Measure control qubit
-                    rr_control = qp.qubit_control.resonator
-                    qp.qubit_control.align()
-                    with for_(*from_array(df, dfs)):
-                        # Update the resonator frequencies for resonator
-                        rr_control.update_frequency(df + rr_control.intermediate_frequency)
-                        # readout the resonator
-                        rr_control.measure("readout", qua_vars=(I[control_idx], Q[control_idx]))
-                        # wait for the resonator to deplete
-                        rr_control.wait(rr_control.depletion_time * u.ns)
-                        # save data
-                        save(I[control_idx], I_st[control_idx])
-                        save(Q[control_idx], Q_st[control_idx])
-                    # Measure target qubit
-                    rr_target = qp.qubit_target.resonator
-                    qp.qubit_target.align()
+                # Measure all qubits
+                for i, qubit in enumerate(qubits):
+                    rr = qubit.resonator
+                    qubit.align()
                     with for_(*from_array(df, dfs)):
-                        # Update the resonator frequencies for resonator
-                        rr_target.update_frequency(df + rr_target.intermediate_frequency)
-                        # readout the resonator
-                        rr_target.measure("readout", qua_vars=(I[target_idx], Q[target_idx]))
-                        # wait for the resonator to deplete
-                        rr_target.wait(rr_target.depletion_time * u.ns)
-                        # save data
-                        save(I[target_idx], I_st[target_idx])
-                        save(Q[target_idx], Q_st[target_idx])
-                    qubit_idx += 2  # Python variable, incremented at Python time during code generation
+                        # Update the resonator frequency
+                        rr.update_frequency(df + rr.intermediate_frequency)
+                        # Readout the resonator
+                        rr.measure("readout", qua_vars=(I[i], Q[i]))
+                        # Wait for the resonator to deplete
+                        rr.wait(rr.depletion_time * u.ns)
+                        # Save data
+                        save(I[i], I_st[i])
+                        save(Q[i], Q_st[i])
 
         with stream_processing():
             n_st.save("n")
@@ -229,15 +236,21 @@ def load_data(node: QualibrationNode[Parameters, Quam]):
     # Load the specified dataset
     node.load_from_id(node.parameters.load_data_id)
     node.parameters.load_data_id = load_data_id
-    # Get the active qubit pairs from the loaded node parameters
-    node.namespace["qubit_pairs"] = get_qubit_pairs(node)
-    # Extract measured qubits
-    measured_qubits = []
-    for qp in node.namespace["qubit_pairs"]:
-        measured_qubits.append(qp.qubit_control)
-        measured_qubits.append(qp.qubit_target)
-    node.namespace["measured_qubits"] = measured_qubits
-    node.namespace["qubits"] = measured_qubits
+
+    machine = node.machine
+
+    # Get the qubits from parameters or use active qubits
+    if not node.parameters.qubits:
+        qubits = machine.active_qubits
+    else:
+        qubits = [machine.qubits[q] for q in node.parameters.qubits]
+
+    # Get the coupler from parameters
+    assert node.parameters.coupler is not None, "Coupler parameter must be specified."
+    coupler = machine.qubit_pairs[node.parameters.coupler].coupler
+
+    node.namespace["qubits"] = qubits
+    node.namespace["coupler"] = coupler
 
 
 # %% {Analyse_data}