Repository: QCoDeS/Qcodes
Pull request #5589 (2023-12-11 12:53:39)

Smarter selection of multi index export to xarray
* Only automatically select multi index export if shapes are unset
* Allow the user to select which export type to use

- [x] Tests?
- [x] Changelog


Patch:
diff --git a/docs/changes/newsfragments/5589.improved b/docs/changes/newsfragments/5589.improved
new file mode 100644
index 000000000..489f7e8ff
--- /dev/null
+++ b/docs/changes/newsfragments/5589.improved
@@ -0,0 +1,5 @@
+The use of multi index when exporting to xarray (and netcdf files) has been made smarter
+such that any dataset with a known shape (such as those measured by doNd etc) will
+never be automatically exported using multi index even in the case of incomplete datasets (i.e. due to an interrupted measurement).
+Furthermore `to_xarray_dataset` and `to_xarray_dataarray_dict` have gained a key word argument `use_multi_index` to allow the user
+to control the use of multi indexes.
diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py
index 2ad45c82d..c58aef15f 100644
--- a/src/qcodes/dataset/data_set.py
+++ b/src/qcodes/dataset/data_set.py
@@ -11,7 +11,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from queue import Queue
 from threading import Thread
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy
 from tqdm.auto import trange
@@ -971,6 +971,7 @@ class DataSet(BaseDataSet):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> dict[str, xr.DataArray]:
         """
         Returns the values stored in the :class:`.DataSet` for the specified parameters
@@ -1001,6 +1002,18 @@ class DataSet(BaseDataSet):
                 if None
             end: end value of selection range (by results count); ignored if
                 None
+            use_multi_index: Should the data be exported using a multi index
+                rather than regular cartesian indexes. With regular cartesian
+                coordinates, the xarray dimensions are calculated from the sets or all
+                values along the setpoint axis of the QCoDeS dataset. Any position
+                in this grid not corresponding to a measured value will be filled
+                with a placeholder (typically NaN) potentially creating a sparse
+                dataset with significant storage overhead.
+                Multi index avoids this and is therefor better
+                suited for data that is known to not be on a grid.
+                If set to "auto" multi index will be used if projecting the data onto
+                a grid requires filling non measured values with NaN  and the shapes
+                of the data has not been set in the run description.
 
         Returns:
             Dictionary from requested parameter names to :py:class:`xr.DataArray` s
@@ -1012,10 +1025,10 @@ class DataSet(BaseDataSet):
 
                 dataarray_dict = ds.to_xarray_dataarray_dict()
         """
-        data = self.get_parameter_data(*params,
-                                       start=start,
-                                       end=end)
-        datadict = load_to_xarray_dataarray_dict(self, data)
+        data = self.get_parameter_data(*params, start=start, end=end)
+        datadict = load_to_xarray_dataarray_dict(
+            self, data, use_multi_index=use_multi_index
+        )
 
         return datadict
 
@@ -1024,6 +1037,7 @@ class DataSet(BaseDataSet):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> xr.Dataset:
         """
         Returns the values stored in the :class:`.DataSet` for the specified parameters
@@ -1052,7 +1066,18 @@ class DataSet(BaseDataSet):
                 if None
             end: end value of selection range (by results count); ignored if
                 None
-
+            use_multi_index: Should the data be exported using a multi index
+                rather than regular cartesian indexes. With regular cartesian
+                coordinates, the xarray dimensions are calculated from the sets or all
+                values along the setpoint axis of the QCoDeS dataset. Any position
+                in this grid not corresponding to a measured value will be filled
+                with a placeholder (typically NaN) potentially creating a sparse
+                dataset with significant storage overhead.
+                Multi index avoids this and is therefor better
+                suited for data that is known to not be on a grid.
+                If set to "auto" multi index will be used if projecting the data onto
+                a grid requires filling non measured values with NaN  and the shapes
+                of the data has not been set in the run description.
         Returns:
             :py:class:`xr.Dataset` with the requested parameter(s) data as
             :py:class:`xr.DataArray` s and coordinates formed by the dependencies.
@@ -1066,7 +1091,7 @@ class DataSet(BaseDataSet):
                                        start=start,
                                        end=end)
 
-        return load_to_xarray_dataset(self, data)
+        return load_to_xarray_dataset(self, data, use_multi_index=use_multi_index)
 
     def write_data_to_text_file(
         self, path: str, single_file: bool = False, single_file_name: str | None = None
diff --git a/src/qcodes/dataset/data_set_cache.py b/src/qcodes/dataset/data_set_cache.py
index a00f6b77a..49ff2b4f0 100644
--- a/src/qcodes/dataset/data_set_cache.py
+++ b/src/qcodes/dataset/data_set_cache.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, Literal, TypeVar
 
 import numpy as np
 
@@ -177,7 +177,9 @@ class DataSetCache(Generic[DatasetType]):
         """
         return self.to_pandas_dataframe_dict()
 
-    def to_xarray_dataarray_dict(self) -> dict[str, xr.DataArray]:
+    def to_xarray_dataarray_dict(
+        self, *, use_multi_index: Literal["auto", "always", "never"] = "auto"
+    ) -> dict[str, xr.DataArray]:  # noqa: F821
         """
         Returns the values stored in the :class:`.dataset.data_set.DataSet` as a dict of
         :py:class:`xr.DataArray` s
@@ -190,9 +192,13 @@ class DataSetCache(Generic[DatasetType]):
 
         """
         data = self.data()
-        return load_to_xarray_dataarray_dict(self._dataset, data)
+        return load_to_xarray_dataarray_dict(
+            self._dataset, data, use_multi_index=use_multi_index
+        )
 
-    def to_xarray_dataset(self) -> xr.Dataset:
+    def to_xarray_dataset(
+        self, *, use_multi_index: Literal["auto", "always", "never"] = "auto"
+    ) -> xr.Dataset:
         """
         Returns the values stored in the :class:`.dataset.data_set.DataSet` as a
         :py:class:`xr.Dataset` object.
@@ -207,7 +213,9 @@ class DataSetCache(Generic[DatasetType]):
 
         """
         data = self.data()
-        return load_to_xarray_dataset(self._dataset, data)
+        return load_to_xarray_dataset(
+            self._dataset, data, use_multi_index=use_multi_index
+        )
 
 
 def load_new_data_from_db_and_append(
diff --git a/src/qcodes/dataset/data_set_in_memory.py b/src/qcodes/dataset/data_set_in_memory.py
index 855d9ce3c..463214290 100644
--- a/src/qcodes/dataset/data_set_in_memory.py
+++ b/src/qcodes/dataset/data_set_in_memory.py
@@ -8,7 +8,7 @@ import time
 import warnings
 from collections.abc import Mapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import numpy as np
 
@@ -838,6 +838,7 @@ class DataSetInMem(BaseDataSet):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> dict[str, xr.DataArray]:
         self._warn_if_set(*params, start=start, end=end)
         return self.cache.to_xarray_dataarray_dict()
@@ -847,6 +848,7 @@ class DataSetInMem(BaseDataSet):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> xr.Dataset:
         self._warn_if_set(*params, start=start, end=end)
         return self.cache.to_xarray_dataset()
diff --git a/src/qcodes/dataset/data_set_protocol.py b/src/qcodes/dataset/data_set_protocol.py
index ee1fba802..e250d1130 100644
--- a/src/qcodes/dataset/data_set_protocol.py
+++ b/src/qcodes/dataset/data_set_protocol.py
@@ -7,7 +7,15 @@ import warnings
 from collections.abc import Mapping, Sequence
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Protocol, Union, runtime_checkable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    Protocol,
+    Union,
+    runtime_checkable,
+)
 
 import numpy as np
 from typing_extensions import TypeAlias
@@ -257,6 +265,7 @@ class DataSetProtocol(Protocol):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> dict[str, xr.DataArray]:
         ...
 
@@ -265,6 +274,7 @@ class DataSetProtocol(Protocol):
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        use_multi_index: Literal["auto", "always", "never"] = "auto",
     ) -> xr.Dataset:
         ...
 
diff --git a/src/qcodes/dataset/exporters/export_to_xarray.py b/src/qcodes/dataset/exporters/export_to_xarray.py
index 2e3512f90..4ca5793f5 100644
--- a/src/qcodes/dataset/exporters/export_to_xarray.py
+++ b/src/qcodes/dataset/exporters/export_to_xarray.py
@@ -5,7 +5,7 @@ import warnings
 from collections.abc import Hashable, Mapping
 from math import prod
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, cast
 
 import numpy as np
 from tqdm.dask import TqdmCallback
@@ -61,11 +61,19 @@ def _calculate_index_shape(idx: pd.Index | pd.MultiIndex) -> dict[Hashable, int]
 
 
 def _load_to_xarray_dataarray_dict_no_metadata(
-    dataset: DataSetProtocol, datadict: Mapping[str, Mapping[str, np.ndarray]]
+    dataset: DataSetProtocol,
+    datadict: Mapping[str, Mapping[str, np.ndarray]],
+    *,
+    use_multi_index: Literal["auto", "always", "never"] = "auto",
 ) -> dict[str, xr.DataArray]:
     import pandas as pd
     import xarray as xr
 
+    if use_multi_index not in ("auto", "always", "never"):
+        raise ValueError(
+            f"Invalid value for use_multi_index. Expected one of 'auto', 'always', 'never' but got {use_multi_index}"
+        )
+
     data_xrdarray_dict: dict[str, xr.DataArray] = {}
 
     for name, subdict in datadict.items():
@@ -94,8 +102,16 @@ def _load_to_xarray_dataarray_dict_no_metadata(
                 index_prod = prod(calc_index.values())
                 # if the product of the len of individual index dims == len(total_index)
                 # we are on a grid
+
                 on_grid = index_prod == len(index)
-                if not on_grid:
+
+                export_with_multi_index = (
+                    not on_grid
+                    and dataset.description.shapes is None
+                    and use_multi_index == "auto"
+                ) or use_multi_index == "always"
+
+                if export_with_multi_index:
                     assert isinstance(df.index, pd.MultiIndex)
 
                     if hasattr(xr, "Coordinates"):
@@ -115,9 +131,14 @@ def _load_to_xarray_dataarray_dict_no_metadata(
 
 
 def load_to_xarray_dataarray_dict(
-    dataset: DataSetProtocol, datadict: Mapping[str, Mapping[str, np.ndarray]]
+    dataset: DataSetProtocol,
+    datadict: Mapping[str, Mapping[str, np.ndarray]],
+    *,
+    use_multi_index: Literal["auto", "always", "never"] = "auto",
 ) -> dict[str, xr.DataArray]:
-    dataarrays = _load_to_xarray_dataarray_dict_no_metadata(dataset, datadict)
+    dataarrays = _load_to_xarray_dataarray_dict_no_metadata(
+        dataset, datadict, use_multi_index=use_multi_index
+    )
 
     for dataname, dataarray in dataarrays.items():
         _add_param_spec_to_xarray_coords(dataset, dataarray)
@@ -157,7 +178,12 @@ def _add_metadata_to_xarray(
             xrdataset.attrs[metadata_tag] = metadata
 
 
-def load_to_xarray_dataset(dataset: DataSetProtocol, data: ParameterData) -> xr.Dataset:
+def load_to_xarray_dataset(
+    dataset: DataSetProtocol,
+    data: ParameterData,
+    *,
+    use_multi_index: Literal["auto", "always", "never"] = "auto",
+) -> xr.Dataset:
     import xarray as xr
 
     if not _same_setpoints(data):
@@ -168,7 +194,9 @@ def load_to_xarray_dataset(dataset: DataSetProtocol, data: ParameterData) -> xr.
             "independent parameter to its own datarray."
         )
 
-    data_xrdarray_dict = _load_to_xarray_dataarray_dict_no_metadata(dataset, data)
+    data_xrdarray_dict = _load_to_xarray_dataarray_dict_no_metadata(
+        dataset, data, use_multi_index=use_multi_index
+    )
 
     # Casting Hashable for the key type until python/mypy#1114
     # and python/typing#445 are resolved.


Test patch:
diff --git a/tests/dataset/test_dataset_export.py b/tests/dataset/test_dataset_export.py
index 307c3cc49..f9f00195a 100644
--- a/tests/dataset/test_dataset_export.py
+++ b/tests/dataset/test_dataset_export.py
@@ -127,6 +127,74 @@ def _make_mock_dataset_grid(experiment) -> DataSet:
     return dataset
 
 
+@pytest.fixture(name="mock_dataset_grid_with_shapes")
+def _make_mock_dataset_grid_with_shapes(experiment) -> DataSet:
+    dataset = new_data_set("dataset")
+    xparam = ParamSpecBase("x", "numeric")
+    yparam = ParamSpecBase("y", "numeric")
+    zparam = ParamSpecBase("z", "numeric")
+    idps = InterDependencies_(dependencies={zparam: (xparam, yparam)})
+    dataset.set_interdependencies(idps, shapes={"z": (10, 5)})
+
+    dataset.mark_started()
+    for x in range(10):
+        for y in range(20, 25):
+            results = [{"x": x, "y": y, "z": x + y}]
+            dataset.add_results(results)
+    dataset.mark_completed()
+    return dataset
+
+
+@pytest.fixture(name="mock_dataset_grid_incomplete")
+def _make_mock_dataset_grid_incomplete(experiment) -> DataSet:
+    dataset = new_data_set("dataset")
+    xparam = ParamSpecBase("x", "numeric")
+    yparam = ParamSpecBase("y", "numeric")
+    zparam = ParamSpecBase("z", "numeric")
+    idps = InterDependencies_(dependencies={zparam: (xparam, yparam)})
+    dataset.set_interdependencies(idps)
+
+    dataset.mark_started()
+    break_loop = False
+
+    for x in range(10):
+        if break_loop is True:
+            break
+        for y in range(20, 25):
+            results = [{"x": x, "y": y, "z": x + y}]
+            dataset.add_results(results)
+            if y == 23 and x == 7:
+                break_loop = True
+                break
+    dataset.mark_completed()
+    return dataset
+
+
+@pytest.fixture(name="mock_dataset_grid_incomplete_with_shapes")
+def _make_mock_dataset_grid_incomplete_with_shapes(experiment) -> DataSet:
+    dataset = new_data_set("dataset")
+    xparam = ParamSpecBase("x", "numeric")
+    yparam = ParamSpecBase("y", "numeric")
+    zparam = ParamSpecBase("z", "numeric")
+    idps = InterDependencies_(dependencies={zparam: (xparam, yparam)})
+    dataset.set_interdependencies(idps, shapes={"z": (10, 5)})
+
+    dataset.mark_started()
+    break_loop = False
+
+    for x in range(10):
+        if break_loop is True:
+            break
+        for y in range(20, 25):
+            results = [{"x": x, "y": y, "z": x + y}]
+            dataset.add_results(results)
+            if y == 23 and x == 7:
+                break_loop = True
+                break
+    dataset.mark_completed()
+    return dataset
+
+
 @pytest.fixture(name="mock_dataset_numpy")
 def _make_mock_dataset_numpy(experiment) -> DataSet:
     dataset = new_data_set("dataset")
@@ -1136,3 +1204,106 @@ def _assert_xarray_metadata_is_as_expected(xarray_ds, qc_dataset):
     assert xarray_ds.parent_dataset_links == links_to_str(
         qc_dataset.parent_dataset_links
     )
+
+
+def test_multi_index_options_grid(mock_dataset_grid) -> None:
+    assert mock_dataset_grid.description.shapes is None
+
+    xds = mock_dataset_grid.to_xarray_dataset()
+    assert xds.dims == {"x": 10, "y": 5}
+
+    xds_never = mock_dataset_grid.to_xarray_dataset(use_multi_index="never")
+    assert xds_never.dims == {"x": 10, "y": 5}
+
+    xds_auto = mock_dataset_grid.to_xarray_dataset(use_multi_index="auto")
+    assert xds_auto.dims == {"x": 10, "y": 5}
+
+    xds_always = mock_dataset_grid.to_xarray_dataset(use_multi_index="always")
+    assert xds_always.dims == {"multi_index": 50}
+
+
+def test_multi_index_options_grid_with_shape(mock_dataset_grid_with_shapes) -> None:
+    assert mock_dataset_grid_with_shapes.description.shapes == {"z": (10, 5)}
+
+    xds = mock_dataset_grid_with_shapes.to_xarray_dataset()
+    assert xds.dims == {"x": 10, "y": 5}
+
+    xds_never = mock_dataset_grid_with_shapes.to_xarray_dataset(use_multi_index="never")
+    assert xds_never.dims == {"x": 10, "y": 5}
+
+    xds_auto = mock_dataset_grid_with_shapes.to_xarray_dataset(use_multi_index="auto")
+    assert xds_auto.dims == {"x": 10, "y": 5}
+
+    xds_always = mock_dataset_grid_with_shapes.to_xarray_dataset(
+        use_multi_index="always"
+    )
+    assert xds_always.dims == {"multi_index": 50}
+
+
+def test_multi_index_options_incomplete_grid(mock_dataset_grid_incomplete) -> None:
+    assert mock_dataset_grid_incomplete.description.shapes is None
+
+    xds = mock_dataset_grid_incomplete.to_xarray_dataset()
+    assert xds.dims == {"multi_index": 39}
+
+    xds_never = mock_dataset_grid_incomplete.to_xarray_dataset(use_multi_index="never")
+    assert xds_never.dims == {"x": 8, "y": 5}
+
+    xds_auto = mock_dataset_grid_incomplete.to_xarray_dataset(use_multi_index="auto")
+    assert xds_auto.dims == {"multi_index": 39}
+
+    xds_always = mock_dataset_grid_incomplete.to_xarray_dataset(
+        use_multi_index="always"
+    )
+    assert xds_always.dims == {"multi_index": 39}
+
+
+def test_multi_index_options_incomplete_grid_with_shapes(
+    mock_dataset_grid_incomplete_with_shapes,
+) -> None:
+    assert mock_dataset_grid_incomplete_with_shapes.description.shapes == {"z": (10, 5)}
+
+    xds = mock_dataset_grid_incomplete_with_shapes.to_xarray_dataset()
+    assert xds.dims == {"x": 8, "y": 5}
+
+    xds_never = mock_dataset_grid_incomplete_with_shapes.to_xarray_dataset(
+        use_multi_index="never"
+    )
+    assert xds_never.dims == {"x": 8, "y": 5}
+
+    xds_auto = mock_dataset_grid_incomplete_with_shapes.to_xarray_dataset(
+        use_multi_index="auto"
+    )
+    assert xds_auto.dims == {"x": 8, "y": 5}
+
+    xds_always = mock_dataset_grid_incomplete_with_shapes.to_xarray_dataset(
+        use_multi_index="always"
+    )
+    assert xds_always.dims == {"multi_index": 39}
+
+
+def test_multi_index_options_non_grid(mock_dataset_non_grid) -> None:
+    assert mock_dataset_non_grid.description.shapes is None
+
+    xds = mock_dataset_non_grid.to_xarray_dataset()
+    assert xds.dims == {"multi_index": 50}
+
+    xds_never = mock_dataset_non_grid.to_xarray_dataset(use_multi_index="never")
+    assert xds_never.dims == {"x": 50, "y": 50}
+
+    xds_auto = mock_dataset_non_grid.to_xarray_dataset(use_multi_index="auto")
+    assert xds_auto.dims == {"multi_index": 50}
+
+    xds_always = mock_dataset_non_grid.to_xarray_dataset(use_multi_index="always")
+    assert xds_always.dims == {"multi_index": 50}
+
+
+def test_multi_index_wrong_option(mock_dataset_non_grid) -> None:
+    with pytest.raises(ValueError, match="Invalid value for use_multi_index"):
+        mock_dataset_non_grid.to_xarray_dataset(use_multi_index=True)
+
+    with pytest.raises(ValueError, match="Invalid value for use_multi_index"):
+        mock_dataset_non_grid.to_xarray_dataset(use_multi_index=False)
+
+    with pytest.raises(ValueError, match="Invalid value for use_multi_index"):
+        mock_dataset_non_grid.to_xarray_dataset(use_multi_index="perhaps")