From 5aea2ca76fff9a6960d242f7e013df887932ead8 Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Thu, 27 Aug 2026 10:46:12 +0200 Subject: [PATCH] Fix inferred parameter export when index cannot be expanded `Series(...).to_xarray()` expands a pandas index into one dimension per index level. That is only compatible with the target xarray dataset when the dataset dimensions are exactly the levels of the index. It is not the case when the dataset uses a single `multi_index` dimension or the flat index created by `DataFrame.reset_index()`, and a non unique MultiIndex cannot be converted at all. Only take the reindexing path when the index is unique and the dataset dimensions match the index level names. In the remaining cases the data is already in index order and can be reshaped directly. Adds tests covering export of an inferred parameter with a `multi_index` dimension and with a non unique MultiIndex. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 92aac34d-1a70-4ca9-900c-4ab563a731df --- .../dataset/exporters/export_to_xarray.py | 13 ++- tests/dataset/test_dataset_export.py | 84 +++++++++++++++++++ 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/src/qcodes/dataset/exporters/export_to_xarray.py b/src/qcodes/dataset/exporters/export_to_xarray.py index 3e564dc08f3..528aca68c7d 100644 --- a/src/qcodes/dataset/exporters/export_to_xarray.py +++ b/src/qcodes/dataset/exporters/export_to_xarray.py @@ -84,6 +84,17 @@ def _add_inferred_data_vars( dep_names = {dep.name for dep in deps} dims = tuple(d for d in xr_dataset.dims) + # ``Series.to_xarray()`` expands the index into one dimension per index + # level. That is only compatible with the target dataset if the dataset + # dimensions are exactly the levels of the index. It is not the case when + # the dataset uses a single ``multi_index`` dimension or the flat index + # created by ``DataFrame.reset_index()``. A non unique MultiIndex cannot be + # converted at all. In those cases the data is already in index order so it + # can be reshaped directly. + index_is_expandable = ( + index is not None and index.is_unique and set(dims) == set(index.names) + ) + for inf in inferred: if inf.name in dep_names: continue @@ -108,7 +119,7 @@ def _add_inferred_data_vars( expected_shape = tuple(xr_dataset.sizes[d] for d in dims) expected_size = prod(expected_shape) if flat.shape[0] == expected_size: - if index is not None: + if index is not None and index_is_expandable: # If an index is provided, we should align the inferred data with the index. # This is necessary because data may be reordered when transforming from a pandas DataFrame to an xarray Dataset. # Passing an index allows the original data ordering to be preserved on reconstruction. diff --git a/tests/dataset/test_dataset_export.py b/tests/dataset/test_dataset_export.py index 16dfc2155a3..217dc94c99e 100644 --- a/tests/dataset/test_dataset_export.py +++ b/tests/dataset/test_dataset_export.py @@ -305,6 +305,62 @@ def _make_mock_dataset_non_grid(experiment: Experiment) -> DataSet: return dataset +@pytest.fixture(name="mock_dataset_non_grid_inferred") +def _make_mock_dataset_non_grid_inferred(experiment: Experiment) -> DataSet: + """Non grid dataset where an inferred parameter is inferred from z.""" + dataset = new_data_set("dataset") + xparam = ParamSpecBase("x", "numeric") + yparam = ParamSpecBase("y", "numeric") + zparam = ParamSpecBase("z", "numeric") + tparam = ParamSpecBase("t", "numeric") + idps = InterDependencies_( + dependencies={zparam: (xparam, yparam)}, + inferences={tparam: (zparam,)}, + ) + dataset.set_interdependencies(idps) + + num_samples = 50 + + rng = np.random.default_rng(1234) + + x_vals = rng.random(num_samples) * 10 + y_vals = 20 + rng.random(num_samples) * 5 + + dataset.mark_started() + + for i, (x, y) in enumerate(zip(x_vals, y_vals)): + dataset.add_results([{"x": x, "y": y, "z": x + y, "t": float(i)}]) + dataset.mark_completed() + return dataset + + +@pytest.fixture(name="mock_dataset_non_unique_index_inferred") +def _make_mock_dataset_non_unique_index_inferred(experiment: Experiment) -> DataSet: + """Dataset with a non unique MultiIndex and an inferred parameter.""" + dataset = new_data_set("dataset") + xparam = ParamSpecBase("x", "numeric") + yparam = ParamSpecBase("y", "numeric") + zparam = ParamSpecBase("z", "numeric") + tparam = ParamSpecBase("t", "numeric") + idps = InterDependencies_( + dependencies={zparam: (xparam, yparam)}, + inferences={tparam: (zparam,)}, + ) + dataset.set_interdependencies(idps) + + num_samples = 20 + # every (x, y) pair is measured twice making the index non unique + x_vals = np.repeat(np.arange(num_samples // 2, dtype=float), 2) + y_vals = np.repeat(np.arange(num_samples // 2, dtype=float), 2) + + dataset.mark_started() + + for i, (x, y) in enumerate(zip(x_vals, y_vals)): + dataset.add_results([{"x": x, "y": y, "z": x + y, "t": float(i)}]) + dataset.mark_completed() + return dataset + + @pytest.fixture(name="mock_dataset_non_grid_in_mem") def _make_mock_dataset_non_grid_in_mem(experiment: Experiment) -> DataSetProtocol: meas = Measurement(exp=experiment, name="in_mem_ds") @@ -1760,6 +1816,34 @@ def test_multi_index_options_non_grid(mock_dataset_non_grid: DataSet) -> None: assert xds_always.sizes == {"multi_index": 50} +@pytest.mark.parametrize("use_multi_index", ["auto", "always"]) +def test_multi_index_export_with_inferred_parameter( + mock_dataset_non_grid_inferred: DataSet, use_multi_index: str +) -> None: + """Inferred parameters must export correctly when a MultiIndex dim is used.""" + xds = mock_dataset_non_grid_inferred.to_xarray_dataset( + use_multi_index=use_multi_index # pyright: ignore[reportArgumentType] + ) + + assert xds.sizes == {"multi_index": 50} + assert "t" in xds.data_vars + assert xds["t"].dims == ("multi_index",) + np.testing.assert_array_equal(xds["t"].values, np.arange(50, dtype=float)) + + +def test_non_unique_multi_index_export_with_inferred_parameter( + mock_dataset_non_unique_index_inferred: DataSet, +) -> None: + """A non unique MultiIndex must not break export of inferred parameters.""" + xds = mock_dataset_non_unique_index_inferred.to_xarray_dataset() + + assert "t" in xds.data_vars + assert xds["t"].dims == xds["z"].dims + np.testing.assert_array_equal( + np.asarray(xds["t"].values).ravel(), np.arange(20, dtype=float) + ) + + def test_multi_index_wrong_option(mock_dataset_non_grid: DataSet) -> None: with pytest.raises(ValueError, match="Invalid value for use_multi_index"): mock_dataset_non_grid.to_xarray_dataset(use_multi_index=True) # pyright: ignore[reportArgumentType]