Insert derived parameters in data

Author: jenshnielsen

src/qcodes/dataset/exporters/export_to_xarray.py

@@ -173,6 +173,7 @@ def _xarray_data_array_direct(
         ]
 
     extra_coords: dict[str, tuple[tuple[str, ...], npt.NDArray]] = {}
+    extra_data_vars: dict[str, tuple[tuple[str, ...], npt.NDArray]] = {}
     for inf in inferred:
         # skip parameters already used as primary coordinate axes
         if inf.name in dep_axis:
@@ -187,26 +188,37 @@ def _xarray_data_array_direct(
         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]
+            # If inferred param is related to the top-level measurement parameter,
+            # add it as a data variable with the full dependency dimensions
+            inf_data_full = subdict[inf.name]
+            inf_dims_full = tuple(dep_axis.keys())
+            extra_data_vars[inf.name] = (inf_dims_full, inf_data_full)
+        else:
+            # Otherwise, add as a coordinate along the related dependency axes only
+            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)
+            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, ...], npt.NDArray] | npt.NDArray]
     coords = {**dep_axis, **extra_coords}
 
-    ds = xr.Dataset(
-        {name: (tuple(dep_axis.keys()), subdict[name])},
-        coords=coords,
-    )
-    return ds[name]
+    # Compose data variables dict including measured var and any inferred data vars
+    data_vars: dict[str, tuple[tuple[str, ...], npt.NDArray]] = {
+        name: (tuple(dep_axis.keys()), subdict[name])
+    }
+    data_vars.update(extra_data_vars)
+
+    ds = xr.Dataset(data_vars, coords=coords)
+    da = ds[name]
+    if len(extra_data_vars) > 0:
+        # stash extra data vars to be added at dataset assembly time
+        # mapping: var_name -> (dims_tuple, numpy array)
+        da.attrs["_qcodes_extra_data_vars"] = extra_data_vars
+    return da
 
 
 def load_to_xarray_dataarray_dict(
@@ -272,6 +284,13 @@ def load_to_xarray_dataset(
     # and python/typing#445 are resolved.
     xrdataset = xr.Dataset(cast("dict[Hashable, xr.DataArray]", data_xrdarray_dict))
 
+    # add any stashed extra data variables created during direct export
+    for _, dataarray in data_xrdarray_dict.items():
+        extras = dataarray.attrs.pop("_qcodes_extra_data_vars", None)
+        if isinstance(extras, dict):
+            for var_name, (dims, values) in extras.items():
+                xrdataset[var_name] = (dims, values)
+
     _add_param_spec_to_xarray_coords(dataset, xrdataset)
     _add_param_spec_to_xarray_data_vars(dataset, xrdataset)
     _add_metadata_to_xarray(dataset, xrdataset)

tests/dataset/test_dataset_export.py

@@ -1971,3 +1971,59 @@ def test_measurement_2d_with_inferred_setpoint(
         inf_idx = xr_ds.coords[name].indexes
         assert set(inf_idx.keys()) == {"y"}
         assert inf_idx["y"].equals(xr_ds.indexes["y"])
+
+
+def test_measurement_2d_top_level_inferred_is_data_var(
+    experiment: Experiment, caplog: LogCaptureFixture
+) -> None:
+    """
+    If an inferred parameter is related to the top-level measured parameter,
+    it must be exported as a data variable (not a coordinate) with the full
+    dependency dimensions.
+    """
+    nx, ny = 2, 3
+    x_vals = np.linspace(0.0, 1.0, nx)
+    y_vals = np.linspace(10.0, 12.0, ny)
+
+    # Define a measured signal and an inferred param both defined on (x, y)
+    # The inferred param is related to the measured top-level param in the graph
+    meas = Measurement(exp=experiment, name="2d_top_level_inferred")
+    meas.register_custom_parameter("x", paramtype="numeric")
+    meas.register_custom_parameter("y", paramtype="numeric")
+    # Register measured top-level
+    meas.register_custom_parameter("signal", setpoints=("x", "y"), paramtype="numeric")
+    # Register inferred related to top-level (basis includes the measured top-level)
+    meas.register_custom_parameter("derived", basis=("signal",), paramtype="numeric")
+    meas.set_shapes({"signal": (nx, ny)})
+
+    with meas.run() as datasaver:
+        for ix in range(nx):
+            for iy in range(ny):
+                x = float(x_vals[ix])
+                y = float(y_vals[iy])
+                signal = x + y
+                derived = 2.0 * signal  # inferred from top-level
+                datasaver.add_result(
+                    ("x", x), ("y", y), ("signal", signal), ("derived", derived)
+                )
+
+    ds = datasaver.dataset
+    caplog.clear()
+    with caplog.at_level(logging.INFO):
+        xr_ds = ds.to_xarray_dataset()
+
+    # Direct path log should be present
+    assert any(
+        "Exporting signal to xarray using direct method" in record.message
+        for record in caplog.records
+    )
+
+    # The derived param should be a data variable with dims (x, y), not a coord
+    assert "derived" in xr_ds.data_vars
+    assert "derived" not in xr_ds.coords
+    assert xr_ds["derived"].dims == ("x", "y")
+
+    expected_signal = x_vals[:, None] + y_vals[None, :]
+    expected_derived = 2.0 * expected_signal
+    np.testing.assert_allclose(xr_ds["signal"].values, expected_signal)
+    np.testing.assert_allclose(xr_ds["derived"].values, expected_derived)