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]