Merge pull request #7407 from jenshnielsen/jenshnielsen/xarray_export

Xarray export cleanup and bugfixes

Author: jenshnielsen

pyproject.toml

@@ -42,7 +42,7 @@ dependencies = [
     "uncertainties>=3.2.0",
     "versioningit>=2.2.1",
     "websockets>=11.0",
-    "xarray>=2022.06.0",
+    "xarray>=2023.08.0",
     "cf_xarray>=0.8.4",
     "opentelemetry-api>=1.17.0",
     "pillow>=9.2.0",

src/qcodes/dataset/exporters/export_to_pandas.py

@@ -108,7 +108,7 @@ def _generate_pandas_index(
             else:
                 index_data.append(data[key].ravel())
 
-        index = pd.MultiIndex.from_arrays(index_data, names=keys[1:])
+        index = pd.MultiIndex.from_arrays(index_data, names=list(deps_data.keys()))
     return index
 
 

src/qcodes/dataset/exporters/export_to_xarray.py

@@ -118,14 +118,10 @@ def _load_to_xarray_dataarray_dict_no_metadata(
                 if export_with_multi_index:
                     assert isinstance(df.index, pd.MultiIndex)
 
-                    if hasattr(xr, "Coordinates"):
-                        coords = xr.Coordinates.from_pandas_multiindex(
-                            df.index, "multi_index"
-                        )
-                        xrdarray = xr.DataArray(df[name], coords=coords)
-                    else:
-                        # support xarray < 2023.8.0, can be removed when we drop support for that
-                        xrdarray = xr.DataArray(df[name], [("multi_index", df.index)])
+                    coords = xr.Coordinates.from_pandas_multiindex(
+                        df.index, "multi_index"
+                    )
+                    xrdarray = xr.DataArray(df[name], coords=coords)
                 else:
                     xrdarray = df.to_xarray().get(name, xr.DataArray())
 

tests/dataset/measurement/test_inferred_parameters_fix.py

@@ -215,6 +215,90 @@ def test_inferred_parameters_in_actual_measurement_1d(experiment, DAC):
     assert len(df) == num_points
 
 
+def test_inferred_parameters_in_actual_measurement_2d(experiment, DAC):
+    """
+    2D version: both axes are DelegateParameters inferred from DAC channels.
+    Ensures inferred parameters on both axes are saved and exported correctly.
+    """
+    num_points_x = 5
+    num_points_y = 7
+
+    # Delegate parameters (axes)
+    del_param_1 = DelegateParameter("del_param_1", label="del param 1", source=DAC.ch1)
+    del_param_2 = DelegateParameter("del_param_2", label="del param 2", source=DAC.ch2)
+
+    # Measurement parameter depending on both delegate parameters
+    meas_parameter = ManualParameter("meas_parameter", initial_value=0.0)
+
+    meas = Measurement(name="test_measurement_2d", exp=experiment)
+
+    # Register underlying basis (standalone) parameters
+    meas.register_parameter(DAC.ch1)
+    meas.register_parameter(DAC.ch2)
+
+    # Register delegate parameters inferred from their respective DAC channels
+    meas.register_parameter(del_param_1, basis=(DAC.ch1,))
+    meas.register_parameter(del_param_2, basis=(DAC.ch2,))
+
+    # Register measurement parameter with 2D setpoints
+    meas.register_parameter(meas_parameter, setpoints=(del_param_1, del_param_2))
+
+    with meas.run() as datasaver:
+        for x in np.linspace(0, 1, num_points_x):
+            for y in np.linspace(0, 2, num_points_y):
+                del_param_1.set(x)
+                del_param_2.set(y)
+                meas_parameter.set(x + y)  # simple deterministic value
+
+                datasaver.add_result(
+                    (DAC.ch1, DAC.ch1()),
+                    (DAC.ch2, DAC.ch2()),
+                    (del_param_1, del_param_1()),
+                    (del_param_2, del_param_2()),
+                    (meas_parameter, meas_parameter()),
+                )
+
+    dataset = datasaver.dataset
+    param_data = dataset.get_parameter_data()
+
+    assert "meas_parameter" in param_data
+    assert len(param_data) == 1  # only one top-level dependent parameter tree
+
+    tree = param_data["meas_parameter"]
+    expected_keys = {
+        "meas_parameter",
+        "del_param_1",
+        "del_param_2",
+        "dummy_dac_ch1",
+        "dummy_dac_ch2",
+    }
+    assert set(tree.keys()) == expected_keys
+
+    total_points = num_points_x * num_points_y
+    for k in expected_keys:
+        assert len(tree[k]) == total_points, f"{k} should have {total_points} entries"
+
+    # xarray export
+    xarr = dataset.to_xarray_dataset()
+    assert "meas_parameter" in xarr.data_vars
+    assert "del_param_1" in xarr.coords
+    assert "del_param_2" in xarr.coords
+    # inferred (basis) parameters not exported
+    assert "dummy_dac_ch1" not in xarr.coords
+    assert "dummy_dac_ch2" not in xarr.coords
+
+    assert xarr.meas_parameter.dims == ("del_param_1", "del_param_2")
+    assert xarr.meas_parameter.shape == (num_points_x, num_points_y)
+
+    # pandas export
+    df = dataset.to_pandas_dataframe()
+    assert "meas_parameter" in df.columns
+    assert df.index.names == ["del_param_1", "del_param_2"]
+    assert "dummy_dac_ch1" not in df.columns
+    assert "dummy_dac_ch2" not in df.columns
+    assert len(df) == total_points
+
+
 def test_multiple_dependent_parameters_no_cross_contamination(experiment):
     """
     Test that multiple dependent parameters that depend on the same independent