Repository: QCoDeS/Qcodes
Pull request #5517 (2023-11-14 11:19:22)

Fix expansion of multi index when loading im memory cache from a netcdf file
- [x] News fragment


Patch:
diff --git a/docs/changes/newsfragments/5517.improved b/docs/changes/newsfragments/5517.improved
new file mode 100644
index 000000000..582cd7a62
--- /dev/null
+++ b/docs/changes/newsfragments/5517.improved
@@ -0,0 +1,2 @@
+QCoDeS in memory datasets exported to netcdf with MultiIndexes are now
+correctly decompressed when reloaded with load_by_id and friends.
diff --git a/src/qcodes/dataset/data_set_in_memory.py b/src/qcodes/dataset/data_set_in_memory.py
index a6a3c1640..855d9ce3c 100644
--- a/src/qcodes/dataset/data_set_in_memory.py
+++ b/src/qcodes/dataset/data_set_in_memory.py
@@ -342,12 +342,14 @@ class DataSetInMem(BaseDataSet):
 
     @classmethod
     def _set_cache_from_netcdf(cls, ds: DataSetInMem, xr_path: str | None) -> bool:
+        import cf_xarray as cfxr
         import xarray as xr
 
         success = True
         if xr_path is not None:
             try:
                 loaded_data = xr.load_dataset(xr_path, engine="h5netcdf")
+                loaded_data = cfxr.coding.decode_compress_to_multi_index(loaded_data)
                 ds._cache = DataSetCacheInMem(ds)
                 ds._cache._data = cls._from_xarray_dataset_to_qcodes_raw_data(
                     loaded_data
diff --git a/src/qcodes/dataset/data_set_protocol.py b/src/qcodes/dataset/data_set_protocol.py
index e1479a4ef..ee1fba802 100644
--- a/src/qcodes/dataset/data_set_protocol.py
+++ b/src/qcodes/dataset/data_set_protocol.py
@@ -218,7 +218,7 @@ class DataSetProtocol(Protocol):
     def export(
         self,
         export_type: DataExportType | str | None = None,
-        path: str | None = None,
+        path: Path | str | None = None,
         prefix: str | None = None,
         automatic_export: bool = False,
     ) -> None:


Test patch:
diff --git a/tests/dataset/test_dataset_export.py b/tests/dataset/test_dataset_export.py
index 2ead10552..24774762a 100644
--- a/tests/dataset/test_dataset_export.py
+++ b/tests/dataset/test_dataset_export.py
@@ -12,7 +12,15 @@ from numpy.testing import assert_allclose
 from pytest import LogCaptureFixture, TempPathFactory
 
 import qcodes
-from qcodes.dataset import get_data_export_path, load_from_netcdf, new_data_set
+from qcodes.dataset import (
+    DataSetProtocol,
+    DataSetType,
+    Measurement,
+    get_data_export_path,
+    load_by_id,
+    load_from_netcdf,
+    new_data_set,
+)
 from qcodes.dataset.data_set import DataSet
 from qcodes.dataset.descriptions.dependencies import InterDependencies_
 from qcodes.dataset.descriptions.param_spec import ParamSpecBase
@@ -180,6 +188,29 @@ def _make_mock_dataset_non_grid(experiment) -> DataSet:
     return dataset
 
 
+@pytest.fixture(name="mock_dataset_non_grid_in_mem")
+def _make_mock_dataset_non_grid_in_mem(experiment) -> DataSetProtocol:
+    meas = Measurement(exp=experiment, name="in_mem_ds")
+
+    num_samples = 50
+
+    rng = np.random.default_rng()
+
+    x_vals = rng.random(num_samples) * 10
+    y_vals = 20 + rng.random(num_samples) * 5
+
+    meas.register_custom_parameter("x")
+    meas.register_custom_parameter("y")
+    meas.register_custom_parameter("z", setpoints=("x", "y"))
+
+    with meas.run(dataset_class=DataSetType.DataSetInMem) as datasaver:
+        for x, y in zip(x_vals, y_vals):
+            results = [("x", x), ("y", y), ("z", x + y)]
+            datasaver.add_result(*results)
+
+    return datasaver.dataset
+
+
 @pytest.fixture(name="mock_dataset_non_grid_in_grid")
 def _make_mock_dataset_non_grid_in_grid(experiment) -> DataSet:
     dataset = new_data_set("dataset")
@@ -864,6 +895,49 @@ def test_export_non_grid_dataset(
     assert xr_ds.identical(xr_ds_reimported)
 
 
+def test_export_non_grid_in_mem_dataset(
+    tmp_path_factory: TempPathFactory, mock_dataset_non_grid_in_mem: DataSetProtocol
+) -> None:
+    tmp_path = tmp_path_factory.mktemp("export_netcdf")
+    path = str(tmp_path)
+    mock_dataset_non_grid_in_mem.export(
+        export_type="netcdf", path=tmp_path, prefix="qcodes_"
+    )
+
+    pdf = mock_dataset_non_grid_in_mem.to_pandas_dataframe()
+    dims = _calculate_index_shape(pdf.index)
+    assert dims == {"x": 50, "y": 50}
+
+    xr_ds = mock_dataset_non_grid_in_mem.to_xarray_dataset()
+    assert len(xr_ds.coords) == 3
+    assert "multi_index" in xr_ds.coords
+    assert "x" in xr_ds.coords
+    assert "y" in xr_ds.coords
+    assert xr_ds.dims == {"multi_index": 50}
+
+    expected_path = f"qcodes_{mock_dataset_non_grid_in_mem.captured_run_id}_{mock_dataset_non_grid_in_mem.guid}.nc"
+    assert os.listdir(path) == [expected_path]
+    file_path = os.path.join(path, expected_path)
+    ds = load_from_netcdf(file_path)
+
+    xr_ds_reimported = ds.to_xarray_dataset()
+
+    assert len(xr_ds_reimported.coords) == 3
+    assert "multi_index" in xr_ds_reimported.coords
+    assert "x" in xr_ds_reimported.coords
+    assert "y" in xr_ds_reimported.coords
+    assert xr_ds.identical(xr_ds_reimported)
+
+    loaded_by_id_ds = load_by_id(mock_dataset_non_grid_in_mem.run_id)
+    loaded_by_id_ds_xr = loaded_by_id_ds.to_xarray_dataset()
+
+    assert len(loaded_by_id_ds_xr.coords) == 3
+    assert "multi_index" in loaded_by_id_ds_xr.coords
+    assert "x" in loaded_by_id_ds_xr.coords
+    assert "y" in loaded_by_id_ds_xr.coords
+    assert xr_ds.identical(loaded_by_id_ds_xr)
+
+
 def test_export_non_grid_in_grid_dataset(
     tmp_path_factory: TempPathFactory, mock_dataset_non_grid_in_grid: DataSet
 ) -> None: