Export inferred parameters related to dependencies

Author: jenshnielsen

src/qcodes/dataset/exporters/export_to_xarray.py

@@ -162,21 +162,52 @@ def _xarray_data_array_direct(
     import xarray as xr
 
     meas_paramspec = dataset.description.interdeps.graph.nodes[name]["value"]
-    _, deps, _ = dataset.description.interdeps.all_parameters_in_tree_by_group(
+    _, deps, inferred = dataset.description.interdeps.all_parameters_in_tree_by_group(
         meas_paramspec
     )
-    dep_axis = {}
+    # Build coordinate axes from direct dependencies preserving their order
+    dep_axis: dict[str, npt.NDArray] = {}
     for axis, dep in enumerate(deps):
         dep_array = subdict[dep.name]
         dep_axis[dep.name] = dep_array[
             tuple(slice(None) if i == axis else 0 for i in range(dep_array.ndim))
         ]
 
-    da = xr.Dataset(
+    extra_coords: dict[str, tuple[tuple[str, ...] | tuple[str], npt.NDArray]] = {}
+    for inf in inferred:
+        # skip parameters already used as primary coordinate axes
+        if inf.name in dep_axis:
+            continue
+        # add only if data for this parameter is available
+        if inf.name not in subdict:
+            continue
+
+        inf_related = dataset.description.interdeps.find_all_parameters_in_tree(inf)
+
+        related_deps = inf_related.intersection(set(deps))
+        related_top_level = inf_related.intersection({meas_paramspec})
+
+        if len(related_top_level) > 0:
+            raise NotImplementedError(
+                "Adding inferred coords related to top level param is not yet supported"
+            )
+
+        inf_data = subdict[inf.name][
+            tuple(slice(None) if dep in related_deps else 0 for dep in deps)
+        ]
+        inf_coords = [dep.name for dep in deps if dep in related_deps]
+
+        extra_coords[inf.name] = (tuple(inf_coords), inf_data)
+
+    # Compose coordinates dict including dependency axes and extra inferred coords
+    coords: dict[str, tuple[tuple[str, ...] | tuple[str], npt.NDArray] | npt.NDArray]
+    coords = {**dep_axis, **extra_coords}
+
+    ds = xr.Dataset(
         {name: (tuple(dep_axis.keys()), subdict[name])},
-        coords=dep_axis,
-    )[name]
-    return da
+        coords=coords,
+    )
+    return ds[name]
 
 
 def load_to_xarray_dataarray_dict(

tests/dataset/measurement/test_inferred_parameters_fix.py

@@ -305,14 +305,25 @@ def test_inferred_parameters_in_actual_measurement_2d(
     )
 
     assert "meas_parameter" in xarr.data_vars
+
     assert "del_param_1" in xarr.coords
+    assert xarr.coords["del_param_1"].shape == (num_points_x,)
+    assert xarr.coords["del_param_1"].dims == ("del_param_1",)
+
     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.coords["del_param_2"].shape == (num_points_y,)
+    assert xarr.coords["del_param_2"].dims == ("del_param_2",)
+
+    assert "dummy_dac_ch1" in xarr.coords
+    assert xarr.coords["dummy_dac_ch1"].shape == (num_points_x,)
+    assert xarr.coords["dummy_dac_ch1"].dims == ("del_param_1",)
+
+    assert "dummy_dac_ch2" in xarr.coords
+    assert xarr.coords["dummy_dac_ch2"].shape == (num_points_y,)
+    assert xarr.coords["dummy_dac_ch2"].dims == ("del_param_2",)
 
-    assert xarr.meas_parameter.dims == ("del_param_1", "del_param_2")
-    assert xarr.meas_parameter.shape == (num_points_x, num_points_y)
+    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()