diff --git a/qualibration_libs/data/fetcher.py b/qualibration_libs/data/fetcher.py index 22750fc..9b44fa2 100644 --- a/qualibration_libs/data/fetcher.py +++ b/qualibration_libs/data/fetcher.py @@ -198,6 +198,25 @@ def update_dataset(self): if dims_order[0] not in ["qubit", "qubit_pair"]: logger.error("first axis must be either qubit or qubit_pair") + # 0-D ndarray streams (e.g. shot counter "n" from n_st.save("n")) sit next to 1-D + # swept I/Q buffers. They must not participate in the uniform-shape check below. + # Expose scalars on self.data for progress_counter(..., data_fetcher.get("n", 0)). + for label in list(raw_data_arrays.keys()): + arr = raw_data_arrays[label] + if isinstance(arr, np.ndarray) and arr.ndim == 0: + self.data[label] = arr.item() + del raw_data_arrays[label] + logger.debug( + "Extracted 0-D stream %r for fetcher.data; excluded from swept-array layout", + label, + ) + + if not raw_data_arrays: + logger.debug( + "No swept-array handles after 0-D extraction; returning current dataset." + ) + return self.dataset + # Determine reference shape from non-None entries. data_arrays = [d for d in raw_data_arrays.values() if isinstance(d, np.ndarray)] if data_arrays: diff --git a/qualibration_libs/tests/test_fetcher.py b/qualibration_libs/tests/test_fetcher.py index 7c331ac..f33e7fe 100644 --- a/qualibration_libs/tests/test_fetcher.py +++ b/qualibration_libs/tests/test_fetcher.py @@ -1,7 +1,10 @@ """Tests for exception handling in data/fetcher.py.""" +import numpy as np import pytest -from unittest.mock import Mock +import xarray as xr +from unittest.mock import MagicMock, Mock, patch + from qualibration_libs.data.fetcher import XarrayDataFetcher @@ -29,3 +32,64 @@ def test_invalid_key_raises_keyerror(self): assert "Data key 'invalid_key' not found in XarrayDataFetcher" in error_msg assert "Available keys: 'dataset'" in error_msg assert exc_info.value.__cause__ is not None + + +@patch("qualibration_libs.data.fetcher.fetching_tool") +def test_zero_dim_shot_counter_with_swept_iq_does_not_raise(mock_fetching_tool): + """Regression: n_st.save('n') yields 0-D 'n' while I*/Q* are 1-D along detuning.""" + mock_result = MagicMock() + mock_result.fetch_all.return_value = [ + np.array(42, dtype=np.int64), + np.linspace(0, 1, 5), + np.linspace(1, 0, 5), + ] + mock_result.is_processing.return_value = False + mock_result.get_start_time.return_value = 0.0 + mock_fetching_tool.return_value = mock_result + + mock_job = Mock() + mock_job.result_handles.keys.return_value = ["n", "I1", "Q1"] + + axes = { + "qubit": xr.DataArray(["q0"]), + "detuning": xr.DataArray(np.linspace(-1e6, 1e6, 5), attrs={"units": "Hz"}), + } + fetcher = XarrayDataFetcher(job=mock_job, axes=axes) + fetcher.retrieve_latest_data() + fetcher.update_dataset() + + assert fetcher.get("n") == 42 + assert "I" in fetcher.dataset.data_vars + assert "Q" in fetcher.dataset.data_vars + assert fetcher.dataset["I"].shape == (1, 5) + assert fetcher.dataset["Q"].shape == (1, 5) + + +@patch("qualibration_libs.data.fetcher.fetching_tool") +def test_zero_dim_n_two_qubits(mock_fetching_tool): + mock_result = MagicMock() + mock_result.fetch_all.return_value = [ + np.array(100, dtype=np.int32), + np.arange(3.0), + np.arange(3.0, 6.0), + np.arange(6.0, 9.0), + np.arange(9.0, 12.0), + ] + mock_result.is_processing.return_value = False + mock_result.get_start_time.return_value = 0.0 + mock_fetching_tool.return_value = mock_result + + mock_job = Mock() + mock_job.result_handles.keys.return_value = ["n", "I1", "Q1", "I2", "Q2"] + + axes = { + "qubit": xr.DataArray(["q0", "q1"]), + "detuning": xr.DataArray(np.zeros(3)), + } + fetcher = XarrayDataFetcher(job=mock_job, axes=axes) + fetcher.retrieve_latest_data() + fetcher.update_dataset() + + assert fetcher.get("n") == 100 + assert fetcher.dataset["I"].shape == (2, 3) + assert fetcher.dataset["Q"].shape == (2, 3)