Feature: modify resonator spec wide

qualibration_graphs/superconducting/calibrations/1Q_calibrations/02_resonator_spectroscopy_wide.py

@@ -26,21 +26,22 @@
     plot_raw_amplitude_with_fit,
     plot_raw_phase,
 )
+from qualibration_libs.analysis import peaks_dips
+from qualibration_libs.data import add_amplitude_and_phase, convert_IQ_to_V
 from qualibration_libs.parameters import get_qubits
+from qualibration_libs.plotting import QubitGrid, grid_iter
 from qualibration_libs.runtime import simulate_and_plot
 from qualibration_libs.data import XarrayDataFetcher
 
 # %% {Node initialisation}
 description = """
-        1D RESONATOR SPECTROSCOPY
-This sequence involves measuring the resonator by sending a readout pulse and demodulating the signals to extract the
-'I' and 'Q' quadratures across varying readout intermediate frequencies for all the active qubits.
-The data is then post-processed to determine the resonator resonance frequency.
-This frequency is used to update the readout frequency in the state.
+        READOUT LINE WIDE RESONATOR SPECTROSCOPY
+This sequence performs a wide frequency sweep per readout line and measures the line response using a probe resonator.
+Each readout line can host multiple resonators. The wide scan identifies multiple resonances and maps them back to the
+qubits on that line, updating each resonator frequency in the state.
 
 Prerequisites:
-    - Having calibrated the IQ mixer/Octave connected to the readout line (node 01a_mixer_calibration.py).
-    - Having calibrated the time of flight, offsets, and gains (node 01a_time_of_flight.py).
+    - Having calibrated the time of flight, offsets, and gains (node 01_time_of_flight_mw_fem.py).
     - Having initialized the QUAM state parameters for the readout pulse amplitude and duration, and the resonators depletion time.
     - Having specified the desired flux point if relevant (qubit.z.flux_point).
 
@@ -50,7 +51,7 @@
 
 # Be sure to include [Parameters, Quam] so the node has proper type hinting
 node = QualibrationNode[Parameters, Quam](
-    name="02a_resonator_spectroscopy",  # Name should be unique
+    name="02_resonator_spectroscopy_wide",  # Name should be unique
     description=description,  # Describe what the node is doing, which is also reflected in the QUAlibrate GUI
     parameters=Parameters(),  # Node parameters defined under quam_experiment/experiments/node_name
 )
@@ -70,15 +71,60 @@ def custom_param(node: QualibrationNode[Parameters, Quam]):
 node.machine = Quam.load()
 
 
+# %% {Helpers}
+def _get_upconverter_frequency(resonator):
+    """Resolve the readout line upconverter frequency from the resonator object."""
+    for attr_name in ("upconverter_frequency", "LO_frequency"):
+        value = getattr(resonator, attr_name, None)
+        if value is not None:
+            return value
+    opx_output = getattr(resonator, "opx_output", None)
+    value = getattr(opx_output, "upconverter_frequency", None)
+    if value is not None:
+        return value
+    rf_freq = getattr(resonator, "RF_frequency", None)
+    if rf_freq is None:
+        return None
+    inter_freq = getattr(resonator, "intermediate_frequency", None)
+    if inter_freq is None:
+        return None
+    return rf_freq - inter_freq
+
+
+def _prepare_line_scan(node: QualibrationNode[Parameters, Quam]):
+    """Prepare readout-line groupings and centers for the wide scan."""
+    all_qubits = get_qubits(node)
+    line_groups = [list(multiplexed_qubits.values()) for multiplexed_qubits in all_qubits.batch()]
+    line_probe_qubits = [group[0] for group in line_groups]
+    line_centers_rf = [float(np.mean([q.resonator.RF_frequency for q in group])) for group in line_groups]
+    line_if_centers = []
+    for group, rf_center in zip(line_groups, line_centers_rf):
+        lo_freq = _get_upconverter_frequency(group[0].resonator)
+        if lo_freq is None:
+            raise ValueError(f"Missing upconverter frequency for readout line probe {group[0].name}")
+        line_if_centers.append(int(round(rf_center - lo_freq)))
+
+    node.namespace["all_qubits"] = all_qubits
+    node.namespace["line_groups"] = line_groups
+    node.namespace["line_probe_qubits"] = line_probe_qubits
+    node.namespace["line_centers_rf"] = line_centers_rf
+    node.namespace["line_if_centers"] = line_if_centers
+    # Use line probes as the dataset "qubits" axis
+    node.namespace["qubits"] = line_probe_qubits
+
+
 # %% {Create_QUA_program}
 @node.run_action(skip_if=node.parameters.load_data_id is not None)
 def create_qua_program(node: QualibrationNode[Parameters, Quam]):
-    """Create the sweep axes and generate the QUA program from the pulse sequence and the node parameters."""
+    """Create readout-line sweep axes and the wide-scan QUA program."""
     # Class containing tools to help handle units and conversions.
     u = unit(coerce_to_integer=True)
-    # Get the active qubits from the node and organize them by batches
-    node.namespace["qubits"] = qubits = get_qubits(node)
-    num_qubits = len(qubits)
+    # Prepare line-level probing and store qubit groupings
+    _prepare_line_scan(node)
+    line_probe_qubits = node.namespace["line_probe_qubits"]
+    line_groups = node.namespace["line_groups"]
+    line_if_centers = node.namespace["line_if_centers"]
+    num_lines = len(line_probe_qubits)
     # Extract the sweep parameters and axes from the node parameters
     n_avg = node.parameters.num_shots
     # The frequency sweep around the resonator resonance frequency
@@ -87,7 +133,7 @@ def create_qua_program(node: QualibrationNode[Parameters, Quam]):
     dfs = np.arange(-span / 2, +span / 2, step)
     # Register the sweep axes to be added to the dataset when fetching data
     node.namespace["sweep_axes"] = {
-        "qubit": xr.DataArray(qubits.get_names()),
+        "qubit": xr.DataArray([q.name for q in line_probe_qubits]),
         "detuning": xr.DataArray(dfs, attrs={"long_name": "readout frequency", "units": "Hz"}),
     }
 
@@ -96,38 +142,37 @@ def create_qua_program(node: QualibrationNode[Parameters, Quam]):
         I, I_st, Q, Q_st, n, n_st = node.machine.declare_qua_variables()
         df = declare(int)  # QUA variable for the readout frequency
 
-        for multiplexed_qubits in qubits.batch():
+        for line_index, line_qubits in enumerate(line_groups):
             # Initialize the QPU in terms of flux points (flux tunable transmons and/or tunable couplers)
-            for qubit in multiplexed_qubits.values():
+            for qubit in line_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(df, dfs)):
-                    for i, qubit in multiplexed_qubits.items():
-                        rr = qubit.resonator
-                        # Update the resonator frequencies for all resonators
-                        rr.update_frequency(df + rr.intermediate_frequency)
-                        # Measure 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])
-                    align()
+                    probe_rr = line_probe_qubits[line_index].resonator
+                    # Sweep the readout line around its center frequency
+                    probe_rr.update_frequency(df + line_if_centers[line_index])
+                    # Measure the readout line response
+                    probe_rr.measure("readout", qua_vars=(I[line_index], Q[line_index]))
+                    # wait for the resonator to deplete
+                    probe_rr.wait(probe_rr.depletion_time * u.ns)
+                    # save data
+                    save(I[line_index], I_st[line_index])
+                    save(Q[line_index], Q_st[line_index])
+                align()
 
         with stream_processing():
             n_st.save("n")
-            for i in range(num_qubits):
+            for i in range(num_lines):
                 I_st[i].buffer(len(dfs)).average().save(f"I{i + 1}")
                 Q_st[i].buffer(len(dfs)).average().save(f"Q{i + 1}")
 
 
 # %% {Simulate}
 @node.run_action(skip_if=node.parameters.load_data_id is not None or not node.parameters.simulate)
 def simulate_qua_program(node: QualibrationNode[Parameters, Quam]):
-    """Connect to the QOP and simulate the QUA program"""
+    """Connect to the QOP and simulate the QUA program."""
     # Connect to the QOP
     qmm = node.machine.connect()
     # Get the config from the machine
@@ -145,6 +190,7 @@ def execute_qua_program(node: QualibrationNode[Parameters, Quam]):
     Connect to the QOP, execute the QUA program and fetch the raw data.
 
     Stores results in an xarray dataset called "ds_raw".
+    This run action performs the wide scan per readout line.
     """
     # Connect to the QOP
     qmm = node.machine.connect()
@@ -171,25 +217,71 @@ def execute_qua_program(node: QualibrationNode[Parameters, Quam]):
 # %% {Load_historical_data}
 @node.run_action(skip_if=node.parameters.load_data_id is None)
 def load_data(node: QualibrationNode[Parameters, Quam]):
-    """Load a previously acquired dataset."""
+    """Load a previously acquired dataset and rebuild line metadata."""
     load_data_id = node.parameters.load_data_id
     # Load the specified dataset
     node.load_from_id(node.parameters.load_data_id)
     node.parameters.load_data_id = load_data_id
-    # Get the active qubits from the loaded node parameters
-    node.namespace["qubits"] = get_qubits(node)
+    # Prepare line-level probing and store qubit groupings
+    _prepare_line_scan(node)
 
 
 # %% {Analyse_data}
 @node.run_action(skip_if=node.parameters.simulate)
 def analyse_data(node: QualibrationNode[Parameters, Quam]):
     """
-    Analyse the raw data and store the fitted data in another xarray dataset "ds_fit"
-    and the fitted results in the "fit_results" dictionary.
+    Analyse the raw data and store the fitted data in "ds_fit".
+    Identify multiple resonators per readout line and store results in "fit_results".
     """
-    node.results["ds_raw"] = process_raw_dataset(node.results["ds_raw"], node)
-    node.results["ds_fit"], fit_results = fit_raw_data(node.results["ds_raw"], node)
-    node.results["fit_results"] = {k: asdict(v) for k, v in fit_results.items()}
+    line_probe_qubits = node.namespace["line_probe_qubits"]
+    line_groups = node.namespace["line_groups"]
+    line_centers_rf = node.namespace["line_centers_rf"]
+
+    ds_raw = convert_IQ_to_V(node.results["ds_raw"], line_probe_qubits)
+    ds_raw = add_amplitude_and_phase(ds_raw, "detuning", subtract_slope_flag=True)
+    line_centers = xr.DataArray(line_centers_rf, coords={"qubit": ds_raw.qubit.values})
+    full_freq = np.array([ds_raw.detuning.values + center for center in line_centers_rf])
+    ds_raw = ds_raw.assign_coords(full_freq=(["qubit", "detuning"], full_freq))
+    ds_raw.full_freq.attrs = {"long_name": "RF frequency", "units": "Hz"}
+    node.results["ds_raw"] = ds_raw
+
+    # Identify multiple resonators per readout line and map to qubits
+    fit_results = {}
+    for line_index, line_qubits in enumerate(line_groups):
+        probe = line_probe_qubits[line_index]
+        probe_name = probe.name
+        rf_center = line_centers_rf[line_index]
+        da = ds_raw.IQ_abs.sel(qubit=probe_name)
+        num_resonators = len(line_qubits)
+
+        peak_candidates = []
+        for peak_number in range(1, num_resonators + 1):
+            peak_fit = peaks_dips(da, "detuning", number=peak_number)
+            detuning = float(peak_fit.position.values)
+            if np.isnan(detuning):
+                continue
+            peak_freq = rf_center + detuning
+            peak_width = float(np.abs(peak_fit.width.values))
+            peak_candidates.append((peak_freq, peak_width))
+
+        peak_candidates.sort(key=lambda item: item[0])
+        line_qubits_sorted = sorted(line_qubits, key=lambda q: q.resonator.RF_frequency)
+
+        for qubit, peak in zip(line_qubits_sorted, peak_candidates):
+            fit_results[qubit.name] = {
+                "frequency": float(peak[0]),
+                "fwhm": float(peak[1]),
+                "success": True,
+            }
+
+        for qubit in line_qubits_sorted[len(peak_candidates) :]:
+            fit_results[qubit.name] = {
+                "frequency": float(qubit.resonator.RF_frequency),
+                "fwhm": float("nan"),
+                "success": False,
+            }
+
+    node.results["fit_results"] = fit_results
 
     # Log the relevant information extracted from the data analysis
     log_fitted_results(node.results["fit_results"], log_callable=node.log)
@@ -202,25 +294,65 @@ def analyse_data(node: QualibrationNode[Parameters, Quam]):
 # %% {Plot_data}
 @node.run_action(skip_if=node.parameters.simulate)
 def plot_data(node: QualibrationNode[Parameters, Quam]):
-    """Plot the raw and fitted data in specific figures whose shape is given by qubit.grid_location."""
+    """Plot line-probe raw and fitted data using the qubit grid layout."""
+    u = unit(coerce_to_integer=True)
+    ds_raw = node.results["ds_raw"]
+    line_groups = node.namespace["line_groups"]
+    line_probe_qubits = node.namespace["line_probe_qubits"]
+    line_centers_rf = node.namespace["line_centers_rf"]
+
     fig_raw_phase = plot_raw_phase(node.results["ds_raw"], node.namespace["qubits"])
-    fig_fit_amplitude = plot_raw_amplitude_with_fit(
-        node.results["ds_raw"], node.namespace["qubits"], node.results["ds_fit"]
-    )
+    grid = QubitGrid(ds_raw, [q.grid_location for q in node.namespace["qubits"]])
+    axes_by_qubit = {}
+    for ax, qubit in grid_iter(grid):
+        (ds_raw.assign_coords(full_freq_GHz=ds_raw.full_freq / u.GHz).loc[qubit].IQ_abs / u.mV).plot(
+            ax=ax, x="full_freq_GHz"
+        )
+        ax.set_xlabel("RF frequency [GHz]")
+        ax.set_ylabel(r"$R=\sqrt{I^2 + Q^2}$ [mV]")
+        ax2 = ax.twiny()
+        (ds_raw.assign_coords(detuning_MHz=ds_raw.detuning / u.MHz).loc[qubit].IQ_abs / u.mV).plot(
+            ax=ax2, x="detuning_MHz"
+        )
+        ax2.set_xlabel("Detuning [MHz]")
+        axes_by_qubit[qubit["qubit"]] = (ax, ax2)
+
+    for line_index, probe in enumerate(line_probe_qubits):
+        line_qubits = line_groups[line_index]
+        rf_center = line_centers_rf[line_index]
+        peak_freqs = [
+            node.results["fit_results"][q.name]["frequency"]
+            for q in line_qubits
+            if node.results["fit_results"][q.name]["success"]
+        ]
+        if not peak_freqs:
+            continue
+        axes = axes_by_qubit.get(probe.name)
+        if axes is None:
+            continue
+        ax, ax2 = axes
+        for peak_freq in peak_freqs:
+            ax.axvline(peak_freq / u.GHz, color="tab:red", linestyle="--", alpha=0.5)
+            ax2.axvline((peak_freq - rf_center) / u.MHz, color="tab:red", linestyle="--", alpha=0.5)
+
+    grid.fig.suptitle("Resonator spectroscopy (amplitude + multi-peak)")
+    grid.fig.set_size_inches(15, 9)
+    grid.fig.tight_layout()
+    fig_raw_amplitude = grid.fig
     plt.show()
     # Store the generated figures
     node.results["figures"] = {
         "phase": fig_raw_phase,
-        "amplitude": fig_fit_amplitude,
+        "amplitude": fig_raw_amplitude,
     }
 
 
 # %% {Update_state}
 @node.run_action(skip_if=node.parameters.simulate)
 def update_state(node: QualibrationNode[Parameters, Quam]):
-    """Update the relevant parameters if the qubit data analysis was successful."""
+    """Update resonator frequencies for all successfully assigned peaks."""
     with node.record_state_updates():
-        for q in node.namespace["qubits"]:
+        for q in node.namespace["all_qubits"]:
             if node.outcomes[q.name] == "failed":
                 continue